astrodatadownloader: only keep what we really need

This commit is contained in:
2024-03-06 11:10:11 +01:00
parent bf8c1d9915
commit 51d0baaccb
@@ -1,4 +1,5 @@
from astroquery.mast import Observations from astroquery.mast import Observations
from astropy import table
import numpy as np import numpy as np
from enum import Enum from enum import Enum
from copy import deepcopy from copy import deepcopy
@@ -70,67 +71,19 @@ def getStarObservations(starName: str, sources: list[ObservationSource], sequenc
return obs[obsWantedFilter] return obs[obsWantedFilter]
def fitsFilenameFilterFunc(table, key_colnames): def downloadStarProducts(obs_wanted):
if(str(table["productFilename"]).lower().endswith(".fits")):
return True
return False
sequencesFilter: list[int] = []
def fullSequenceOnlyFiles(table, key_colnames):
if(len(sequencesFilter) > 1):
seqStr = f"s{sequencesFilter[0]:04d}-s{sequencesFilter[1]:04d}"
else:
seqStr = f"s{sequencesFilter[0]:04d}"
if(seqStr in str(table["productFilename"]).lower()):
return True
return False
def downloadStarProducts(starName: str, sequences: list[int], dataType: DataType,
downloadPath: str = "./", fitsOnly: bool = True,
fullSectorOnly: bool = True) -> str:
obs_wanted = getStarObservations(starName, sequences)
dataProducts = Observations.get_product_list(obs_wanted) dataProducts = Observations.get_product_list(obs_wanted)
productSubGroups = [] # LC: TESS lightcurve
match dataType: # SLC: Kepler short cadence lightcurve
case DataType.DataValidation: # LLC: Kepler long cadence lightcurve
productSubGroups.append("DVT") productSubGroups = ["LC", "SLC"]
productSubGroups.append("DVM")
productSubGroups.append("DVS")
productSubGroups.append("DVR")
case DataType.LightCurve:
productSubGroups.append("LC")
case DataType.TargetPixel:
productSubGroups.append("TP")
productsWanted = Observations.filter_products(dataProducts, productsWanted = Observations.filter_products(dataProducts,
productSubGroupDescription=productSubGroups) productSubGroupDescription=productSubGroups)
productsWanted = productsWanted.group_by("productFilename") productsWanted = productsWanted.group_by("productFilename")
if(fitsOnly):
productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)
if(len(sequences) == 0):
fullSectorOnly = False
if(fullSectorOnly):
global sequencesFilter
if(len(sequences) == 1 and not
(dataType == DataType.LightCurve or dataType == DataType.TargetPixel)):
sequencesFilter = (sequences[0], sequences[0])
elif(len(sequences) > 1 and
(dataType == DataType.LightCurve or dataType == DataType.TargetPixel)):
sequencesFilter = [sequences[0]]
else:
sequencesFilter = sequences
productsWanted = productsWanted.groups.filter(fullSequenceOnlyFiles)
filenames = Observations.download_products(productsWanted) filenames = Observations.download_products(productsWanted)
return parse_manifest(filenames) return parse_manifest(filenames)