diff --git a/download_tess_kepler.py b/download_tess_kepler.py new file mode 100644 index 0000000..d15e976 --- /dev/null +++ b/download_tess_kepler.py @@ -0,0 +1,273 @@ +from astroquery.mast import Observations +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 + +DEFAULT_DB = "stars.db" + +class StarDB(): + __instances = list() + __dbNames = list() + + @staticmethod + def getAllInstances(): + return StarDB.__instances + + @staticmethod + def getInstance(dbName: str): + if(dbName in StarDB.__dbNames): + return StarDB.__instances[StarDB.__dbNames.index(dbName)] + else: + instance = StarDB(dbName) + StarDB.__instances.append(instance) + StarDB.__dbNames.append(dbName) + return instance + + def __init__(self, dbName: str): + if(dbName in StarDB.__dbNames): + raise Exception(f"StarDB {dbName} cannot be instantiated more than once!") + else: + self.dbName = dbName + StarDB.__dbNames.append(dbName) + StarDB.__instances.append(self) + + self.error = StarDBError.NO_ERROR + + try: + self.connection = sqlite3.connect(dbName) + self.dbCursor = self.connection.cursor() + except: + self.error = StarDBError.DB_CONNECTION_FAILED + + if(self.error is not StarDBError.NO_ERROR): + raise Exception(f"StarDB cannot connect to the database: {self.error}") + else: + self.initializeDB() + + def close(self): + StarDB.__instances.remove(self) + StarDB.__dbNames.remove(self.dbName) + self.connection.close() + + def initializeDB(self): + print("Initializing database") + self.dbCursor.execute("""CREATE TABLE IF NOT EXISTS starnames + ( + mainName TEXT NOT NULL, + altName TEXT UNIQUE, + PRIMARY KEY (mainName, altName) + );""") + + self.dbCursor.execute("""CREATE TABLE IF NOT EXISTS sequenceSourceNames + ( + sourceName TEXT, + PRIMARY KEY (sourceName) + );""") + + self.dbCursor.execute("""CREATE TABLE IF NOT EXISTS stars + ( + mainName TEXT NOT NULL, + sourceName TEXT NOT NULL, + sequence INTEGER NOT NULL, + filename TEXT NOT NULL, + PRIMARY KEY (mainName, sourceName, sequence), + FOREIGN KEY (mainName) REFERENCES starnames(mainName), + FOREIGN KEY (sourceName) REFERENCES sequenceSourceNames(sourceName) + );""") + + self.dbCursor.execute("""CREATE TABLE IF NOT EXISTS starInfo + ( + mainName TEXT, + spType TEXT, + rotVel REAL, + rotVelUnit TEXT, + dist REAL, + distUnit TEXT, + PRIMARY KEY (mainName), + FOREIGN KEY (mainName) REFERENCES starnames(mainName) + );""") + + self.dbCursor.execute("""CREATE TABLE IF NOT EXISTS starFoldedFitType + ( + mainName TEXT, + foldedFitType TEXT, + PRIMARY KEY (mainName), + FOREIGN KEY (mainName) REFERENCES starnames(mainName) + );""") + + self.insertDefaultValues() + self.dbCorrection() + + self.connection.commit() + + def dbCorrection(self): + self.dbCursor.execute("""UPDATE starInfo + SET spType = "-" + WHERE spType IS NULL + OR length(spType) = 0;""") + self.dbCursor.execute("""UPDATE stars + SET filename = replace(filename, '\\', '/') + WHERE filename LIKE '%\\%';""") + self.dbCursor.execute("""UPDATE starFoldedFitType + SET foldedFitType = 'sine' + WHERE foldedFitType IS NULL OR foldedFitType = '';""") + + + def insertDefaultValues(self): + sources = ["TESS", "Kepler", "K2"] + for s in sources: + self.dbCursor.execute(f"""INSERT OR IGNORE INTO + sequenceSourceNames(sourceName) + VALUES(\"{s}\")""") + + def insertStarAlternativeNames(self, mainName: str, altNames): + for an in altNames: + self.dbCursor.execute(f"""INSERT OR IGNORE INTO + starnames(mainName, altName) + VALUES(\"{mainName}\", \"{an["ID"]}\")""") + + def insertStarInformations(self, mainName, spType, rotVel, rotVelUnit, + dist, distUnit): + print(f"""INSERT OR IGNORE INTO + starInfo(mainName, spType, rotVel, rotVelUnit, dist, distUnit) + VALUES(\"{mainName}\", \"{spType}\", {rotVel}, \"{rotVelUnit}\", {dist}, \"{distUnit}\")""") + self.dbCursor.execute(f"""INSERT OR IGNORE INTO + starInfo(mainName, spType, rotVel, rotVelUnit, dist, distUnit) + VALUES(\"{mainName}\", \"{spType}\", {rotVel}, \"{rotVelUnit}\", {dist}, \"{distUnit}\")""") + + def insertStar(self, mainName: str, altNames, pm, + spType, rotVel, rotVelUnit, dist, distUnit): + self.insertStarAlternativeNames(mainName, altNames) + self.insertStarInformations(mainName, spType, rotVel, rotVelUnit, + dist, distUnit) + + file = pm["Local Path"].replace("\\", "/") + source = pm["source"] + sequence = pm["sector"] + + self.dbCursor.execute(f"""INSERT OR IGNORE INTO + stars(mainName, sourceName, sequence, filename) + VALUES(\"{mainName}\", \"{source}\", {sequence}, \"{file}\")""") + + self.connection.commit() + + +def parse_manifest(manifest): + """ + Parse manifest and add back columns that are useful for TESS DV exploration. + """ + results = deepcopy(manifest) + filenames = [] + sources = [] + sector_range = [] + exts = [] + starIDs = [] + for i,f in enumerate(manifest['Local Path']): + file_parts = np.array(np.unique(f.split(sep = '-'))) + try: + with fits.open(f, mode="readonly") as hdu: + if("SECTOR" in hdu[0].header): + sector_range.append(f"{hdu[0].header['SECTOR']}") + elif("TTABLEID" in hdu[0].header): + sector_range.append(f"{hdu[0].header['TTABLEID']}") + starIDs.append(hdu[0].header['OBJECT']) + hdu.close() + except: + sector_range.append("-") + starIDs.append("-") + + if("TESS" in f): + sources.append("TESS") + elif("Kepler" in f): + sources.append("Kepler") + elif("K2" in f): + sources.append("K2") + + path_parts = np.array(f.split(sep = '/')) + filenames.append(path_parts[-1]) + exts.append(path_parts[-1][-8:]) + + results.add_column(table.Column(name = "starID", data = starIDs)) + results.add_column(table.Column(name = "filename", data = filenames)) + results.add_column(table.Column(name = "source", data = sources)) + results.add_column(table.Column(name = "sector", 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 fitsFilenameFilterFunc(table, key_colnames): + if(str(table["productFilename"]).lower().endswith(".fits") or + str(table["productFilename"]).lower().endswith(".fit")): + 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") + +tessObs = Observations.query_criteria(obs_collection="TESS", intentType="science", calib_level=2) +tessDataProducts = Observations.filter_products(Observations.get_product_list(tessObs), productSubGroupDescription=["LC"]) + +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)) + +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) +