Cls indices fix

gallery/tutorials/tutorials/class_averaging.py

@@ -226,6 +226,7 @@
 est_shifts = avgs.averager.shifts
 est_dot_products = avgs.averager.dot_products
 
+# These are dictionaries mapping each class to arrays of attributes.
 print(f"Estimated Rotations: {est_rotations}")
 print(f"Estimated Shifts: {est_shifts}")
 print(f"Estimated Dot Products: {est_dot_products}")
@@ -241,7 +242,12 @@
 original_img_nbr = noisy_src.images[original_img_nbr_idx].asnumpy()[0]
 
 # Rotate using estimated rotations.
-angle = est_rotations[0, nbr] * 180 / np.pi
+# First retrieve all angles for the `review_class` (original_img_0_idx),
+#   then lookup the specific neighbor `nbr`
+assert (
+    original_img_0_idx == review_class
+), "DebugClassAvgSource should retain original source image ordering"
+angle = est_rotations[original_img_0_idx][nbr] * 180 / np.pi
 if reflections[nbr]:
     print("Reflection reported.")
     original_img_nbr = np.flipud(original_img_nbr)

src/aspire/classification/averager2d.py

@@ -84,7 +84,7 @@ def average(
 
         :param classes: class indices, refering to src. (src.n, n_nbor).
         :param reflections: Bool representing whether to reflect image in `classes`.
-            (n_clases, n_nbor)
+            (n_classes, n_nbor)
         :param coefs: Optional basis coefs (could avoid recomputing).
             (src.n, coef_count)
         :return: Stack of synthetic class average images as Image instance.
@@ -152,6 +152,16 @@ def __init__(
                 f"{self.__class__.__name__}'s composite_basis {self.composite_basis} must provide a `shift` method."
             )
 
+        # Instantiate dicts to hold alignment results.
+        # Note dicts are used in place of arrays because:
+        #   The entire set of `src.n` classes may not always need to be computed,
+        #   and the order/batching where results are computed is potentially arbitrary.
+        #   We may not know apriori how many nbors are in each class,
+        #   and this may be variable with future methods.
+        self.rotations = dict()
+        self.shifts = dict()
+        self.dot_products = dict()
+
     @abstractmethod
     def align(self, classes, reflections, basis_coefficients=None):
         """
@@ -192,9 +202,15 @@ def average(
         classes = np.atleast_2d(classes)
         reflections = np.atleast_2d(reflections)
 
-        self.rotations, self.shifts, self.dot_products = self.align(
-            classes, reflections, coefs
-        )
+        rotations, shifts, dot_products = self.align(classes, reflections, coefs)
+
+        # Assign batch results
+        src_indices = classes[:, 0]  # First column of class table
+        for i, k in enumerate(src_indices):
+            self.rotations[k] = rotations[i]
+            if shifts is not None:
+                self.shifts[k] = shifts[i]
+            self.dot_products[k] = dot_products[i]
 
         n_classes, n_nbor = classes.shape
 
@@ -212,29 +228,29 @@ def _innerloop(i):
                 neighbors_imgs = Image(self._cls_images(classes[i]))
 
                 # Do shifts
-                if self.shifts is not None:
-                    neighbors_imgs = neighbors_imgs.shift(self.shifts[i])
+                if shifts is not None:
+                    neighbors_imgs = neighbors_imgs.shift(shifts[i])
 
                 neighbors_coefs = self.composite_basis.evaluate_t(neighbors_imgs)
             else:
                 # Get the neighbors
                 neighbors_ids = classes[i]
                 neighbors_coefs = coefs[neighbors_ids]
-                if self.shifts is not None:
+                if shifts is not None:
                     neighbors_coefs = self.composite_basis.shift(
-                        neighbors_coefs, self.shifts[i]
+                        neighbors_coefs, shifts[i]
                     )
 
             # Rotate in composite_basis
             neighbors_coefs = self.composite_basis.rotate(
-                neighbors_coefs, self.rotations[i], reflections[i]
+                neighbors_coefs, rotations[i], reflections[i]
             )
 
             # Averaging in composite_basis
             return self.image_stacker(neighbors_coefs.asnumpy())
 
-        desc = f"Stacking and evaluating class averages from {self.composite_basis.__class__.__name__} to Cartesian"
-        for start in trange(0, n_classes, self.batch_size, desc=desc):
+        desc = f"Stacking and evaluating batch of class averages from {self.composite_basis.__class__.__name__} to Cartesian"
+        for start in trange(0, n_classes, self.batch_size, desc=desc, leave=False):
             end = min(start + self.batch_size, n_classes)
             for i, cls in enumerate(
                 trange(start, end, desc="Stacking batch", leave=False)
@@ -362,7 +378,7 @@ def align(self, classes, reflections, basis_coefficients=None):
         # This is done primarily in case of a tie later, we would take unshifted.
         test_shifts = self._shift_search_grid(self.src.L, self.radius, roll_zero=True)
 
-        for k in trange(n_classes, desc="Rotationally aligning classes"):
+        for k in trange(n_classes, desc="Rotationally aligning classes", leave=False):
             # We want to locally cache the original images,
             #  because we will mutate them with shifts in the next loop.
             #  This avoids recomputing them before each shift
@@ -564,7 +580,7 @@ def _innerloop(k):
                 dtype=self.dtype,
             )
 
-        for k in trange(n_classes, desc="Rotationally aligning classes"):
+        for k in trange(n_classes, desc="Rotationally aligning classes", leave=False):
             rotations[k], shifts[k], dot_products[k] = _innerloop(k)
 
         return rotations, shifts, dot_products
@@ -580,9 +596,15 @@ def average(
         Otherwise is similar to `AligningAverager2D.average`.
         """
 
-        self.rotations, self.shifts, self.dot_products = self.align(
-            classes, reflections, coefs
-        )
+        rotations, shifts, dot_products = self.align(classes, reflections, coefs)
+
+        # Assign batch results
+        src_indices = classes[:, 0]  # First column of class table
+        for i, k in enumerate(src_indices):
+            self.rotations[k] = rotations[i]
+            if shifts is not None:
+                self.shifts[k] = shifts[i]
+            self.dot_products[k] = dot_products[i]
 
         n_classes, n_nbor = classes.shape
 
@@ -601,19 +623,17 @@ def _innerloop(i):
 
             # Rotate in composite_basis
             neighbors_coefs = self.composite_basis.rotate(
-                neighbors_coefs, self.rotations[i], reflections[i]
+                neighbors_coefs, rotations[i], reflections[i]
             )
 
             # Note shifts are after rotation for this approach!
-            if self.shifts is not None:
-                neighbors_coefs = self.composite_basis.shift(
-                    neighbors_coefs, self.shifts[i]
-                )
+            if shifts is not None:
+                neighbors_coefs = self.composite_basis.shift(neighbors_coefs, shifts[i])
 
             # Averaging in composite_basis
             return self.image_stacker(neighbors_coefs.asnumpy())
 
-        for i in trange(n_classes, desc="Stacking class averages"):
+        for i in trange(n_classes, desc="Stacking class averages", leave=False):
             b_avgs[i] = _innerloop(i)
 
         # Now we convert the averaged images from Basis to Cartesian.
@@ -732,7 +752,7 @@ def _innerloop(k):
 
             return _rotations, _shifts, _dot_products
 
-        for k in trange(n_classes, desc="Rotationally aligning classes"):
+        for k in trange(n_classes, desc="Rotationally aligning classes", leave=False):
             rotations[k], shifts[k], dot_products[k] = _innerloop(k)
 
         return rotations, shifts, dot_products
@@ -880,7 +900,7 @@ def align(self, classes, reflections, basis_coefficients=None):
         )
         _images = xp.empty((n_nbor - 1, self.src.L, self.src.L), dtype=self.dtype)
 
-        for k in trange(n_classes, desc="Rotationally aligning classes"):
+        for k in trange(n_classes, desc="Rotationally aligning classes", leave=False):
             # We want to locally cache the original images,
             #  because we will mutate them with shifts in the next loop.
             #  This avoids recomputing them before each shift

src/aspire/denoising/class_avg.py

@@ -223,6 +223,7 @@ def _class_select(self):
             self._classify()
 
         # Perform class selection
+        logger.info("Performing class selection")
         _selection_indices = self.class_selector.select(
             self.class_indices,
             self.class_refl,