Files
flaredetector/main/astrodatadownloader/astrodatadownloader.py
T

90 lines
3.1 KiB
Python

from astroquery.mast import Observations
from astropy import table
import numpy as np
from enum import Enum
from copy import deepcopy
class DataType(Enum):
DataValidation = 1
LightCurve = 2
TargetPixel = 3
class ObservationSource(Enum):
KEPLER = 1
K2 = 2
TESS = 3
# base from https://spacetelescope.github.io/mast_notebooks/notebooks/TESS/beginner_astroquery_dv/beginner_astroquery_dv.html
# modified for sector span
def parse_manifest(manifest):
"""
Parse manifest and add back columns that are useful for TESS DV exploration.
"""
results = deepcopy(manifest)
filenames = []
sector_range = []
exts = []
for i,f in enumerate(manifest['Local Path']):
file_parts = np.array(np.unique(f.split(sep = '-')))
sectors = list( map ( lambda x: x[0:2] == 's0', file_parts))
s1 = file_parts[sectors][0]
try:
s2 = file_parts[sectors][1]
except:
s2 = s1
if(s1 == s2):
sector_range.append("%s" % s1)
else:
sector_range.append("%s-%s" % (s1,s2))
path_parts = np.array(f.split(sep = '/'))
filenames.append(path_parts[-1])
exts.append(path_parts[-1][-8:])
results.add_column(table.Column(name = "filename", data = filenames))
results.add_column(table.Column(name = "sectors", data = sector_range))
results.add_column(table.Column(name = "fileType", data = exts))
results.add_column(table.Column(name = "index", data = np.arange(0,len(manifest))))
return results
def getStarObservations(starName: str, sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
obs = Observations.query_object(objectname=starName, radius="0 deg")
obsWantedFilter = obs["dataproduct_type"] == "timeseries"
if(len(sequences) > 1):
obsWantedFilter &= ((obs['sequence_number'] >= max(sequences)) &
(obs['sequence_number'] <= min(sequences)))
elif(len(sequences) == 1):
obsWantedFilter &= (obs['sequence_number'] == sequences[0])
obsSourceFilter = np.ma.MaskedArray(data=np.full(obsWantedFilter.shape, False),
mask=False, fill_value=True)
if(ObservationSource.KEPLER in sources):
obsSourceFilter |= (obs['obs_collection'] == "Kepler")
if(ObservationSource.K2 in sources):
obsSourceFilter |= (obs['obs_collection'] == "K2")
if(ObservationSource.TESS in sources):
obsSourceFilter |= (obs['obs_collection'] == "TESS")
obsWantedFilter &= obsSourceFilter
return obs[obsWantedFilter]
def downloadStarProducts(obs_wanted):
dataProducts = Observations.get_product_list(obs_wanted)
# LC: TESS lightcurve
# SLC: Kepler short cadence lightcurve
# LLC: Kepler long cadence lightcurve
productSubGroups = ["LC", "SLC"]
productsWanted = Observations.filter_products(dataProducts,
productSubGroupDescription=productSubGroups)
productsWanted = productsWanted.group_by("productFilename")
filenames = Observations.download_products(productsWanted)
return parse_manifest(filenames)