srg/heasarc/management/commands/00_heasarc_object_class.py
2024-04-25 17:50:51 +03:00

89 lines
3.4 KiB
Python

from django.core.management.base import BaseCommand, CommandError
from datetime import date
from heasarc.models import INPUT_DATA_DIR
import datetime
from django.utils import timezone
import astropy
from astropy.io import ascii
import pandas as pd
import pymysql
from sqlalchemy import create_engine
from heasarc.tdat import tDat
from heasarc.models import HeasarcTable, TableColumn, HeasarcObjectClass
def load_heasarc_table(filename, engine):
tdatfile=INPUT_DATA_DIR+'/dump/'+filename
data = astropy.table.Table.read(tdatfile, format='ascii.tdat')
keywords = data.meta['keywords']
for key, value in keywords.items():
print(key, '->', value)
try:
table=HeasarcTable.objects.get(name__exact=keywords['table_name'])
table.delete()
except:
pass
table_name=keywords['table_name']
table = HeasarcTable(name=keywords['table_name'])
table.description=keywords['table_description'].replace('\"', '')
table.document_url=keywords['table_document_url']
if 'default_search_radius' in keywords:
table.search_radius=int(keywords['default_search_radius'])
if 'frequency_regime' in keywords:
table.frequency_regime=keywords['frequency_regime']
table.observatory_name=keywords['observatory_name']
table.security=keywords['table_security']
if 'table_author' in keywords:
table.author=keywords['table_author']
if 'catalog_bibcode' in keywords:
table.bibcode=keywords['catalog_bibcode']
if 'declination' in keywords:
table.declination=keywords['declination'].replace('@', '')
if 'right_ascension' in keywords:
table.right_ascension=keywords['right_ascension'].replace('@', '')
table.observatory_name = keywords['observatory_name']
if 'parameter_defaults' in keywords:
table.parameter_defaults = keywords['parameter_defaults']
table.save()
cols = data.meta['cols']
for key, value in cols.items():
column = TableColumn(table=table)
column.name=key
column.tdat_type=value['type']
column.description=value['description']
column.save()
print(key, '->', value)
# convert data frame to pandas
df = data.to_pandas()
sql = 'DROP TABLE IF EXISTS '+table_name+';'
result = engine.execute(sql)
# Insert whole DataFrame into MySQL
df.to_sql(table_name, con = engine, if_exists = 'append', chunksize = 200000)
for item in data:
obj = HeasarcObjectClass(table=table,
class_id=int(item['class_id']),
class_name=item['class_name'])
obj.save()
print('--> Successfully loaded "%s"' % table_name)
pass
class Command(BaseCommand):
help = 'Initiates data dase'
# def add_arguments(self, parser):
# parser.add_argument('poll_id', nargs='+', type=int)
def handle(self, *args, **options):
load_heasarc_table('heasarc_class.tdat.gz', engine)
self.stdout.write(self.style.SUCCESS('Done'))