diff --git a/main/astrodatadownloader/astrodatadownloader.py b/main/astrodatadownloader/astrodatadownloader.py index 9333403..4430c29 100644 --- a/main/astrodatadownloader/astrodatadownloader.py +++ b/main/astrodatadownloader/astrodatadownloader.py @@ -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)