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

217 lines
8.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
import numpy.ma as ma
from heasarc.tdat import tDat
from heasarc.models import HeasarcTable, TableColumn, HeasarcBase, HeasarcGAIADR2
from heasarc.models import HeasarcObjectClass, GaiaSourceFile
from heasarc.models import NSIDE_SOURCES, ORDER
from astropy_healpix import HEALPix
from astropy.coordinates import SkyCoord # High-level coordinates
from astropy.coordinates import ICRS, Galactic, FK4, FK5 # Low-level frames
from astropy.coordinates import Angle, Latitude, Longitude # Angles
import astropy.units as u
from django.db import DatabaseError, transaction
import os
import glob
import logging
def load_heasarc_table_atomic(filepath):
filename_w_ext = os.path.basename(filepath)
filename, file_extension = os.path.splitext(filename_w_ext)
#print(filename_w_ext)
try:
gaia = GaiaSourceFile.objects.get(filename=filename)
print("Gaia: %s is already loaded, skip" % filename)
return 0
except GaiaSourceFile.DoesNotExist:
print("Gaia: %s is not loaded" % filename)
pass
gaia = GaiaSourceFile(filename=filename)
gaia.save()
return
data = astropy.table.Table.read(filepath, format='ascii.csv',encoding='latin1')
print(data.info)
tables = HeasarcGAIADR2.objects.all()
tables.delete()
hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
try:
with transaction.atomic():
for item in data:
crd = SkyCoord(item['ra'], item['dec'], frame=FK5(), unit="deg")
healpix = hp.skycoord_to_healpix(crd)
obj = HeasarcGAIADR2.objects.create(healpix=healpix,
error_radius = max(item['ra_error'], item['dec_error']),
ra=item['ra'],
dec=item['dec'],
lii = item['l'],
bii = item['b'],
name=item['designation'],
solution_id = item['solution_id'],
source_id = item['source_id'],
ref_epoch = item['ref_epoch'],
ra_error = item['ra_error'],
dec_error = item['dec_error'],
parallax = item['parallax'],
parallax_error = item['parallax_error'],
pmra = item['pmra'],
pmra_error = item['pmra_error'],
pmdec = item['pmdec'],
pmdec_error = item['pmdec_error'],
phot_g_mean_mag = item['phot_g_mean_mag'],
phot_bp_mean_mag = item['phot_bp_mean_mag'],
phot_rp_mean_mag = item['phot_rp_mean_mag'])
obj.save()
except DatabaseError:
print('--> Database error "%s"' % filename)
print('--> Successfully loaded "%s"' % filename)
pass
def load_heasarc_table(filename):
data = astropy.table.Table.read(filename, format='ascii.csv',encoding='latin1')
print(data.info)
tables = HeasarcGAIADR2.objects.all()
tables.delete()
hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
for item in data:
#crd = SkyCoord(item['ra'], item['dec'], frame=FK5(), unit="deg")
#healpix = hp.skycoord_to_healpix(crd)
healpix = 0
ra = 0.0
dec = 0.0
lon = 0.0
lat = 0.0
if not (ma.is_masked(item['ra']) or ma.is_masked(item['dec'])):
ra = float(item['ra'])
dec = float(item['dec'])
crd = SkyCoord(ra, dec, frame=FK5(), unit="deg")
healpix = hp.skycoord_to_healpix(crd)
else:
continue
obj = HeasarcGAIADR2.objects.create(healpix=healpix,
ra=ra,
dec=dec,
lii = item['l'],
bii = item['b'],
name=item['designation'],
solution_id = item['solution_id'],
source_id = item['source_id'],
ref_epoch = item['ref_epoch'],
ra_error = item['ra_error'],
dec_error = item['dec_error'],
parallax = item['parallax'],
parallax_error = item['parallax_error'],
pmra = item['pmra'],
pmra_error = item['pmra_error'],
pmdec = item['pmdec'],
pmdec_error = item['pmdec_error'],
phot_g_mean_mag = item['phot_g_mean_mag'],
phot_bp_mean_mag = item['phot_bp_mean_mag'],
phot_rp_mean_mag = item['phot_rp_mean_mag'])
obj.save()
print('--> Successfully loaded "%s"' % filename)
pass
def load_heasarc_table_bulk(filename):
data = astropy.table.Table.read(filename, format='ascii.csv',encoding='latin1')
print(data.info)
tables = HeasarcGAIADR2.objects.all()
tables.delete()
hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
batch_size = 200
objs = (HeasarcGAIADR2(ra=data[i]['ra'],
dec=data[i]['dec'],
name=data[i]['designation'],
solution_id = data[i]['solution_id'],
source_id = data[i]['source_id'],
ref_epoch = data[i]['ref_epoch'],
ra_error = data[i]['ra_error'],
dec_error = data[i]['dec_error'],
parallax = data[i]['parallax'],
parallax_error = data[i]['parallax_error'],
pmra = data[i]['pmra'],
pmra_error = data[i]['pmra_error'],
pmdec = data[i]['pmdec'],
pmdec_error = data[i]['pmdec_error'],
phot_g_mean_mag = data[i]['phot_g_mean_mag'],
phot_bp_mean_mag = data[i]['phot_bp_mean_mag'],
phot_rp_mean_mag = data[i]['phot_rp_mean_mag'],)
for i in range(len(data)))
HeasarcBase.objects.bulk_create(objs,batch_size)
print('--> Successfully loaded "%s"' % filename)
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_atomic('/data/Gaia/gaia_source/csv/GaiaSource_970579492690449408_970786789287579392.csv.gz')
files = []
for file in glob.glob("/data/Gaia/gaia_source/csv/*.csv.gz"):
files.append(file)
for file in files:
print("Loading ",file)
load_heasarc_table_atomic(file)
return
hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5())
# update healpix only
srcs = HeasarcXMMSSC.objects.all()
for src in srcs:
crd = SkyCoord(src.ra, src.dec, frame=FK5(), unit="deg")
src.healpix = hp.skycoord_to_healpix(crd)
src.save()
self.stdout.write(self.style.SUCCESS('Done'))