diff --git a/main/astrodatadownloader/astrodatadownloader.py b/main/astrodatadownloader/astrodatadownloader.py index 8d67c96..9333403 100644 --- a/main/astrodatadownloader/astrodatadownloader.py +++ b/main/astrodatadownloader/astrodatadownloader.py @@ -1,12 +1,51 @@ from astroquery.mast import Observations 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") @@ -29,4 +68,69 @@ def getStarObservations(starName: str, sources: list[ObservationSource], sequenc obsWantedFilter &= obsSourceFilter - return obs[obsWantedFilter] \ No newline at end of file + 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) + 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") + + 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) +