From 298cba6616bce01c2681b11ee46f4ae34a8d5a64 Mon Sep 17 00:00:00 2001 From: SGCMarkus Date: Thu, 7 Aug 2025 17:06:19 +0200 Subject: [PATCH] download_tess_kepler: rework to send all requests in batches --- download_tess_kepler.py | 157 ++++++++++++++++++++++++++-------------- 1 file changed, 102 insertions(+), 55 deletions(-) diff --git a/download_tess_kepler.py b/download_tess_kepler.py index d15e976..914761f 100644 --- a/download_tess_kepler.py +++ b/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,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) -tessDataProducts = Observations.filter_products(Observations.get_product_list(tessObs), productSubGroupDescription=["LC"]) +# 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) -keplerObs = Observations.query_criteria(obs_collection="Kepler", intentType="science", calib_level=2) -keplerDataProducts = Observations.filter_products(Observations.get_product_list(keplerObs), productSubGroupDescription=["SLC"]) +# 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) -k2Obs = Observations.query_criteria(obs_collection="K2", intentType="science", calib_level=2) -k2DataProducts = Observations.filter_products(Observations.get_product_list(k2Obs), productSubGroupDescription=["SLC"]) +# 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) -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)