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 enum import Enum
-#from main.astrodatagui.db.StarsDB import StarDB
+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,6 +213,80 @@ def fitsFilenameFilterFunc(table, key_colnames):
         return True
     return False
 
+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]
+
+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)
+
+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 = pm["starID"]
+        dist = -1
+        distUnit = ""
+        rotVel = -1
+        rotVelUnit = ""
+        specType = "-"
+
+        try:
+            result = simbad.query_object(pm["starID"])
+            mainName = result["MAIN_ID"][0]
+        except:
+            print("Main Identifier Error", "Failed to grab main identifier, using TIC...")
+            noAltNames = True
+
+        try:
+            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")
+
+        try:
+            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"][0])
+            if(np.isnan(dist)):
+                dist = -1
+            distUnit = result["Distance_unit"][0]
+        except:
+            dist = -1
+            distUnit = ""
+
+        starDB.insertStar(mainName, altNames, pm,
+                          specType, rotVel, rotVelUnit, dist, distUnit)
+
+
 starDB: StarDB = StarDB.getInstance(DEFAULT_DB)
 simbad = Simbad()
 simbad.add_votable_fields("sptype")
@@ -213,61 +294,27 @@ simbad.add_votable_fields("velocity")
 simbad.add_votable_fields("diameter")
 simbad.add_votable_fields("distance")
 
-tessObs = Observations.query_criteria(obs_collection="TESS", intentType="science", calib_level=2)
+# TESS
-tessDataProducts = Observations.filter_products(Observations.get_product_list(tessObs), productSubGroupDescription=["LC"])
+tessObs = Observations.query_criteria(obs_collection="TESS")
+tessProductList = getProductList(tessObs)
+tessDataProducts = Observations.filter_products(tessProductList, productSubGroupDescription=["LC"], extension="fits")
+tessManifest = getManifest(tessDataProducts)
+insertManifestIntoDB(tessManifest)
 
-keplerObs = Observations.query_criteria(obs_collection="Kepler", intentType="science", calib_level=2)
+# Kepler
-keplerDataProducts = Observations.filter_products(Observations.get_product_list(keplerObs), productSubGroupDescription=["SLC"])
+keplerObs = Observations.query_criteria(obs_collection="Kepler")
+keplerProductList = getProductList(keplerObs)
+keplerDataProducts = Observations.filter_products(keplerProductList, productSubGroupDescription=["SLC"], extension="fits")
+keplerManifest = getManifest(keplerDataProducts)
+insertManifestIntoDB(keplerManifest)
 
-k2Obs = Observations.query_criteria(obs_collection="K2", intentType="science", calib_level=2)
+# K2
-k2DataProducts = Observations.filter_products(Observations.get_product_list(k2Obs), productSubGroupDescription=["SLC"])
+k2Obs = Observations.query_criteria(obs_collection="K2")
+k2ProductList = getProductList(k2Obs)
+k2DataProducts = Observations.filter_products(k2ProductList, productSubGroupDescription=["SLC"], extension="fits")
+k2Manifest = getManifest(k2DataProducts)
+insertManifestIntoDB(k2Manifest)
 
-productsWanted = table.vstack([tessDataProducts, keplerDataProducts, k2DataProducts])
-productsWanted = productsWanted.group_by("productFilename")
-productsWanted = productsWanted.groups.filter(fitsFilenameFilterFunc)
 
-manifest = parse_manifest(Observations.download_products(productsWanted))
 
-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
-
-    try:
-        altNames = Simbad.query_objectids(mainName)
-    except:
-        self.showErrorMessage("Identifier Error", "Failed to grab all identifiers, aborting...")
-        continue
-
-    try:
-        specType = result["SP_TYPE"]
-        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)
-    except:
-        print("No radial velocity found")
-
-    try:
-        dist = float(result["Distance_distance"])
-        if(np.isnan(dist)):
-            dist = -1
-        distUnit = result["Distance_unit"]
-    except:
-        dist = -1
-        distUnit = ""
-
-    starDB.insertStar(mainName, altNames, pm,
-                      specType, rotVel, rotVelUnit, dist, distUnit)