From a1561993cccbfc744e226afd7890fd47476a76b9 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 4 Dec 2025 07:51:09 -0500 Subject: [PATCH 1/4] update class alignment attrs fixed bug introduced via batching --- src/aspire/classification/averager2d.py | 54 +++++++++++++++++-------- 1 file changed, 37 insertions(+), 17 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index d982e746e5..301bc9bc37 100644 --- a/src/aspire/classification/averager2d.py +++ b/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,22 +228,22 @@ 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 @@ -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,14 +623,12 @@ 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()) From d980c91dbd1978c5e069de42d84399872d4d7972 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 4 Dec 2025 09:27:29 -0500 Subject: [PATCH 2/4] leave progress False --- src/aspire/classification/averager2d.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 301bc9bc37..3bbaff05b2 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -249,8 +249,8 @@ def _innerloop(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) @@ -378,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 @@ -580,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 @@ -633,7 +633,7 @@ def _innerloop(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. @@ -752,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 @@ -900,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 From 0535e1a1bee1d02b11cf3ba3e830b90e3e6f0751 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 4 Dec 2025 09:27:40 -0500 Subject: [PATCH 3/4] add selection message --- src/aspire/denoising/class_avg.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py index 564549342d..646cdc9e2a 100644 --- a/src/aspire/denoising/class_avg.py +++ b/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, From 59cf2becd17c5d86047e36082bb765a92b53f87c Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 5 Dec 2025 08:26:22 -0500 Subject: [PATCH 4/4] Update class avg tutorial documentation --- gallery/tutorials/tutorials/class_averaging.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/gallery/tutorials/tutorials/class_averaging.py b/gallery/tutorials/tutorials/class_averaging.py index 7108a97846..063e10c1e0 100644 --- a/gallery/tutorials/tutorials/class_averaging.py +++ b/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)