srg/astrobasis/management/commands/03_healpix_gaia.py
2024-04-26 12:43:00 +03:00

109 lines
2.9 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 astrobasis.models import GaiaSourceFileDR3, GAIADR3
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 heasarc.models import NSIDE_SOURCES, ORDER
import os
import glob
import logging
import time
from multiprocessing import Pool, cpu_count
USE_NCORES=1
def myfunc(n):
print(n)
batch_size=500
hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5)
plates=GaiaSourceFileDR3.objects.filter(status__exact="ready")
total_ready = len(plates)
print('Total <ready> Gaia files: '+str(total_ready))
step=int(total_ready/USE_NCORES)+1
idx1 = step*n
idx2 = step*(n+1)
if idx2 > total_ready:
idx2=total_ready
for i in range(idx1,idx2):
print(str(n)+": "+plates[i].filename)
phots=plates[i].gaiadr3_set.exclude(healpix__gt=0.0)
ra=phots.values_list('ra', flat=True)
dec=phots.values_list('dec', flat=True)
crd = SkyCoord(ra, dec, frame="fk5", unit="deg")
heal = hp.skycoord_to_healpix(crd)
for j, phot in enumerate(phots):
phot.healpix=heal[j]
GAIADR3.objects.bulk_update(phots, ['healpix'], batch_size=batch_size)
plates[i].status='healpix'
plates[i].save()
pass
def run_healpix():
start_time = time.time()
batch_size=500
hp = HEALPix(nside=NSIDE_SOURCES, order=ORDER, frame=FK5)
plates=GaiaSourceFileDR3.objects.filter(status__exact="ready")
total_ready = len(plates)
print('Total <ready> Gaia files: '+str(total_ready))
for plate in plates:
print(plate.filename)
phots=plate.gaiadr3_set.exclude(healpix__gt=0.0)
ra=phots.values_list('ra', flat=True)
dec=phots.values_list('dec', flat=True)
crd = SkyCoord(ra, dec, frame="fk5", unit="deg")
heal = hp.skycoord_to_healpix(crd)
for j, phot in enumerate(phots):
phot.healpix=heal[j]
GAIADR3.objects.bulk_update(phots, ['healpix'], batch_size=batch_size)
plate.status='healpix'
plate.save()
print("--- %s seconds ---" % (time.time() - start_time))
pass
class Command(BaseCommand):
help = 'Initiates healpix for Gaia'
def handle(self, *args, **options):
run_healpix()
"""
does not work for multitheading for some reason
with Pool(10) as p:
p.map(myfunc, range(USE_NCORES))
"""