astrodatadownloader: add downloader function and all dependencies from old project

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]
+    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)
+