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) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index d982e746e5..3bbaff05b2 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,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 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,