download_tess_kepler: rework to send all requests in batches

download_tess_kepler.py

@@ -1,14 +1,21 @@
-from astroquery.mast import Observations
+from astroquery.mast import Observations, Catalogs
 from astropy import table
 from astropy.io import fits
 import numpy as np
 from copy import deepcopy
 from astroquery.simbad import Simbad
-
-#from main.astrodatagui.db.StarsDB import StarDB
+from enum import Enum
+import sqlite3
+import time
 
 DEFAULT_DB = "stars.db"
 
+BATCH_SIZE = 500
+
+class StarDBError(Enum):
+    NO_ERROR = 0
+    DB_CONNECTION_FAILED = 1
+
 class StarDB():
     __instances = list()
     __dbNames = list()
@@ -206,64 +213,72 @@ def fitsFilenameFilterFunc(table, key_colnames):
         return True
     return False
 
-starDB: StarDB = StarDB.getInstance(DEFAULT_DB)
-simbad = Simbad()
-simbad.add_votable_fields("sptype")
-simbad.add_votable_fields("velocity")
-simbad.add_votable_fields("diameter")
-simbad.add_votable_fields("distance")
+def split_table(table, batch_size):
+    """Yield successive batch_size-sized chunks from an Astropy Table."""
+    for start in range(0, len(table), batch_size):
+        yield table[start:start + batch_size]
 
-tessObs = Observations.query_criteria(obs_collection="TESS", intentType="science", calib_level=2)
-tessDataProducts = Observations.filter_products(Observations.get_product_list(tessObs), productSubGroupDescription=["LC"])
+def getProductList(obsList):
+    productList = []
+    for i, batch in enumerate(split_table(obsList, BATCH_SIZE)):
+        print(f"[{i+1}] Getting products for batch of {len(batch)} observations...")
+        try:
+            prod = Observations.get_product_list(batch)
+            productList.append(prod)
+        except Exception as e:
+            print(f"Failed on batch {i+1}: {e}")
+        time.sleep(100)  # delay so we dont immediately send another request
+    return table.vstack(productList)
 
-keplerObs = Observations.query_criteria(obs_collection="Kepler", intentType="science", calib_level=2)
-keplerDataProducts = Observations.filter_products(Observations.get_product_list(keplerObs), productSubGroupDescription=["SLC"])
-
-k2Obs = Observations.query_criteria(obs_collection="K2", intentType="science", calib_level=2)
-k2DataProducts = Observations.filter_products(Observations.get_product_list(k2Obs), productSubGroupDescription=["SLC"])
-
-productsWanted = table.vstack([tessDataProducts, keplerDataProducts, k2DataProducts])
-productsWanted = productsWanted.group_by("productFilename")
-productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)
-
-manifest = parse_manifest(Observations.download_products(productsWanted))
+def getManifest(dataProduct):
+    manifest = []
+    for batch in split_table(dataProduct, BATCH_SIZE//10):
+        partial_manifest = parse_manifest(Observations.download_products(batch))
+        manifest.append(partial_manifest)
+        time.sleep(100)  # delay so we dont immediately send another request
+    return table.vstack(manifest)
 
+def insertManifestIntoDB(manifest):
     for pm in manifest:
-    result = self.simbad.query_object(pm["starID"])[0]
-    try:
-        mainName = result["MAIN_ID"]
-    except:
-        self.showErrorMessage("Main Identifier Error", "Failed to grab main identifier, aborting...")
-        continue
+        result = pm["starID"]
+        dist = -1
+        distUnit = ""
+        rotVel = -1
+        rotVelUnit = ""
+        specType = "-"
 
         try:
-        altNames = Simbad.query_objectids(mainName)
+            result = simbad.query_object(pm["starID"])
+            mainName = result["MAIN_ID"][0]
         except:
-        self.showErrorMessage("Identifier Error", "Failed to grab all identifiers, aborting...")
-        continue
+            print("Main Identifier Error", "Failed to grab main identifier, using TIC...")
+            noAltNames = True
 
         try:
-        specType = result["SP_TYPE"]
+            altNames = simbad.query_objectids(mainName)
+        except:
+            print("Identifier Error", "Failed to grab all identifiers, aborting...")
+            altNames = []
+
+        try:
+            specType = result["SP_TYPE"][0]
             if(specType == ""):
                 specType = "-"
         except:
             print("No spectral type found")
-        specType = "-"
 
-    rotVel = -1
-    rotVelUnit = ""
         try:
-        if(result["RVZ_TYPE"] == "v"):
-            rotVel = result["RV_VALUE"]
-            rotVelUnit = str(result["RV_VALUE"].unit)
+            if(result["RVZ_TYPE"][0] == "v"):
+                rotVel = result["RV_VALUE"][0]
+                rotVelUnit = str(result["RV_VALUE"][0].unit)
         except:
             print("No radial velocity found")
 
         try:
-        dist = float(result["Distance_distance"])
+            dist = float(result["Distance_distance"][0])
             if(np.isnan(dist)):
                 dist = -1
-        distUnit = result["Distance_unit"]
+            distUnit = result["Distance_unit"][0]
         except:
             dist = -1
             distUnit = ""
@@ -271,3 +286,35 @@ for pm in manifest:
         starDB.insertStar(mainName, altNames, pm,
                           specType, rotVel, rotVelUnit, dist, distUnit)
 
+
+starDB: StarDB = StarDB.getInstance(DEFAULT_DB)
+simbad = Simbad()
+simbad.add_votable_fields("sptype")
+simbad.add_votable_fields("velocity")
+simbad.add_votable_fields("diameter")
+simbad.add_votable_fields("distance")
+
+# TESS
+tessObs = Observations.query_criteria(obs_collection="TESS")
+tessProductList = getProductList(tessObs)
+tessDataProducts = Observations.filter_products(tessProductList, productSubGroupDescription=["LC"], extension="fits")
+tessManifest = getManifest(tessDataProducts)
+insertManifestIntoDB(tessManifest)
+
+# Kepler
+keplerObs = Observations.query_criteria(obs_collection="Kepler")
+keplerProductList = getProductList(keplerObs)
+keplerDataProducts = Observations.filter_products(keplerProductList, productSubGroupDescription=["SLC"], extension="fits")
+keplerManifest = getManifest(keplerDataProducts)
+insertManifestIntoDB(keplerManifest)
+
+# K2
+k2Obs = Observations.query_criteria(obs_collection="K2")
+k2ProductList = getProductList(k2Obs)
+k2DataProducts = Observations.filter_products(k2ProductList, productSubGroupDescription=["SLC"], extension="fits")
+k2Manifest = getManifest(k2DataProducts)
+insertManifestIntoDB(k2Manifest)
+
+
+
+