136 lines
5.0 KiB
Python
136 lines
5.0 KiB
Python
from astroquery.mast import Observations
|
|
from astropy import table
|
|
from astropy.io import fits
|
|
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 = []
|
|
sources = []
|
|
sector_range = []
|
|
exts = []
|
|
for i,f in enumerate(manifest['Local Path']):
|
|
file_parts = np.array(np.unique(f.split(sep = '-')))
|
|
try:
|
|
with fits.open(f, mode="readonly") as hdu:
|
|
if("SECTOR" in hdu[0].header):
|
|
sector_range.append(f"{hdu[0].header['SECTOR']}")
|
|
elif("TTABLEID" in hdu[0].header):
|
|
sector_range.append(f"{hdu[0].header['TTABLEID']}")
|
|
hdu.close()
|
|
except:
|
|
sector_range.append("-")
|
|
|
|
if("TESS" in f):
|
|
sources.append("TESS")
|
|
elif("Kepler" in f):
|
|
sources.append("Kepler")
|
|
elif("K2" in f):
|
|
sources.append("K2")
|
|
|
|
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 = "source", data = sources))
|
|
results.add_column(table.Column(name = "sector", 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 fitsFilenameFilterFunc(table, key_colnames):
|
|
if(str(table["productFilename"]).lower().endswith(".fits") or
|
|
str(table["productFilename"]).lower().endswith(".fit")):
|
|
return True
|
|
return False
|
|
|
|
def getStarObservations(starName: str, keplerCadences, k2Cadences,
|
|
sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
|
|
if(len(keplerCadences) != 2):
|
|
raise Exception("keplerCadences needs to have length 2")
|
|
if(len(k2Cadences) != 2):
|
|
raise Exception("k2Cadences need to have length 2")
|
|
|
|
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
|
|
|
|
dataProducts = Observations.get_product_list(obs[obsWantedFilter])
|
|
|
|
# LC: TESS lightcurve
|
|
# SLC: Kepler short cadence lightcurve
|
|
# LLC: Kepler long cadence lightcurve
|
|
productSubGroups = ["LC", "SLC"]
|
|
tessProducts = Observations.filter_products(dataProducts, obs_collection="TESS",
|
|
productSubGroupDescription=["LC"])
|
|
keplerSubGroups = []
|
|
if(keplerCadences[0]):
|
|
keplerSubGroups.append("SLC")
|
|
if(keplerCadences[1]):
|
|
keplerSubGroups.append("LLC")
|
|
|
|
keplerProducts = Observations.filter_products(dataProducts, obs_collection="Kepler",
|
|
productSubGroupDescription=keplerSubGroups)
|
|
|
|
k2SubGroups = []
|
|
if(k2Cadences[0]):
|
|
k2SubGroups.append("SLC")
|
|
if(k2Cadences[1]):
|
|
k2SubGroups.append("LLC")
|
|
|
|
k2Products = Observations.filter_products(dataProducts, obs_collection="K2",
|
|
productSubGroupDescription=k2SubGroups)
|
|
|
|
productsWanted = table.vstack([tessProducts, keplerProducts, k2Products])
|
|
productsWanted = productsWanted.group_by("productFilename")
|
|
productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)
|
|
|
|
newCol = []
|
|
for pwCol in productsWanted:
|
|
for owCol in obs[obsWantedFilter]:
|
|
if(pwCol["obsID"] == owCol["obsid"]):
|
|
newCol.append(owCol["sequence_number"])
|
|
|
|
productsWanted.add_column(newCol, name="sequence_number")
|
|
return productsWanted
|
|
|
|
def downloadStarProducts(productsWanted):
|
|
filenames = Observations.download_products(productsWanted)
|
|
return parse_manifest(filenames)
|