274 lines
9.4 KiB
Python
274 lines
9.4 KiB
Python
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)
|
|
|