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 astropy import table
import numpy as np
from enum import Enum
from copy import deepcopy
@@ -70,67 +71,19 @@ def getStarObservations(starName: str, sources: list[ObservationSource], sequenc
return obs[obsWantedFilter]
def fitsFilenameFilterFunc(table, key_colnames):
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)
def downloadStarProducts(obs_wanted):
dataProducts = Observations.get_product_list(obs_wanted)
productSubGroups = []
match dataType:
case DataType.DataValidation:
productSubGroups.append("DVT")
productSubGroups.append("DVM")
productSubGroups.append("DVS")
productSubGroups.append("DVR")
case DataType.LightCurve:
productSubGroups.append("LC")
case DataType.TargetPixel:
productSubGroups.append("TP")
# 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")
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)
return parse_manifest(filenames)