SGCMarkus
2606d49093
* the Tables can become huge, and therefor require a lot of RAM -> * cleanup whenever we dont need a table anymore
345 lines
11 KiB
Python
345 lines
11 KiB
Python
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
|
|
import sqlite3
|
|
import time
|
|
import gc
|
|
|
|
DEFAULT_DB = "stars.db"
|
|
|
|
BATCH_SIZE = 500
|
|
|
|
class StarDBError(Enum):
|
|
NO_ERROR = 0
|
|
DB_CONNECTION_FAILED = 1
|
|
|
|
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
|
|
|
|
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")
|
|
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)
|
|
del tessObs
|
|
gc.collect()
|
|
tessDataProducts = Observations.filter_products(tessProductList, productSubGroupDescription=["LC"], extension="fits")
|
|
del tessProductList
|
|
gc.collect()
|
|
tessManifest = getManifest(tessDataProducts)
|
|
del tessDataProducts
|
|
gc.collect()
|
|
insertManifestIntoDB(tessManifest)
|
|
del tessManifest
|
|
gc.collect()
|
|
|
|
# Kepler
|
|
keplerObs = Observations.query_criteria(obs_collection="Kepler")
|
|
keplerProductList = getProductList(keplerObs)
|
|
del keplerObs
|
|
gc.collect()
|
|
keplerDataProducts = Observations.filter_products(keplerProductList, productSubGroupDescription=["SLC"], extension="fits")
|
|
del keplerProductList
|
|
gc.collect()
|
|
keplerManifest = getManifest(keplerDataProducts)
|
|
del keplerDataProducts
|
|
gc.collect()
|
|
insertManifestIntoDB(keplerManifest)
|
|
del keplerManifest
|
|
gc.collect()
|
|
|
|
# K2
|
|
k2Obs = Observations.query_criteria(obs_collection="K2")
|
|
k2ProductList = getProductList(k2Obs)
|
|
del k2Obs
|
|
gc.collect()
|
|
k2DataProducts = Observations.filter_products(k2ProductList, productSubGroupDescription=["SLC"], extension="fits")
|
|
del k2ProductList
|
|
gc.collect()
|
|
k2Manifest = getManifest(k2DataProducts)
|
|
del k2DataProducts
|
|
gc.collect()
|
|
insertManifestIntoDB(k2Manifest)
|
|
del k2Manifest
|
|
gc.collect()
|
|
|
|
|
|
|