102 lines
2.6 KiB
Python
102 lines
2.6 KiB
Python
from django.core.management.base import BaseCommand
|
|
from GaiaDBInterface.models import GaiaSource, CatalogFile
|
|
from astropy.coordinates import SkyCoord
|
|
import astropy_healpix as ah
|
|
import numpy as np
|
|
from datetime import datetime, timedelta
|
|
|
|
|
|
|
|
|
|
def current_time():
|
|
return (datetime.now() + timedelta(hours=3)).strftime("%H:%M:%S")
|
|
|
|
|
|
|
|
def update_catalog_file_status(catalog_file, status): #catalog file status updater
|
|
catalog_file.status = status
|
|
catalog_file.save()
|
|
|
|
|
|
|
|
def healpix(nside): #NSIDE 2048, ORDER ring or nested
|
|
|
|
for catalog_file in CatalogFile.objects.filter(status='INGESTED'):
|
|
|
|
|
|
print(f'[{current_time()}] Loading sources from {catalog_file.name}...')
|
|
sources = list(catalog_file.sources.all())
|
|
|
|
|
|
if catalog_file.status == 'INDEX_IN_PROGRESS':
|
|
#sources.update(healpix_ring_index=None, healpix_nested_index=None)
|
|
GaiaSource.objects.filter(catalog_file=catalog_file).update(
|
|
healpix_ring_index=None,
|
|
healpix_nested_index=None
|
|
)
|
|
|
|
|
|
ra_list = np.array([source.ra for source in sources])
|
|
dec_list = np.array([source.dec for source in sources])
|
|
|
|
|
|
skycoord = SkyCoord(
|
|
ra=ra_list,
|
|
dec=dec_list,
|
|
unit='deg',
|
|
frame='fk5'
|
|
)
|
|
|
|
|
|
ring_healpix = ah.HEALPix(
|
|
nside=nside,
|
|
order=ring,
|
|
frame='fk5'
|
|
)
|
|
ring_healpix_indices = healpix.skycoord_to_healpix(skycoord)
|
|
|
|
|
|
nested_healpix = ah.HEALPix(
|
|
nside=nside,
|
|
order=nested,
|
|
frame='fk5'
|
|
)
|
|
nested_healpix_indices = healpix.skycoord_to_healpix(skycoord)
|
|
|
|
|
|
for source, healpix_index in zip(sources, ring_healpix_indices):
|
|
source.healpix_ring_index = healpix_index
|
|
|
|
for source, healpix_index in zip(sources, nested_healpix_indices):
|
|
source.healpix_nested_index = healpix_index
|
|
|
|
|
|
update_catalog_file_status(catalog_file, 'INDEX_IN_PROGRESS')
|
|
|
|
GaiaSource.objects.bulk_update(sources, ['healpix_ring_index', 'healpix_nested_index'])
|
|
|
|
update_catalog_file_status(catalog_file, 'INDEXED')
|
|
|
|
|
|
print(f'[{current_time()}] Database updated, sources indexed successfully.')
|
|
|
|
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = 'Index sources using healpix.'
|
|
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument(
|
|
'nside',
|
|
type=int,
|
|
help='NSIDE parameter for HEALPix'
|
|
)
|
|
|
|
|
|
def handle(self, *args, **options):
|
|
nside = options['nside']
|
|
|
|
healpix(nside)
|
|
|