code formatting

management/commands/load_survey.py

@@ -1,21 +1,21 @@
 # uplim/management/commands/load_survey.py
 
 import numpy as np
+
 from astropy.io import fits
+from itertools import islice
+from datetime import datetime
 
 from django.core.management.base import BaseCommand
 from django.db import transaction
 from uplim.models import Pixel
 from django.db.models import Max
 
-from itertools import islice
-
-from datetime import datetime
-
 # DEFINE BATCH SIZE AND BATCH
 # **************************************************************
 
-#BATCH_SIZE = 1000000
+# BATCH_SIZE = 1000000
+
 
 def batch(iterable, size):
     """
@@ -29,8 +29,6 @@ def batch(iterable, size):
         yield chunk
 
 
-
-
 class Command(BaseCommand):
     help = "Process FITS files and store the data in the database"
 
@@ -39,40 +37,33 @@ class Command(BaseCommand):
 
     def add_arguments(self, parser):
         parser.add_argument(
-            '--counts', 
-            type=str, 
-            required=True,
-            help='Path of the counts file'
+            "--counts", type=str, required=True, help="Path of the counts file"
         )
         parser.add_argument(
-            '--exposure', 
-            type=str, 
-            required=True,
-            help='Path of the exposure file'
+            "--exposure", type=str, required=True, help="Path of the exposure file"
         )
         parser.add_argument(
-            '--survey_number', 
-            type=int, 
+            "--survey_number",
+            type=int,
             required=True,
-            help='Integer ID of the survey being read'
+            help="Integer ID of the survey being read",
         )
         parser.add_argument(
-            '--batch_size', 
-            type=int, 
+            "--batch_size",
+            type=int,
             default=1000,
-            help='Integer number of pixels to be inserted into the database at once'
+            help="Integer number of pixels to be inserted into the database at once",
         )
-        
 
     def handle(self, *args, **options):
-        
+
         # GET FILENAMES FROM ARGUMENTS
         # **************************************************************
-        
-        counts_file = options['counts']
-        exposure_file = options['exposure']
-        survey_number = options['survey_number']
-        BATCH_SIZE = options['batch_size']
+
+        counts_file = options["counts"]
+        exposure_file = options["exposure"]
+        survey_number = options["survey_number"]
+        BATCH_SIZE = options["batch_size"]
 
         self.stdout.write(f"\nCounts file:\t{counts_file}")
         self.stdout.write(f"Exposure file:\t{exposure_file}")
@@ -87,7 +78,6 @@ class Command(BaseCommand):
 
         counts_data = counts_map.ravel()
 
-
         with fits.open(exposure_file) as hdul:
 
             column_name = "T"
@@ -100,49 +90,48 @@ class Command(BaseCommand):
 
         self.stdout.write(f"\nCounts Data Shape:\t{counts_data.shape}")
         self.stdout.write(f"Exposure Data Shape:\t{exposure_data.shape}")
-        
+
         total_pixels = counts_data.shape[0]
         self.stdout.write(f"\nTotal pixels to insert:\t{total_pixels}")
 
-        assert counts_data.shape == exposure_data.shape, "Counts and exposure maps must have the same shape"
+        assert (
+            counts_data.shape == exposure_data.shape
+        ), "Counts and exposure maps must have the same shape"
 
         # CREATE THE SURVEY IF IT DOES NOT EXIST
         # **************************************************************
-        
+
         # with transaction.atomic():
-            
+
         #     survey,created = Survey.objects.get_or_create(number=survey_number)
-            
+
         #     if created:
         #         self.stdout.write(f"Created a new survey instance with number: {survey.number}")
         #     else:
         #         self.stdout.write(f"Using existing survey instance with the number: {survey.number}")
-        
+
         # FETCH THE LAST PROCESSED HPID AND CONTINUE FROM IT
         # **************************************************************
-        
+
         last_hpid = (
-            Pixel.objects
-                .filter(survey=survey_number)
-                .aggregate(max_hpid=Max('hpid'))['max_hpid']
+            Pixel.objects.filter(survey=survey_number).aggregate(max_hpid=Max("hpid"))[
+                "max_hpid"
+            ]
             or -1
         )
         start_index = last_hpid + 1
-        
-        
-        
+
         pixel_generator = (
             Pixel(
                 hpid=i,
                 counts=int(count),
                 exposure=float(exposure),
-                survey=survey_number
+                survey=survey_number,
             )
             for i, (count, exposure) in enumerate(zip(counts_data, exposure_data))
-            if i >= start_index                     
+            if i >= start_index
         )
 
-
         total_inserted = start_index
         # Process in batches
         for pixel_batch in batch(pixel_generator, BATCH_SIZE):
@@ -151,8 +140,6 @@ class Command(BaseCommand):
             total_inserted += len(pixel_batch)
             percentage = total_inserted / total_pixels * 100
             timestamp = datetime.now().strftime("%H:%M:%S")
-            self.stdout.write(
-                f"[{timestamp}] {percentage:.2f}% inserted"
-            )
+            self.stdout.write(f"[{timestamp}] {percentage:.2f}% inserted")
 
         self.stdout.write(f"Inserted a total of {total_inserted} pixels.")