Removed Survey model, it is now a field in the Pixel model; all code rewritten for the change; added set_contaminated management command to set the source contamination flag; added CatalogSource model for the art-xc source catalog storage, the table is filled by the set_contaminated management script; the upper limit view now outputs extra fields, including a list of sources and their properties within 120 arcseconds

management/commands/set_contaminated.py

@@ -1,17 +1,34 @@
 # uplim/management/commands/set_contaminated.py
+
 # add custom flux-radius mapping?
 # add specifying the columns?
 # do contamination setting per survey?
 # include nside for surveys?
 
 from django.core.management.base import BaseCommand
+from django.db import transaction
 
 import pandas as pd
 import healpy as hp
 import numpy as np
 from astropy.coordinates import SkyCoord
 
-from uplim.models import Pixel
+from uplim.models import Pixel, CatalogSource
+
+from itertools import islice
+
+from datetime import datetime
+
+BATCH_SIZE=900
+
+def batch(iterable, size):
+    iterable = iter(iterable)
+    while True:
+        chunk = list(islice(iterable, size))
+        if not chunk:
+            break
+        yield chunk
+    
 
 class Command(BaseCommand):
     help = "Set the 'contaminated' flag for all pixels based on the fluxes in the provided catalog."
@@ -67,7 +84,7 @@ class Command(BaseCommand):
         # READ THE CATALOG FILE USING PANDAS READ_FWF
         # *******************************************
         
-        # Define column positions based on the byte ranges in your table
+        # Define column positions based on the byte ranges 
         colspecs = [
             (0, 4),    # SrcID (1-4)
             (5, 26),   # Name (6-26)
@@ -91,8 +108,49 @@ class Command(BaseCommand):
         # Read the file using the fixed-width format
         catalog = pd.read_fwf(catalog_file, colspecs=colspecs, names=colnames)
         
+        for col in ['Name', 'CName', 'Type']:
+            catalog[col] = catalog[col].fillna('')
+        
         self.stdout.write(str(catalog.head()))
                 
+        # LOAD THE CATALOG INTO THE DATABASE
+        # **********************************
+        
+        existing_srcids = set(
+            CatalogSource.objects.values_list('srcid', flat=True)
+        )
+        
+        to_create = []
+        
+        for _, row in catalog.iterrows():
+            
+            srcid = int(row['SrcID'])
+            if srcid in existing_srcids:
+                continue
+            to_create.append(
+                CatalogSource(
+                    srcid        = srcid,
+                    name         = row['Name'].strip(),
+                    ra_deg       = float(row['RAdeg']),
+                    dec_deg      = float(row['DEdeg']),
+                    pos_error    = float(row['ePos']),
+                    significance = float(row['Signi']),
+                    flux         = float(row['Flux']),
+                    flux_error   = float(row['e_Flux']),
+                    catalog_name = row['CName'].strip(),
+                    new_xray     = bool(int(row['NewXray'])),
+                    source_type  = row['Type'].strip()
+                )
+            )
+        
+        if to_create:
+            self.stdout.write(f'Inserting {len(to_create)} new catalog rows.')
+            for chunk in batch(to_create, BATCH_SIZE):
+                CatalogSource.objects.bulk_create(chunk, ignore_conflicts=True)
+            self.stdout.write('Catalog update complete.')
+        else:
+            self.stdout.write('All catalog rows already exist in the database.')
+        
         # hard coded nside and flux-radius mapping
         # maybe change that
         
@@ -168,9 +226,24 @@ class Command(BaseCommand):
 
         self.stdout.write("\nList ready, updating the database...")
         
+    
         
-        if masked_pixels_list:
-            Pixel.objects.filter(hpid__in=masked_pixels_list).update(contaminated=True)
-            self.stdout.write(f"\nMarked {len(masked_pixels_list)} pixels as contaminated.")
-        else:
-            self.stdout.write("No pixels marked as contaminated, exiting.")
+        if not masked_pixels_list:
+            self.stdout.write("No pixels marked as contaminated, exiting.")
+            return
+        
+        total = len(masked_pixels_list)
+        updated = 0
+        self.stdout.write(f'\nUpdating contaminated flag in batches of {BATCH_SIZE}')
+        
+        for chunk in batch(masked_pixels_list, BATCH_SIZE):
+            with transaction.atomic():
+                Pixel.objects.filter(hpid__in=chunk).update(contaminated=True)
+                
+            updated += len(chunk)
+            percentage = updated / total * 100
+            
+            timestamp = datetime.now().strftime("%H:%M:%S")
+            self.stdout.write(f'[{timestamp}] {updated}/{total} ({percentage:.1f}%) updated')
+        
+        self.stdout.write(f'\n Marked {updated} pixels as contaminated.')