reformatted

dirac-institute · Aug 6, 2024 · 866cfaf · 866cfaf
1 parent 0ad7aa8
commit 866cfaf
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/src/kbmod/reprojection.py b/src/kbmod/reprojection.py
@@ -124,9 +124,12 @@ def _reproject_work_unit(work_unit, common_wcs, frame="original"):
     unique_obstimes, unique_obstime_indices = work_unit.get_unique_obstimes_and_indices()
 
     stack = ImageStack()
-    for time in tqdm(unique_obstimes, bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Reprojecting"):
-        indices = list(np.where(obstimes == time)[0])
-        per_image_indices.append(indices)
+
+    for time, indices in tqdm(
+        zip(unique_obstimes, unique_obstime_indices),
+        bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT,
+        desc="Reprojecting",
+    ):
 
         science_add = np.zeros(common_wcs.array_shape, dtype=np.float32)
         variance_add = np.zeros(common_wcs.array_shape, dtype=np.float32)

diff --git a/src/kbmod/work_unit.py b/src/kbmod/work_unit.py
@@ -289,7 +289,9 @@ def from_fits(cls, filename):
 
             per_image_indices = []
             # Read in all the image files.
-            for i in tqdm(range(num_images), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Loading images"):
+            for i in tqdm(
+                range(num_images), bar_format=TQDMUtils.DEFAULT_TQDM_BAR_FORMAT, desc="Loading images"
+            ):
                 sci_hdu = hdul[f"SCI_{i}"]
 
                 # Read in the layered image from different extensions.