From 16ec967e659e942eb2a354fb12d407107b75d665 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 10 Dec 2024 15:10:44 -0500 Subject: [PATCH 001/216] Initial port of LUD w/ spectral norm constraint. --- src/aspire/abinitio/__init__.py | 1 + src/aspire/abinitio/commonline_lud.py | 448 +++++++++++++++++++++++++- 2 files changed, 435 insertions(+), 14 deletions(-) diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index e8115ea185..1848f2b572 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -1,4 +1,5 @@ from .commonline_base import CLOrient3D +from .commonline_lud import CommonlineLUD from .commonline_sdp import CommonlineSDP from .sync_voting import SyncVotingMixin diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 4a3ece1dfe..6ca1645818 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -1,32 +1,452 @@ import logging +import numpy as np +from scipy.sparse import csr_array +from scipy.sparse.linalg import eigs + from aspire.abinitio import CLOrient3D +from aspire.utils import nearest_rotations +from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) -class CommLineLUD(CLOrient3D): +class CommonlineLUD(CLOrient3D): """ - Define a derived class to estimate 3D orientations using Least Unsquared Deviations described as below: - L. Wang, A. Singer, and Z. Wen, Orientation Determination of Cryo-EM Images Using Least Unsquared Deviations, - SIAM J. Imaging Sciences, 6, 2450-2483 (2013). - + Define a derived class to estimate 3D orientations using Least Unsquared + Deviations described as below: + L. Wang, A. Singer, and Z. Wen, Orientation Determination of Cryo-EM Images Using + Least Unsquared Deviations, SIAM J. Imaging Sciences, 6, 2450-2483 (2013). """ - def __init__(self, src): + def __init__(self, *args, **kwargs): + # Call the parent class initializer + super().__init__(*args, **kwargs) + + # Handle additional parameters specific to CommonlineLUD + self.alpha = kwargs.get("alpha", 2 / 3) + self.tol = kwargs.get("tol", 1e-3) + self.mu = kwargs.get("mu", 1) + self.gam = kwargs.get("gam", 1.618) + self.EPS = kwargs.get("EPS", 1e-12) + self.maxit = kwargs.get("maxit", 1000) + self.adp_proj = kwargs.get("adp_proj", 1) + self.max_rankZ = kwargs.get("max_rankZ", None) + self.max_rankW = kwargs.get("max_rankW", None) + + # Parameters for adjusting mu + self.adp_mu = kwargs.get("adp_mu", 1) + self.dec_mu = kwargs.get("dec_mu", 0.5) + self.inc_mu = kwargs.get("inc_mu", 2) + self.mu_min = kwargs.get("mu_min", 1e-4) + self.mu_max = kwargs.get("mu_max", 1e4) + self.min_mu_itr = kwargs.get("min_mu_itr", 5) + self.max_mu_itr = kwargs.get("max_mu_itr", 20) + self.delta_mu_l = kwargs.get("delta_mu_l", 0.1) + self.delta_mu_u = kwargs.get("delta_mu_u", 10) + + def estimate_rotations(self): + """ + Estimate rotation matrices using the common lines method with semi-definite programming. + """ + logger.info("Computing the common lines matrix.") + self.build_clmatrix() + + C = self.cl_to_C(self.clmatrix) + gram, _, _, _, _ = self.cryoEMSDPL12N(C) + gram = self._restructure_Gram(gram) + self.rotations = self._deterministic_rounding(gram) + + return self.rotations + + def cryoEMSDPL12N(self, C): + # Initialize problem parameters + lambda_ = self.alpha * self.n_img + n = 2 * self.n_img + b = np.concatenate([np.ones(n), np.zeros(self.n_img)]) + + # Adjust rank limits + self.max_rankZ = self.max_rankZ or max(6, self.n_img // 2) + self.max_rankW = self.max_rankW or max(6, self.n_img // 2) + + # Initialize variables + G = np.eye(n, dtype=self.dtype) + W = np.eye(n, dtype=self.dtype) + Z = W + Phi = G / self.mu + + # Compute initial values + S, theta = self.Qtheta(Phi, C, self.mu) + S = (S + S.T) / 2 + AS = self.ComputeAX(S) + resi = self.ComputeAX(G) - b + + kk = 0 + nev = 0 + + itmu_pinf = 0 + itmu_dinf = 0 + + # Maybe initialize values for first iteration prior to loop. + zz = 0 # take this out later + dH = 0 # take this out later + for itr in range(self.maxit): + y = -(AS + self.ComputeAX(W) - self.ComputeAX(Z)) - resi / self.mu + ATy = self.ComputeATy(y) + Phi = W + ATy - Z + G / self.mu + S, theta = self.Qtheta(Phi, C, self.mu) + S = (S + S.T) / 2 + + B = S + W + ATy + G / self.mu + B = (B + B.T) / 2 + + if self.adp_proj == 0: + U, pi = np.linalg.eigh(B) + else: + if itr == 0: + kk = self.max_rankZ + else: + if kk > 0: + # Weird logic here to account for matlab behavior + # for empty array or divide by zero. + rel_drp = 0 + if len(zz) == 2: + rel_drp = np.inf + if len(zz) > 2: + drops = zz[:-1] / zz[1:] + dmx, imx = max( + (val, idx) for idx, val in enumerate(drops) + ) # Find max drop and its index + rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) + if rel_drp > 10: + kk = max(imx, 6) + else: + kk += 3 + else: + kk = 6 + + kk = min(kk, n) + pi, U = eigs( + B, k=kk, which="LM" + ) # Compute top `kk` eigenvalues and eigenvectors + + # Sort by eigenvalue magnitude. + idx = np.argsort(np.abs(pi))[::-1] + pi = pi[idx] + U = U[:, idx].real + pi = pi.real # Ensure real eigenvalues for subsequent calculations + + zz = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) + nD = zz > 0 + kk = np.count_nonzero(nD) + if kk > 0: + zz = zz[nD] + Z = U[:, nD] @ np.diag(zz) @ U[:, nD].T + else: + Z = np.zeros_like(B) + + H = Z - S - ATy - G / self.mu + H = (H + H.T) / 2 + + if self.adp_proj == 0: + D, V = np.linalg.eigh(H) + D = np.diag(D) + W = V[:, D > self.EPS] @ np.diag(D[D > self.EPS]) @ V[:, D > self.EPS].T + else: + if itr == 0: + nev = self.max_rankW + else: + if nev > 0: + drops = dH[:-1] / dH[1:] + dmx, imx = max((val, idx) for idx, val in enumerate(drops)) + rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) + + if rel_drp > 50: + nev = max(imx + 1, 6) + else: + nev = nev + 5 + else: + nev = 6 + + dH, V = eigs(-H, k=min(nev, n), which="LR") + + # Sort by eigenvalue magnitude. + dH = dH.real + idx = np.argsort(dH)[::-1] + dH = dH[idx] + V = V[:, idx].real + nD = dH > self.EPS + dH = dH[nD] + nev = np.count_nonzero(nD) + W = V[:, nD] @ np.diag(dH) @ V[:, nD].T + H if nD.any() else H + + G = (1 - self.gam) * G + self.gam * self.mu * (W - H) + resi = self.ComputeAX(G) - b + pinf = np.linalg.norm(resi) / max(np.linalg.norm(b), 1) + dinf = np.linalg.norm(S + W + ATy - Z, "fro") / max( + np.linalg.norm(S, np.inf), 1 + ) + + if max(pinf, dinf) <= self.tol: + return ( + G, + Z, + W, + y, + {"exit": "optimal", "itr": itr, "pinf": pinf, "dinf": dinf}, + ) + + if self.adp_mu: + if pinf / dinf <= self.delta_mu_l: + itmu_pinf = itmu_pinf + 1 + itmu_dinf = 0 + if itmu_pinf > self.max_mu_itr: + self.mu = max(self.mu * self.inc_mu, self.mu_min) + itmu_pinf = 0 + elif pinf / dinf > self.delta_mu_u: + itmu_dinf = itmu_dinf + 1 + itmu_pinf = 0 + if itmu_dinf > self.max_mu_itr: + self.mu = min(self.mu * self.dec_mu, self.mu_max) + itmu_dinf = 0 + + return G, Z, W, y, {"exit": "max_iter_reached", "itr": self.maxit} + + def Qtheta(self, phi, C, mu): + """ + Python equivalent of Qtheta MEX function. + """ + # Initialize outputs + S = np.zeros((2 * self.n_img, 2 * self.n_img)) + theta = np.zeros_like(C) + + # Main routine + for i in range(self.n_img - 1): + for j in range(i + 1, self.n_img): + t = 0 + for k in range(2): + theta[i, j, k] = C[i, j, k] - mu * ( + phi[2 * i + k, 2 * j] * C[j, i, 0] + + phi[2 * i + k, 2 * j + 1] * C[j, i, 1] + ) + t += theta[i, j, k] ** 2 + + t = np.sqrt(t) + for k in range(2): + theta[i, j, k] /= t + S[2 * i + k, 2 * j] = theta[i, j, k] * C[j, i, 0] + S[2 * i + k, 2 * j + 1] = theta[i, j, k] * C[j, i, 1] + + return S, theta + + def cl_to_C(self, clmatrix): + C = np.zeros((self.n_img, self.n_img, 2), dtype=self.dtype) + for i in range(self.n_img): + for j in range(i + 1, self.n_img): # Only process i < j + cl_ij = clmatrix[i, j] + cl_ji = clmatrix[j, i] + + # Compute (xij, yij) and (xji, yji) from common lines + C[i, j, 0] = np.cos(2 * np.pi * cl_ij / self.n_theta) + C[i, j, 1] = np.sin(2 * np.pi * cl_ij / self.n_theta) + C[j, i, 0] = np.cos(2 * np.pi * cl_ji / self.n_theta) + C[j, i, 1] = np.sin(2 * np.pi * cl_ji / self.n_theta) + + return C + + def ComputeAX(self, X): + n = 2 * self.n_img + rows = np.arange(1, n, 2) + cols = np.arange(0, n, 2) + + # Create diagonal matrix with X on the main diagonal + diags = np.diag(X) + + # Compute the second part of AX + sqrt_2_X_col = np.sqrt(2) * X[rows, cols] + + # Concatenate results vertically + AX = np.concatenate((diags, sqrt_2_X_col)) + + return AX + + def ComputeATy(self, y): + n = 2 * self.n_img + m = 3 * self.n_img + idx = np.arange(n) + rows = np.arange(1, n, 2) + cols = np.arange(0, n, 2) + + ATy = csr_array((n, n), dtype=self.dtype) + ATy[rows, cols] = (np.sqrt(2) / 2) * y[n:m] + ATy = ATy + ATy.T + ATy += csr_array((y[:n], (idx, idx)), shape=(n, n)) + + return ATy.tocsr() + + def _lud_prep(self): + """ + Prepare optimization problem constraints. + + The constraints for the SDP optimization, max tr(SG), performed in `_compute_gram_matrix()` + as min tr(-SG), are that the Gram matrix, G, is semidefinite positive and G11_ii = G22_ii = 1, + G12_ii = G21_ii = 0, i=1,2,...,N, for the block representation of G = [[G11, G12], [G21, G22]]. + + We build a corresponding constraint for CVXPY in the form of tr(A_j @ G) = b_j, j = 1,...,p. + For the constraint G11_ii = G22_ii = 1, we have A_j[i, i] = 1 (zeros elsewhere) and b_j = 1. + For the constraint G12_ii = G21_ii = 0, we have A_j[i, i] = 1 (zeros elsewhere) and b_j = 0. + + :returns: Constraint data A, b. + """ + logger.info("Preparing SDP optimization constraints.") + + n = 2 * self.n_img + A = [] + b = [] + data = np.ones(1, dtype=self.dtype) + for i in range(n): + row_ind = np.array([i]) + col_ind = np.array([i]) + A_i = csr_array((data, (row_ind, col_ind)), shape=(n, n), dtype=self.dtype) + A.append(A_i) + b.append(1) + + for i in range(self.n_img): + row_ind = np.array([i]) + col_ind = np.array([self.n_img + i]) + A_i = csr_array((data, (row_ind, col_ind)), shape=(n, n), dtype=self.dtype) + A.append(A_i) + b.append(0) + + b = np.array(b, dtype=self.dtype) + + return A, b + + def _restructure_Gram(self, G): """ - constructor of an object for estimating 3D orientations + Restructures the input Gram matrix into a block structure based on odd and even + indexed rows and columns. + + The new structure is: + New G = [[Top Left Block, Top Right Block], + [Bottom Left Block, Bottom Right Block]] + + Blocks: + - Top Left Block: Rows and columns with odd indices. + - Top Right Block: Odd rows and even columns. + - Bottom Left Block: Even rows and odd columns. + - Bottom Right Block: Even rows and columns. """ - pass + # Get odd and even indices + odd_indices = np.arange(0, G.shape[0], 2) + even_indices = np.arange(1, G.shape[0], 2) - def estimate(self): + # Extract blocks + top_left = G[np.ix_(odd_indices, odd_indices)] + top_right = G[np.ix_(odd_indices, even_indices)] + bottom_left = G[np.ix_(even_indices, odd_indices)] + bottom_right = G[np.ix_(even_indices, even_indices)] + + # Combine blocks into the new structure + restructured_G = np.block([[top_left, top_right], [bottom_left, bottom_right]]) + + return restructured_G + + def _deterministic_rounding(self, gram): """ - perform estimation of orientations + Deterministic rounding procedure to recover the rotations from the Gram matrix. + + The Gram matrix contains information about the first two columns of every rotation + matrix. These columns are extracted and used to form the remaining column of every + rotation matrix. + + :param gram: A 2n_img x 2n_img Gram matrix. + + :return: An n_img x 3 x 3 stack of rotation matrices. """ - pass + logger.info("Recovering rotations from Gram matrix.") + + # Obtain top eigenvectors from Gram matrix. + d, v = stable_eigsh(gram, 5) + sort_idx = np.argsort(-d) + logger.info(f"Top 5 eigenvalues from (rank-3) Gram matrix: {d[sort_idx]}") + + # Only need the top 3 eigen-vectors. + v = v[:, sort_idx[:3]] + + # According to the structure of the Gram matrix, the first `n_img` rows, denoted v1, + # correspond to the linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of + # column 1 of all rotation matrices. Similarly, the second `n_img` rows of v, + # denoted v2, are linear combinations of R_{i}^{2}, i=1,...,K, that is, the second + # column of all rotation matrices. + v1 = v[: self.n_img].T + v2 = v[self.n_img : 2 * self.n_img].T + + # Use a least-squares method to get A.T*A and a Cholesky decomposition to find A. + A = self._ATA_solver(v1, v2) - def output(self): + # Recover the rotations. The first two columns of all rotation + # matrices are given by unmixing V1 and V2 using A. The third + # column is the cross product of the first two. + r1 = np.dot(A.T, v1) + r2 = np.dot(A.T, v2) + r3 = np.cross(r1, r2, axis=0) + rotations = np.stack((r1.T, r2.T, r3.T), axis=-1) + + # Make sure that we got rotations by enforcing R to be + # a rotation (in case the error is large) + rotations = nearest_rotations(rotations) + + return rotations + + @staticmethod + def _ATA_solver(v1, v2): """ - Output the 3D orientations + Uses a least squares method to solve for the linear transformation A + such that A*v1=R1 and A*v2=R2 correspond to the first and second columns + of a sequence of rotation matrices. + + :param v1: 3 x n_img array corresponding to linear combinations of the first + columns of all rotation matrices. + :param v2: 3 x n_img array corresponding to linear combinations of the second + columns of all rotation matrices. + + :return: 3x3 linear transformation mapping v1, v2 to first two columns of rotations. """ - pass + # We look for a linear transformation (3 x 3 matrix) A such that + # A*v1'=R1 and A*v2=R2 are the columns of the rotations matrices. + # Therefore: + # v1 * A'*A v1' = 1 + # v2 * A'*A v2' = 1 + # v1 * A'*A v2' = 0 + # These are 3*K linear equations for 9 matrix entries of A'*A + # Actually, there are only 6 unknown variables, because A'*A is symmetric. + # So we will truncate from 9 variables to 6 variables corresponding + # to the upper half of the matrix A'*A + n_img = v1.shape[-1] + truncated_equations = np.zeros((3 * n_img, 9), dtype=v1.dtype) + k = 0 + for i in range(3): + for j in range(3): + truncated_equations[0::3, k] = v1[i] * v1[j] + truncated_equations[1::3, k] = v2[i] * v2[j] + truncated_equations[2::3, k] = v1[i] * v2[j] + k += 1 + + # b = [1 1 0 1 1 0 ...]' is the right hand side vector + b = np.ones(3 * n_img) + b[2::3] = 0 + + # Find the least squares approximation of A'*A in vector form + ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] + + # Construct the matrix A'*A from the vectorized matrix. + # Note, this is only the lower triangle of A'*A. + ATA = ATA_vec.reshape(3, 3) + + # The Cholesky decomposition of A'*A gives A (lower triangle). + # Note, that `np.linalg.cholesky()` only uses the lower-triangular + # and diagonal elements of ATA. + A = np.linalg.cholesky(ATA) + + return A From f1fb779b566ffd69eae64bc97f085d4201483b96 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 10 Dec 2024 15:21:22 -0500 Subject: [PATCH 002/216] Use pop in init --- src/aspire/abinitio/commonline_lud.py | 42 +++++++++++++-------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 6ca1645818..0a6f53d6ef 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -20,30 +20,30 @@ class CommonlineLUD(CLOrient3D): """ def __init__(self, *args, **kwargs): - # Call the parent class initializer - super().__init__(*args, **kwargs) - # Handle additional parameters specific to CommonlineLUD - self.alpha = kwargs.get("alpha", 2 / 3) - self.tol = kwargs.get("tol", 1e-3) - self.mu = kwargs.get("mu", 1) - self.gam = kwargs.get("gam", 1.618) - self.EPS = kwargs.get("EPS", 1e-12) - self.maxit = kwargs.get("maxit", 1000) - self.adp_proj = kwargs.get("adp_proj", 1) - self.max_rankZ = kwargs.get("max_rankZ", None) - self.max_rankW = kwargs.get("max_rankW", None) + self.alpha = kwargs.pop("alpha", 2 / 3) + self.tol = kwargs.pop("tol", 1e-3) + self.mu = kwargs.pop("mu", 1) + self.gam = kwargs.pop("gam", 1.618) + self.EPS = kwargs.pop("EPS", 1e-12) + self.maxit = kwargs.pop("maxit", 1000) + self.adp_proj = kwargs.pop("adp_proj", 1) + self.max_rankZ = kwargs.pop("max_rankZ", None) + self.max_rankW = kwargs.pop("max_rankW", None) # Parameters for adjusting mu - self.adp_mu = kwargs.get("adp_mu", 1) - self.dec_mu = kwargs.get("dec_mu", 0.5) - self.inc_mu = kwargs.get("inc_mu", 2) - self.mu_min = kwargs.get("mu_min", 1e-4) - self.mu_max = kwargs.get("mu_max", 1e4) - self.min_mu_itr = kwargs.get("min_mu_itr", 5) - self.max_mu_itr = kwargs.get("max_mu_itr", 20) - self.delta_mu_l = kwargs.get("delta_mu_l", 0.1) - self.delta_mu_u = kwargs.get("delta_mu_u", 10) + self.adp_mu = kwargs.pop("adp_mu", 1) + self.dec_mu = kwargs.pop("dec_mu", 0.5) + self.inc_mu = kwargs.pop("inc_mu", 2) + self.mu_min = kwargs.pop("mu_min", 1e-4) + self.mu_max = kwargs.pop("mu_max", 1e4) + self.min_mu_itr = kwargs.pop("min_mu_itr", 5) + self.max_mu_itr = kwargs.pop("max_mu_itr", 20) + self.delta_mu_l = kwargs.pop("delta_mu_l", 0.1) + self.delta_mu_u = kwargs.pop("delta_mu_u", 10) + + # Call the parent class initializer + super().__init__(*args, **kwargs) def estimate_rotations(self): """ From 28c0d315d28e72ed32e7b57c5ebf58cb10f076b8 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 11 Dec 2024 13:54:34 -0500 Subject: [PATCH 003/216] Add initial test file. --- tests/test_orient_lud.py | 87 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 tests/test_orient_lud.py diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py new file mode 100644 index 0000000000..aeb958baef --- /dev/null +++ b/tests/test_orient_lud.py @@ -0,0 +1,87 @@ +import numpy as np +import pytest + +from aspire.abinitio import CommonlineLUD +from aspire.nufft import backend_available +from aspire.source import Simulation +from aspire.utils import ( + Rotation, + get_aligned_rotations, + mean_aligned_angular_distance, + register_rotations, + rots_to_clmatrix, +) +from aspire.volume import AsymmetricVolume + +RESOLUTION = [ + 32, + 33, +] + +OFFSETS = [ + None, # Defaults to random offsets. + 0, +] + +DTYPES = [ + np.float32, + pytest.param(np.float64, marks=pytest.mark.expensive), +] + + +@pytest.fixture(params=RESOLUTION, ids=lambda x: f"resolution={x}") +def resolution(request): + return request.param + + +@pytest.fixture(params=OFFSETS, ids=lambda x: f"offsets={x}") +def offsets(request): + return request.param + + +@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}") +def dtype(request): + return request.param + + +@pytest.fixture +def src_orient_est_fixture(resolution, offsets, dtype): + """Fixture for simulation source and orientation estimation object.""" + src = Simulation( + n=60, + L=resolution, + vols=AsymmetricVolume(L=resolution, C=1, K=100, seed=10).generate(), + offsets=offsets, + amplitudes=1, + seed=0, + ) + + # Increase max_shift and set shift_step to be sub-pixel when using + # random offsets in the Simulation. This improves common-line detection. + max_shift = 0.20 + shift_step = 0.25 + + # Set max_shift 1 pixel and shift_step to 1 pixel when using 0 offsets. + if np.all(src.offsets == 0.0): + max_shift = 1 / src.L + shift_step = 1 + + orient_est = CommonlineLUD( + src, max_shift=max_shift, shift_step=shift_step, mask=False + ) + + return src, orient_est + + +def test_estimate_rotations(src_orient_est_fixture): + src, orient_est = src_orient_est_fixture + + if backend_available("cufinufft") and src.dtype == np.float32: + pytest.skip("CI on GPU fails for singles.") + + orient_est.estimate_rotations() + + # Register estimates to ground truth rotations and compute the + # angular distance between them (in degrees). + # Assert that mean aligned angular distance is less than 3 degrees. + mean_aligned_angular_distance(orient_est.rotations, src.rotations, degree_tol=3) From 8118936ac62be4c6aa1c596c1dffbe70ece44fd5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 11 Dec 2024 13:57:27 -0500 Subject: [PATCH 004/216] remove unused import --- tests/test_orient_lud.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index aeb958baef..b8db2c735a 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -4,13 +4,7 @@ from aspire.abinitio import CommonlineLUD from aspire.nufft import backend_available from aspire.source import Simulation -from aspire.utils import ( - Rotation, - get_aligned_rotations, - mean_aligned_angular_distance, - register_rotations, - rots_to_clmatrix, -) +from aspire.utils import mean_aligned_angular_distance from aspire.volume import AsymmetricVolume RESOLUTION = [ From 246700bc7d871a046411c24ea2cd98299396de6c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 12 Dec 2024 09:11:05 -0500 Subject: [PATCH 005/216] build sparse matrix directly for ComputeATy. --- src/aspire/abinitio/commonline_lud.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 0a6f53d6ef..a899418856 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -274,15 +274,21 @@ def ComputeATy(self, y): n = 2 * self.n_img m = 3 * self.n_img idx = np.arange(n) - rows = np.arange(1, n, 2) - cols = np.arange(0, n, 2) + rows = np.concatenate([np.arange(1, n, 2), np.arange(0, n, 2)]) + cols = np.concatenate([np.arange(0, n, 2), np.arange(1, n, 2)]) + data = np.concatenate([(np.sqrt(2) / 2) * y[n:m], (np.sqrt(2) / 2) * y[n:m]]) + + # Combine diagonal elements + diag_data = y[:n] + diag_idx = idx - ATy = csr_array((n, n), dtype=self.dtype) - ATy[rows, cols] = (np.sqrt(2) / 2) * y[n:m] - ATy = ATy + ATy.T - ATy += csr_array((y[:n], (idx, idx)), shape=(n, n)) + # Construct the full matrix + data = np.concatenate([data, diag_data]) + rows = np.concatenate([rows, diag_idx]) + cols = np.concatenate([cols, diag_idx]) - return ATy.tocsr() + ATy = csr_array((data, (rows, cols)), shape=(n, n)) + return ATy def _lud_prep(self): """ From 08ae3256d443bdfc3a5447be7e494c4d63a35468 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 18 Dec 2024 11:19:36 -0500 Subject: [PATCH 006/216] init docstring --- src/aspire/abinitio/commonline_lud.py | 63 +++++++++++++++++++++------ tests/test_orient_lud.py | 2 +- 2 files changed, 51 insertions(+), 14 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index a899418856..70fbd2d708 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -20,6 +20,51 @@ class CommonlineLUD(CLOrient3D): """ def __init__(self, *args, **kwargs): + """ + Initialize a class for estimating 3D orientations using a Least Unsquared Deviations algorithm. + + This class extends the `CLOrient3D` class, inheriting its initialization parameters. + + :param alpha: Spectral norm constraint for ADMM algorithm. Default is 2/3. + :param tol: Tolerance for convergence. The algorithm stops when conditions reach this threshold. + Default is 1e-3. + :param mu: The penalty parameter (or dual variable scaling factor) in the optimization problem. + Default is 1. + :param gam: Relaxation factor for updating variables in the algorithm (typically between 1 and 2). + Default is 1.618. + :param EPS: Small positive value used to filter out negligible eigenvalues or avoid numerical issues. + Default is 1e-12. + :param maxit: Maximum number of iterations allowed for the algorithm. + Default is 1000. + :param adp_proj: Flag for using adaptive projection during eigenvalue computation: + - 0: Full eigenvalue decomposition. + - 1: Adaptive rank selection (Default). + :param max_rankZ: Maximum rank used for projecting the Z matrix (for adaptive projection). + Default is None (will be computed based on `n_img`). + :param max_rankW: Maximum rank used for projecting the W matrix (for adaptive projection). + Default is None (will be computed based on `n_img`). + :param adp_mu: Adaptive adjustment of the penalty parameter `mu`: + - 1: Enabled. + - 0: Disabled. + Default is 1. + :param dec_mu: Scaling factor for decreasing `mu` when conditions warrant. + Default is 0.5. + :param inc_mu: Scaling factor for increasing `mu` when conditions warrant. + Default is 2. + :param mu_min: Minimum allowable value for `mu`. + Default is 1e-4. + :param mu_max: Maximum allowable value for `mu`. + Default is 1e4. + :param min_mu_itr: Minimum number of iterations before `mu` is adjusted. + Default is 5. + :param max_mu_itr: Maximum number of iterations allowed for `mu` adjustment. + Default is 20. + :param delta_mu_l: Lower bound for relative drop ratio to trigger a decrease in `mu`. + Default is 0.1. + :param delta_mu_u: Upper bound for relative drop ratio to trigger an increase in `mu`. + Default is 10. + """ + # Handle additional parameters specific to CommonlineLUD self.alpha = kwargs.pop("alpha", 2 / 3) self.tol = kwargs.pop("tol", 1e-3) @@ -53,7 +98,7 @@ def estimate_rotations(self): self.build_clmatrix() C = self.cl_to_C(self.clmatrix) - gram, _, _, _, _ = self.cryoEMSDPL12N(C) + gram = self.cryoEMSDPL12N(C) gram = self._restructure_Gram(gram) self.rotations = self._deterministic_rounding(gram) @@ -81,7 +126,6 @@ def cryoEMSDPL12N(self, C): AS = self.ComputeAX(S) resi = self.ComputeAX(G) - b - kk = 0 nev = 0 itmu_pinf = 0 @@ -150,7 +194,6 @@ def cryoEMSDPL12N(self, C): if self.adp_proj == 0: D, V = np.linalg.eigh(H) - D = np.diag(D) W = V[:, D > self.EPS] @ np.diag(D[D > self.EPS]) @ V[:, D > self.EPS].T else: if itr == 0: @@ -188,13 +231,7 @@ def cryoEMSDPL12N(self, C): ) if max(pinf, dinf) <= self.tol: - return ( - G, - Z, - W, - y, - {"exit": "optimal", "itr": itr, "pinf": pinf, "dinf": dinf}, - ) + return G if self.adp_mu: if pinf / dinf <= self.delta_mu_l: @@ -210,7 +247,7 @@ def cryoEMSDPL12N(self, C): self.mu = min(self.mu * self.dec_mu, self.mu_max) itmu_dinf = 0 - return G, Z, W, y, {"exit": "max_iter_reached", "itr": self.maxit} + return G def Qtheta(self, phi, C, mu): """ @@ -294,7 +331,7 @@ def _lud_prep(self): """ Prepare optimization problem constraints. - The constraints for the SDP optimization, max tr(SG), performed in `_compute_gram_matrix()` + The constraints for the LUD optimization, max tr(SG), performed in `_compute_gram_matrix()` as min tr(-SG), are that the Gram matrix, G, is semidefinite positive and G11_ii = G22_ii = 1, G12_ii = G21_ii = 0, i=1,2,...,N, for the block representation of G = [[G11, G12], [G21, G22]]. @@ -304,7 +341,7 @@ def _lud_prep(self): :returns: Constraint data A, b. """ - logger.info("Preparing SDP optimization constraints.") + logger.info("Preparing LUD optimization constraints.") n = 2 * self.n_img A = [] diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index b8db2c735a..11a3b6f198 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -61,7 +61,7 @@ def src_orient_est_fixture(resolution, offsets, dtype): shift_step = 1 orient_est = CommonlineLUD( - src, max_shift=max_shift, shift_step=shift_step, mask=False + src, max_shift=max_shift, shift_step=shift_step, mask=False, ) return src, orient_est From 32fe7a652fd4a7f089f66a9f7f209b527a385536 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 18 Dec 2024 13:54:30 -0500 Subject: [PATCH 007/216] reformat --- tests/test_orient_lud.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 11a3b6f198..89ea1d4dc1 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -61,7 +61,10 @@ def src_orient_est_fixture(resolution, offsets, dtype): shift_step = 1 orient_est = CommonlineLUD( - src, max_shift=max_shift, shift_step=shift_step, mask=False, + src, + max_shift=max_shift, + shift_step=shift_step, + mask=False, ) return src, orient_est From 2f44d94cbf16419bb83df4e3fddc2773e17b898e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 6 Jan 2025 15:54:02 -0500 Subject: [PATCH 008/216] docstrings --- src/aspire/abinitio/commonline_lud.py | 46 ++++++++++++++++++++++++--- 1 file changed, 42 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 70fbd2d708..23512d99e9 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -14,7 +14,7 @@ class CommonlineLUD(CLOrient3D): """ Define a derived class to estimate 3D orientations using Least Unsquared - Deviations described as below: + Deviations as described in the following publication: L. Wang, A. Singer, and Z. Wen, Orientation Determination of Cryo-EM Images Using Least Unsquared Deviations, SIAM J. Imaging Sciences, 6, 2450-2483 (2013). """ @@ -105,8 +105,19 @@ def estimate_rotations(self): return self.rotations def cryoEMSDPL12N(self, C): + """ + Perform the alternating direction method of multipliers (ADMM) for the SDP + problem: + + min sum_{i Date: Tue, 7 Jan 2025 14:33:51 -0500 Subject: [PATCH 009/216] code comments --- src/aspire/abinitio/commonline_lud.py | 41 +++++++++++++++++++-------- 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 23512d99e9..19387e155e 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -138,20 +138,27 @@ def cryoEMSDPL12N(self, C): resi = self.ComputeAX(G) - b nev = 0 - itmu_pinf = 0 itmu_dinf = 0 - - # Maybe initialize values for first iteration prior to loop. - zz = 0 # take this out later - dH = 0 # take this out later + zz = 0 + dH = 0 for itr in range(self.maxit): + ############# + # Compute y # + ############# y = -(AS + self.ComputeAX(W) - self.ComputeAX(Z)) - resi / self.mu + + ################# + # Compute theta # + ################# ATy = self.ComputeATy(y) Phi = W + ATy - Z + G / self.mu S, theta = self.Qtheta(Phi, C, self.mu) S = (S + S.T) / 2 + ############# + # Compute Z # + ############# B = S + W + ATy + G / self.mu B = (B + B.T) / 2 @@ -162,21 +169,21 @@ def cryoEMSDPL12N(self, C): kk = self.max_rankZ else: if kk > 0: - # Weird logic here to account for matlab behavior - # for empty array or divide by zero. + # Initialize relative drop rel_drp = 0 + + # Calculate relative drop based on `zz` if len(zz) == 2: rel_drp = np.inf - if len(zz) > 2: + elif len(zz) > 2: drops = zz[:-1] / zz[1:] dmx, imx = max( (val, idx) for idx, val in enumerate(drops) ) # Find max drop and its index rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) - if rel_drp > 10: - kk = max(imx, 6) - else: - kk += 3 + + # Update `kk` based on relative drop + kk = max(imx, 6) if rel_drp > 10 else kk + 3 else: kk = 6 @@ -191,6 +198,7 @@ def cryoEMSDPL12N(self, C): U = U[:, idx].real pi = pi.real # Ensure real eigenvalues for subsequent calculations + # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. zz = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) nD = zz > 0 kk = np.count_nonzero(nD) @@ -200,6 +208,9 @@ def cryoEMSDPL12N(self, C): else: Z = np.zeros_like(B) + ############# + # Compute W # + ############# H = Z - S - ATy - G / self.mu H = (H + H.T) / 2 @@ -234,7 +245,12 @@ def cryoEMSDPL12N(self, C): nev = np.count_nonzero(nD) W = V[:, nD] @ np.diag(dH) @ V[:, nD].T + H if nD.any() else H + ############ + # Update G # + ############ G = (1 - self.gam) * G + self.gam * self.mu * (W - H) + + # Check optimality resi = self.ComputeAX(G) - b pinf = np.linalg.norm(resi) / max(np.linalg.norm(b), 1) dinf = np.linalg.norm(S + W + ATy - Z, "fro") / max( @@ -244,6 +260,7 @@ def cryoEMSDPL12N(self, C): if max(pinf, dinf) <= self.tol: return G + # Update mu adaptively if self.adp_mu: if pinf / dinf <= self.delta_mu_l: itmu_pinf = itmu_pinf + 1 From 2be3439438c32ab8f35402b9cd1d3ddb8234a403 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 7 Jan 2025 15:02:48 -0500 Subject: [PATCH 010/216] Move deterministic rounding method to base class --- src/aspire/abinitio/commonline_base.py | 108 ++++++++++++++++++++++++- src/aspire/abinitio/commonline_lud.py | 101 ----------------------- src/aspire/abinitio/commonline_sdp.py | 101 ----------------------- 3 files changed, 107 insertions(+), 203 deletions(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index fe456c645e..f9a221b382 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -7,7 +7,14 @@ from aspire.image import Image from aspire.operators import PolarFT -from aspire.utils import common_line_from_rots, complex_type, fuzzy_mask, tqdm +from aspire.utils import ( + common_line_from_rots, + complex_type, + fuzzy_mask, + nearest_rotations, + tqdm, +) +from aspire.utils.matlab_compat import stable_eigsh from aspire.utils.random import choice logger = logging.getLogger(__name__) @@ -809,3 +816,102 @@ def __init_cupy_module(): backend="nvcc", options=("-O3", "--use_fast_math", "--extra-device-vectorization"), ) + + def _deterministic_rounding(self, gram): + """ + Deterministic rounding procedure to recover the rotations from the Gram matrix. + + The Gram matrix contains information about the first two columns of every rotation + matrix. These columns are extracted and used to form the remaining column of every + rotation matrix. + + :param gram: A 2n_img x 2n_img Gram matrix. + + :return: An n_img x 3 x 3 stack of rotation matrices. + """ + logger.info("Recovering rotations from Gram matrix.") + + # Obtain top eigenvectors from Gram matrix. + d, v = stable_eigsh(gram, 5) + sort_idx = np.argsort(-d) + logger.info(f"Top 5 eigenvalues from (rank-3) Gram matrix: {d[sort_idx]}") + + # Only need the top 3 eigen-vectors. + v = v[:, sort_idx[:3]] + + # According to the structure of the Gram matrix, the first `n_img` rows, denoted v1, + # correspond to the linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of + # column 1 of all rotation matrices. Similarly, the second `n_img` rows of v, + # denoted v2, are linear combinations of R_{i}^{2}, i=1,...,K, that is, the second + # column of all rotation matrices. + v1 = v[: self.n_img].T + v2 = v[self.n_img : 2 * self.n_img].T + + # Use a least-squares method to get A.T*A and a Cholesky decomposition to find A. + A = self._ATA_solver(v1, v2) + + # Recover the rotations. The first two columns of all rotation + # matrices are given by unmixing V1 and V2 using A. The third + # column is the cross product of the first two. + r1 = np.dot(A.T, v1) + r2 = np.dot(A.T, v2) + r3 = np.cross(r1, r2, axis=0) + rotations = np.stack((r1.T, r2.T, r3.T), axis=-1) + + # Make sure that we got rotations by enforcing R to be + # a rotation (in case the error is large) + rotations = nearest_rotations(rotations) + + return rotations + + @staticmethod + def _ATA_solver(v1, v2): + """ + Uses a least squares method to solve for the linear transformation A + such that A*v1=R1 and A*v2=R2 correspond to the first and second columns + of a sequence of rotation matrices. + + :param v1: 3 x n_img array corresponding to linear combinations of the first + columns of all rotation matrices. + :param v2: 3 x n_img array corresponding to linear combinations of the second + columns of all rotation matrices. + + :return: 3x3 linear transformation mapping v1, v2 to first two columns of rotations. + """ + # We look for a linear transformation (3 x 3 matrix) A such that + # A*v1'=R1 and A*v2=R2 are the columns of the rotations matrices. + # Therefore: + # v1 * A'*A v1' = 1 + # v2 * A'*A v2' = 1 + # v1 * A'*A v2' = 0 + # These are 3*K linear equations for 9 matrix entries of A'*A + # Actually, there are only 6 unknown variables, because A'*A is symmetric. + # So we will truncate from 9 variables to 6 variables corresponding + # to the upper half of the matrix A'*A + n_img = v1.shape[-1] + truncated_equations = np.zeros((3 * n_img, 9), dtype=v1.dtype) + k = 0 + for i in range(3): + for j in range(3): + truncated_equations[0::3, k] = v1[i] * v1[j] + truncated_equations[1::3, k] = v2[i] * v2[j] + truncated_equations[2::3, k] = v1[i] * v2[j] + k += 1 + + # b = [1 1 0 1 1 0 ...]' is the right hand side vector + b = np.ones(3 * n_img) + b[2::3] = 0 + + # Find the least squares approximation of A'*A in vector form + ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] + + # Construct the matrix A'*A from the vectorized matrix. + # Note, this is only the lower triangle of A'*A. + ATA = ATA_vec.reshape(3, 3) + + # The Cholesky decomposition of A'*A gives A (lower triangle). + # Note, that `np.linalg.cholesky()` only uses the lower-triangular + # and diagonal elements of ATA. + A = np.linalg.cholesky(ATA) + + return A diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 19387e155e..c6127c1324 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -5,8 +5,6 @@ from scipy.sparse.linalg import eigs from aspire.abinitio import CLOrient3D -from aspire.utils import nearest_rotations -from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) @@ -449,102 +447,3 @@ def _restructure_Gram(self, G): restructured_G = np.block([[top_left, top_right], [bottom_left, bottom_right]]) return restructured_G - - def _deterministic_rounding(self, gram): - """ - Deterministic rounding procedure to recover the rotations from the Gram matrix. - - The Gram matrix contains information about the first two columns of every rotation - matrix. These columns are extracted and used to form the remaining column of every - rotation matrix. - - :param gram: A 2n_img x 2n_img Gram matrix. - - :return: An n_img x 3 x 3 stack of rotation matrices. - """ - logger.info("Recovering rotations from Gram matrix.") - - # Obtain top eigenvectors from Gram matrix. - d, v = stable_eigsh(gram, 5) - sort_idx = np.argsort(-d) - logger.info(f"Top 5 eigenvalues from (rank-3) Gram matrix: {d[sort_idx]}") - - # Only need the top 3 eigen-vectors. - v = v[:, sort_idx[:3]] - - # According to the structure of the Gram matrix, the first `n_img` rows, denoted v1, - # correspond to the linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of - # column 1 of all rotation matrices. Similarly, the second `n_img` rows of v, - # denoted v2, are linear combinations of R_{i}^{2}, i=1,...,K, that is, the second - # column of all rotation matrices. - v1 = v[: self.n_img].T - v2 = v[self.n_img : 2 * self.n_img].T - - # Use a least-squares method to get A.T*A and a Cholesky decomposition to find A. - A = self._ATA_solver(v1, v2) - - # Recover the rotations. The first two columns of all rotation - # matrices are given by unmixing V1 and V2 using A. The third - # column is the cross product of the first two. - r1 = np.dot(A.T, v1) - r2 = np.dot(A.T, v2) - r3 = np.cross(r1, r2, axis=0) - rotations = np.stack((r1.T, r2.T, r3.T), axis=-1) - - # Make sure that we got rotations by enforcing R to be - # a rotation (in case the error is large) - rotations = nearest_rotations(rotations) - - return rotations - - @staticmethod - def _ATA_solver(v1, v2): - """ - Uses a least squares method to solve for the linear transformation A - such that A*v1=R1 and A*v2=R2 correspond to the first and second columns - of a sequence of rotation matrices. - - :param v1: 3 x n_img array corresponding to linear combinations of the first - columns of all rotation matrices. - :param v2: 3 x n_img array corresponding to linear combinations of the second - columns of all rotation matrices. - - :return: 3x3 linear transformation mapping v1, v2 to first two columns of rotations. - """ - # We look for a linear transformation (3 x 3 matrix) A such that - # A*v1'=R1 and A*v2=R2 are the columns of the rotations matrices. - # Therefore: - # v1 * A'*A v1' = 1 - # v2 * A'*A v2' = 1 - # v1 * A'*A v2' = 0 - # These are 3*K linear equations for 9 matrix entries of A'*A - # Actually, there are only 6 unknown variables, because A'*A is symmetric. - # So we will truncate from 9 variables to 6 variables corresponding - # to the upper half of the matrix A'*A - n_img = v1.shape[-1] - truncated_equations = np.zeros((3 * n_img, 9), dtype=v1.dtype) - k = 0 - for i in range(3): - for j in range(3): - truncated_equations[0::3, k] = v1[i] * v1[j] - truncated_equations[1::3, k] = v2[i] * v2[j] - truncated_equations[2::3, k] = v1[i] * v2[j] - k += 1 - - # b = [1 1 0 1 1 0 ...]' is the right hand side vector - b = np.ones(3 * n_img) - b[2::3] = 0 - - # Find the least squares approximation of A'*A in vector form - ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] - - # Construct the matrix A'*A from the vectorized matrix. - # Note, this is only the lower triangle of A'*A. - ATA = ATA_vec.reshape(3, 3) - - # The Cholesky decomposition of A'*A gives A (lower triangle). - # Note, that `np.linalg.cholesky()` only uses the lower-triangular - # and diagonal elements of ATA. - A = np.linalg.cholesky(ATA) - - return A diff --git a/src/aspire/abinitio/commonline_sdp.py b/src/aspire/abinitio/commonline_sdp.py index 11457534ee..b006036be1 100644 --- a/src/aspire/abinitio/commonline_sdp.py +++ b/src/aspire/abinitio/commonline_sdp.py @@ -5,8 +5,6 @@ from scipy.sparse import csr_array from aspire.abinitio import CLOrient3D -from aspire.utils import nearest_rotations -from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) @@ -150,102 +148,3 @@ def _compute_gram_matrix(self, S, A, b): prob.solve() return G.value - - def _deterministic_rounding(self, gram): - """ - Deterministic rounding procedure to recover the rotations from the Gram matrix. - - The Gram matrix contains information about the first two columns of every rotation - matrix. These columns are extracted and used to form the remaining column of every - rotation matrix. - - :param gram: A 2n_img x 2n_img Gram matrix. - - :return: An n_img x 3 x 3 stack of rotation matrices. - """ - logger.info("Recovering rotations from Gram matrix.") - - # Obtain top eigenvectors from Gram matrix. - d, v = stable_eigsh(gram, 5) - sort_idx = np.argsort(-d) - logger.info(f"Top 5 eigenvalues from (rank-3) Gram matrix: {d[sort_idx]}") - - # Only need the top 3 eigen-vectors. - v = v[:, sort_idx[:3]] - - # According to the structure of the Gram matrix, the first `n_img` rows, denoted v1, - # correspond to the linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of - # column 1 of all rotation matrices. Similarly, the second `n_img` rows of v, - # denoted v2, are linear combinations of R_{i}^{2}, i=1,...,K, that is, the second - # column of all rotation matrices. - v1 = v[: self.n_img].T - v2 = v[self.n_img : 2 * self.n_img].T - - # Use a least-squares method to get A.T*A and a Cholesky decomposition to find A. - A = self._ATA_solver(v1, v2) - - # Recover the rotations. The first two columns of all rotation - # matrices are given by unmixing V1 and V2 using A. The third - # column is the cross product of the first two. - r1 = np.dot(A.T, v1) - r2 = np.dot(A.T, v2) - r3 = np.cross(r1, r2, axis=0) - rotations = np.stack((r1.T, r2.T, r3.T), axis=-1) - - # Make sure that we got rotations by enforcing R to be - # a rotation (in case the error is large) - rotations = nearest_rotations(rotations) - - return rotations - - @staticmethod - def _ATA_solver(v1, v2): - """ - Uses a least squares method to solve for the linear transformation A - such that A*v1=R1 and A*v2=R2 correspond to the first and second columns - of a sequence of rotation matrices. - - :param v1: 3 x n_img array corresponding to linear combinations of the first - columns of all rotation matrices. - :param v2: 3 x n_img array corresponding to linear combinations of the second - columns of all rotation matrices. - - :return: 3x3 linear transformation mapping v1, v2 to first two columns of rotations. - """ - # We look for a linear transformation (3 x 3 matrix) A such that - # A*v1'=R1 and A*v2=R2 are the columns of the rotations matrices. - # Therefore: - # v1 * A'*A v1' = 1 - # v2 * A'*A v2' = 1 - # v1 * A'*A v2' = 0 - # These are 3*K linear equations for 9 matrix entries of A'*A - # Actually, there are only 6 unknown variables, because A'*A is symmetric. - # So we will truncate from 9 variables to 6 variables corresponding - # to the upper half of the matrix A'*A - n_img = v1.shape[-1] - truncated_equations = np.zeros((3 * n_img, 9), dtype=v1.dtype) - k = 0 - for i in range(3): - for j in range(3): - truncated_equations[0::3, k] = v1[i] * v1[j] - truncated_equations[1::3, k] = v2[i] * v2[j] - truncated_equations[2::3, k] = v1[i] * v2[j] - k += 1 - - # b = [1 1 0 1 1 0 ...]' is the right hand side vector - b = np.ones(3 * n_img) - b[2::3] = 0 - - # Find the least squares approximation of A'*A in vector form - ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] - - # Construct the matrix A'*A from the vectorized matrix. - # Note, this is only the lower triangle of A'*A. - ATA = ATA_vec.reshape(3, 3) - - # The Cholesky decomposition of A'*A gives A (lower triangle). - # Note, that `np.linalg.cholesky()` only uses the lower-triangular - # and diagonal elements of ATA. - A = np.linalg.cholesky(ATA) - - return A From 99fff8a6811890de13dc4ba60be6be41deffa558 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 9 Jan 2025 09:08:20 -0500 Subject: [PATCH 011/216] Method for spectral norm constraint subproblem, _compute_Z. --- src/aspire/abinitio/commonline_lud.py | 109 ++++++++++++++------------ 1 file changed, 59 insertions(+), 50 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index c6127c1324..12cb088b0f 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -63,7 +63,7 @@ def __init__(self, *args, **kwargs): Default is 10. """ - # Handle additional parameters specific to CommonlineLUD + # Handle parameters specific to CommonlineLUD self.alpha = kwargs.pop("alpha", 2 / 3) self.tol = kwargs.pop("tol", 1e-3) self.mu = kwargs.pop("mu", 1) @@ -115,7 +115,6 @@ def cryoEMSDPL12N(self, C): :return: The gram matrix G. """ # Initialize problem parameters - lambda_ = self.alpha * self.n_img # Spectral norm bound n = 2 * self.n_img b = np.concatenate([np.ones(n), np.zeros(self.n_img)]) @@ -139,6 +138,7 @@ def cryoEMSDPL12N(self, C): itmu_pinf = 0 itmu_dinf = 0 zz = 0 + kk = 0 dH = 0 for itr in range(self.maxit): ############# @@ -157,54 +157,7 @@ def cryoEMSDPL12N(self, C): ############# # Compute Z # ############# - B = S + W + ATy + G / self.mu - B = (B + B.T) / 2 - - if self.adp_proj == 0: - U, pi = np.linalg.eigh(B) - else: - if itr == 0: - kk = self.max_rankZ - else: - if kk > 0: - # Initialize relative drop - rel_drp = 0 - - # Calculate relative drop based on `zz` - if len(zz) == 2: - rel_drp = np.inf - elif len(zz) > 2: - drops = zz[:-1] / zz[1:] - dmx, imx = max( - (val, idx) for idx, val in enumerate(drops) - ) # Find max drop and its index - rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) - - # Update `kk` based on relative drop - kk = max(imx, 6) if rel_drp > 10 else kk + 3 - else: - kk = 6 - - kk = min(kk, n) - pi, U = eigs( - B, k=kk, which="LM" - ) # Compute top `kk` eigenvalues and eigenvectors - - # Sort by eigenvalue magnitude. - idx = np.argsort(np.abs(pi))[::-1] - pi = pi[idx] - U = U[:, idx].real - pi = pi.real # Ensure real eigenvalues for subsequent calculations - - # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. - zz = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) - nD = zz > 0 - kk = np.count_nonzero(nD) - if kk > 0: - zz = zz[nD] - Z = U[:, nD] @ np.diag(zz) @ U[:, nD].T - else: - Z = np.zeros_like(B) + Z, kk, zz = self._compute_Z(S, W, ATy, G, zz, itr, kk, nev) ############# # Compute W # @@ -275,6 +228,62 @@ def cryoEMSDPL12N(self, C): return G + def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): + """ + Update ADMM subproblem for enforcing the spectral norm constraint. + """ + lambda_ = self.alpha * self.n_img # Spectral norm bound + B = S + W + ATy + G / self.mu + B = (B + B.T) / 2 + + if self.adp_proj == 0: + U, pi = np.linalg.eigh(B) + else: + if itr == 0: + kk = self.max_rankZ + else: + if kk > 0: + # Initialize relative drop + rel_drp = 0 + imx = 0 + # Calculate relative drop based on `zz` + if len(zz) == 2: + rel_drp = np.inf + elif len(zz) > 2: + drops = zz[:-1] / zz[1:] + dmx, imx = max( + (val, idx) for idx, val in enumerate(drops) + ) # Find max drop and its index + rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) + + # Update `kk` based on relative drop + kk = max(imx, 6) if rel_drp > 10 else kk + 3 + else: + kk = 6 + + kk = min(kk, 2 * self.n_img) + pi, U = eigs( + B, k=kk, which="LM" + ) # Compute top `kk` eigenvalues and eigenvectors + + # Sort by eigenvalue magnitude. + idx = np.argsort(np.abs(pi))[::-1] + pi = pi[idx] + U = U[:, idx].real + pi = pi.real # Ensure real eigenvalues for subsequent calculations + + # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. + zz = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) + nD = zz > 0 + kk = np.count_nonzero(nD) + if kk > 0: + zz = zz[nD] + Z = U[:, nD] @ np.diag(zz) @ U[:, nD].T + else: + Z = np.zeros_like(B) + + return Z, kk, zz + def Qtheta(self, phi, C, mu): """ Python equivalent of Qtheta MEX function. From 7a911d5611bd2495fad9f0df727f8f1ea6085f20 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 9 Jan 2025 16:02:15 -0500 Subject: [PATCH 012/216] Logic for ADMM w/o spectral norm constraint. Still needs validating against matlab. --- src/aspire/abinitio/commonline_lud.py | 65 +++++++++++++++++++++------ tests/test_orient_lud.py | 13 +++++- 2 files changed, 63 insertions(+), 15 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 12cb088b0f..eb40624db9 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -23,7 +23,9 @@ def __init__(self, *args, **kwargs): This class extends the `CLOrient3D` class, inheriting its initialization parameters. - :param alpha: Spectral norm constraint for ADMM algorithm. Default is 2/3. + :param alpha: Spectral norm constraint for ADMM algorithm. Default is None, which + does not apply a spectral norm constraint. To apply a spectral norm constraint provide + a value in the range [2/3, 1), 2/3 is recommended. :param tol: Tolerance for convergence. The algorithm stops when conditions reach this threshold. Default is 1e-3. :param mu: The penalty parameter (or dual variable scaling factor) in the optimization problem. @@ -64,7 +66,23 @@ def __init__(self, *args, **kwargs): """ # Handle parameters specific to CommonlineLUD - self.alpha = kwargs.pop("alpha", 2 / 3) + self.alpha = kwargs.pop("alpha", None) # Spectral norm constraint bound + if self.alpha is not None: + if not (2 / 3 <= self.alpha < 1): + raise ValueError( + "Spectral norm constraint, alpha, must be in [2/3, 1)." + ) + else: + logger.info( + f"Initializing LUD algorithm using ADMM with spectral norm constraint {self.alpha}." + ) + self.spectral_norm_constraint = True + else: + logger.info( + "Initializing LUD algorithm using ADMM without spectral norm constraint." + ) + self.spectral_norm_constraint = False + self.tol = kwargs.pop("tol", 1e-3) self.mu = kwargs.pop("mu", 1) self.gam = kwargs.pop("gam", 1.618) @@ -119,14 +137,18 @@ def cryoEMSDPL12N(self, C): b = np.concatenate([np.ones(n), np.zeros(self.n_img)]) # Adjust rank limits - self.max_rankZ = self.max_rankZ or max(6, self.n_img // 2) self.max_rankW = self.max_rankW or max(6, self.n_img // 2) + if self.spectral_norm_constraint: + self.max_rankZ = self.max_rankZ or max(6, self.n_img // 2) # Initialize variables G = np.eye(n, dtype=self.dtype) W = np.eye(n, dtype=self.dtype) - Z = W - Phi = G / self.mu + if self.spectral_norm_constraint: + Z = W + Phi = G / self.mu + else: + Phi = W + G / self.mu # Compute initial values S, theta = self.Qtheta(Phi, C, self.mu) @@ -144,25 +166,32 @@ def cryoEMSDPL12N(self, C): ############# # Compute y # ############# - y = -(AS + self.ComputeAX(W) - self.ComputeAX(Z)) - resi / self.mu + y = -(AS + self.ComputeAX(W)) - resi / self.mu + if self.spectral_norm_constraint: + y += self.ComputeAX(Z) ################# # Compute theta # ################# ATy = self.ComputeATy(y) - Phi = W + ATy - Z + G / self.mu + Phi = W + ATy + G / self.mu + if self.spectral_norm_constraint: + Phi -= Z S, theta = self.Qtheta(Phi, C, self.mu) S = (S + S.T) / 2 ############# # Compute Z # ############# - Z, kk, zz = self._compute_Z(S, W, ATy, G, zz, itr, kk, nev) + if self.spectral_norm_constraint: + Z, kk, zz = self._compute_Z(S, W, ATy, G, zz, itr, kk, nev) ############# # Compute W # ############# - H = Z - S - ATy - G / self.mu + H = -S - ATy - G / self.mu + if self.spectral_norm_constraint: + H += Z H = (H + H.T) / 2 if self.adp_proj == 0: @@ -184,7 +213,10 @@ def cryoEMSDPL12N(self, C): else: nev = 6 - dH, V = eigs(-H, k=min(nev, n), which="LR") + n_eigs = nev + if self.spectral_norm_constraint: + n_eigs = min(nev, n) + dH, V = eigs(-H, k=n_eigs, which="LR") # Sort by eigenvalue magnitude. dH = dH.real @@ -204,9 +236,11 @@ def cryoEMSDPL12N(self, C): # Check optimality resi = self.ComputeAX(G) - b pinf = np.linalg.norm(resi) / max(np.linalg.norm(b), 1) - dinf = np.linalg.norm(S + W + ATy - Z, "fro") / max( - np.linalg.norm(S, np.inf), 1 - ) + + dinf_term = S + W + ATy + if self.spectral_norm_constraint: + dinf_term -= Z + dinf = np.linalg.norm(dinf_term, "fro") / max(np.linalg.norm(S, np.inf), 1) if max(pinf, dinf) <= self.tol: return G @@ -325,7 +359,10 @@ def Qtheta(self, phi, C, mu): t = np.sqrt(t) for k in range(2): - theta[i, j, k] /= t + if self.spectral_norm_constraint: + theta[i, j, k] /= t + else: + theta[i, j, k] /= max(t, self.mu) S[2 * i + k, 2 * j] = theta[i, j, k] * C[j, i, 0] S[2 * i + k, 2 * j + 1] = theta[i, j, k] * C[j, i, 1] diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 89ea1d4dc1..c735773231 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -22,6 +22,11 @@ pytest.param(np.float64, marks=pytest.mark.expensive), ] +SPECTRAL_NORM_CONSTRAINT = [ + 2 / 3, + pytest.param(None, marks=pytest.mark.expensive), +] + @pytest.fixture(params=RESOLUTION, ids=lambda x: f"resolution={x}") def resolution(request): @@ -38,8 +43,13 @@ def dtype(request): return request.param +@pytest.fixture(params=SPECTRAL_NORM_CONSTRAINT, ids=lambda x: f"alpha={x}") +def alpha(request): + return request.param + + @pytest.fixture -def src_orient_est_fixture(resolution, offsets, dtype): +def src_orient_est_fixture(resolution, offsets, dtype, alpha): """Fixture for simulation source and orientation estimation object.""" src = Simulation( n=60, @@ -62,6 +72,7 @@ def src_orient_est_fixture(resolution, offsets, dtype): orient_est = CommonlineLUD( src, + alpha=alpha, max_shift=max_shift, shift_step=shift_step, mask=False, From 9dbcebb2dd43790fc89c6fbda4ba938cbef1b9de Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 10 Jan 2025 15:31:15 -0500 Subject: [PATCH 013/216] snake case --- src/aspire/abinitio/commonline_lud.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index eb40624db9..e713ef8eea 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -151,10 +151,10 @@ def cryoEMSDPL12N(self, C): Phi = W + G / self.mu # Compute initial values - S, theta = self.Qtheta(Phi, C, self.mu) + S, theta = self._Q_theta(Phi, C, self.mu) S = (S + S.T) / 2 - AS = self.ComputeAX(S) - resi = self.ComputeAX(G) - b + AS = self._compute_AX(S) + resi = self._compute_AX(G) - b nev = 0 itmu_pinf = 0 @@ -166,18 +166,18 @@ def cryoEMSDPL12N(self, C): ############# # Compute y # ############# - y = -(AS + self.ComputeAX(W)) - resi / self.mu + y = -(AS + self._compute_AX(W)) - resi / self.mu if self.spectral_norm_constraint: - y += self.ComputeAX(Z) + y += self._compute_AX(Z) ################# # Compute theta # ################# - ATy = self.ComputeATy(y) + ATy = self._compute_ATy(y) Phi = W + ATy + G / self.mu if self.spectral_norm_constraint: Phi -= Z - S, theta = self.Qtheta(Phi, C, self.mu) + S, theta = self._Q_theta(Phi, C, self.mu) S = (S + S.T) / 2 ############# @@ -234,7 +234,7 @@ def cryoEMSDPL12N(self, C): G = (1 - self.gam) * G + self.gam * self.mu * (W - H) # Check optimality - resi = self.ComputeAX(G) - b + resi = self._compute_AX(G) - b pinf = np.linalg.norm(resi) / max(np.linalg.norm(b), 1) dinf_term = S + W + ATy @@ -318,7 +318,7 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): return Z, kk, zz - def Qtheta(self, phi, C, mu): + def _Q_theta(self, phi, C, mu): """ Python equivalent of Qtheta MEX function. @@ -390,7 +390,7 @@ def cl_to_C(self, clmatrix): return C - def ComputeAX(self, X): + def _compute_AX(self, X): n = 2 * self.n_img rows = np.arange(1, n, 2) cols = np.arange(0, n, 2) @@ -406,7 +406,7 @@ def ComputeAX(self, X): return AX - def ComputeATy(self, y): + def _compute_ATy(self, y): n = 2 * self.n_img m = 3 * self.n_img idx = np.arange(n) From ebe6d7378a2ff14367d50e95a4b2cdef8b716112 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 13 Jan 2025 16:01:27 -0500 Subject: [PATCH 014/216] cleanup --- src/aspire/abinitio/commonline_lud.py | 101 ++++++++++++++++---------- 1 file changed, 61 insertions(+), 40 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index e713ef8eea..62fe67fbc4 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -21,7 +21,8 @@ def __init__(self, *args, **kwargs): """ Initialize a class for estimating 3D orientations using a Least Unsquared Deviations algorithm. - This class extends the `CLOrient3D` class, inheriting its initialization parameters. + This class extends the `CLOrient3D` class, inheriting its initialization parameters. Additional + parameters detailed below. :param alpha: Spectral norm constraint for ADMM algorithm. Default is None, which does not apply a spectral norm constraint. To apply a spectral norm constraint provide @@ -32,24 +33,23 @@ def __init__(self, *args, **kwargs): Default is 1. :param gam: Relaxation factor for updating variables in the algorithm (typically between 1 and 2). Default is 1.618. - :param EPS: Small positive value used to filter out negligible eigenvalues or avoid numerical issues. + :param EPS: Small positive value used to filter out negligible eigenvalues. Default is 1e-12. :param maxit: Maximum number of iterations allowed for the algorithm. Default is 1000. :param adp_proj: Flag for using adaptive projection during eigenvalue computation: - - 0: Full eigenvalue decomposition. - - 1: Adaptive rank selection (Default). + - True: Adaptive rank selection (Default). + - False: Full eigenvalue decomposition. :param max_rankZ: Maximum rank used for projecting the Z matrix (for adaptive projection). Default is None (will be computed based on `n_img`). :param max_rankW: Maximum rank used for projecting the W matrix (for adaptive projection). Default is None (will be computed based on `n_img`). :param adp_mu: Adaptive adjustment of the penalty parameter `mu`: - - 1: Enabled. - - 0: Disabled. - Default is 1. - :param dec_mu: Scaling factor for decreasing `mu` when conditions warrant. + - True: Enabled (Default). + - False: Disabled. + :param dec_mu: Scaling factor for decreasing `mu`. Default is 0.5. - :param inc_mu: Scaling factor for increasing `mu` when conditions warrant. + :param inc_mu: Scaling factor for increasing `mu`. Default is 2. :param mu_min: Minimum allowable value for `mu`. Default is 1e-4. @@ -88,12 +88,12 @@ def __init__(self, *args, **kwargs): self.gam = kwargs.pop("gam", 1.618) self.EPS = kwargs.pop("EPS", 1e-12) self.maxit = kwargs.pop("maxit", 1000) - self.adp_proj = kwargs.pop("adp_proj", 1) + self.adp_proj = kwargs.pop("adp_proj", True) self.max_rankZ = kwargs.pop("max_rankZ", None) self.max_rankW = kwargs.pop("max_rankW", None) # Parameters for adjusting mu - self.adp_mu = kwargs.pop("adp_mu", 1) + self.adp_mu = kwargs.pop("adp_mu", True) self.dec_mu = kwargs.pop("dec_mu", 0.5) self.inc_mu = kwargs.pop("inc_mu", 2) self.mu_min = kwargs.pop("mu_min", 1e-4) @@ -108,19 +108,19 @@ def __init__(self, *args, **kwargs): def estimate_rotations(self): """ - Estimate rotation matrices using the common lines method with semi-definite programming. + Estimate rotation matrices using the common lines method with LUD optimization. """ logger.info("Computing the common lines matrix.") self.build_clmatrix() - C = self.cl_to_C(self.clmatrix) - gram = self.cryoEMSDPL12N(C) + C = self._cl_to_C(self.clmatrix) + gram = self._compute_Gram(C) gram = self._restructure_Gram(gram) self.rotations = self._deterministic_rounding(gram) return self.rotations - def cryoEMSDPL12N(self, C): + def _compute_Gram(self, C): """ Perform the alternating direction method of multipliers (ADMM) for the SDP problem: @@ -129,7 +129,11 @@ def cryoEMSDPL12N(self, C): s.t. A(G) = b, G psd ||G||_2 <= lambda - :param C: + Equivalent to matlab functions cryoEMSDPL12N/cryoEMSDPL12N_vsimple. + + :param C: ndarray, A 3D array (n_img x n_img x 2) containing commonline coordinates (in Cartesian form) + between pairs of images. Each C[i, j] stores the x and y coordinates of the common + line between image i and image j. :return: The gram matrix G. """ # Initialize problem parameters @@ -194,7 +198,7 @@ def cryoEMSDPL12N(self, C): H += Z H = (H + H.T) / 2 - if self.adp_proj == 0: + if not self.adp_proj: D, V = np.linalg.eigh(H) W = V[:, D > self.EPS] @ np.diag(D[D > self.EPS]) @ V[:, D > self.EPS].T else: @@ -270,7 +274,7 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): B = S + W + ATy + G / self.mu B = (B + B.T) / 2 - if self.adp_proj == 0: + if not self.adp_proj: U, pi = np.linalg.eigh(B) else: if itr == 0: @@ -320,8 +324,6 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): def _Q_theta(self, phi, C, mu): """ - Python equivalent of Qtheta MEX function. - Compute the matrix S and auxiliary variables theta for the optimization problem. This function calculates the fidelity matrix S and the auxiliary variable theta as part of @@ -329,6 +331,8 @@ def _Q_theta(self, phi, C, mu): Least Unsquared Deviations (LUD) problem. It ensures consistency between the Gram matrix G and the detected common-line coordinates. + Python equivalent of matlab Qtheta MEX function. + :param phi: ndarray, A 2*n_img x 2*n_img scaled dual variable matrix (Phi) used in the ADMM iterations. :param C: ndarray, A 3D array (n_img x n_img x 2) containing commonline coordinates (in Cartesian form) between pairs of images. Each C[i, j] stores the x and y coordinates of the common @@ -368,7 +372,7 @@ def _Q_theta(self, phi, C, mu): return S, theta - def cl_to_C(self, clmatrix): + def _cl_to_C(self, clmatrix): """ For each pair of commonline indices cl[i, j] and cl[j, i], convert from polar commonline indices to cartesion coordinates. @@ -391,6 +395,12 @@ def cl_to_C(self, clmatrix): return C def _compute_AX(self, X): + """ + Compute the application of the linear operator A to the input matrix X. + + The operator A extracts diagonal elements of X and computes a scaled + subset of off-diagonal elements, combining them into a single vector. + """ n = 2 * self.n_img rows = np.arange(1, n, 2) cols = np.arange(0, n, 2) @@ -407,16 +417,21 @@ def _compute_AX(self, X): return AX def _compute_ATy(self, y): + """ + Compute the application of the adjoint operator A^T to the input vector y. + + The adjoint operator reconstructs a sparse matrix from the input vector, + placing the values onto the diagonal and selected off-diagonal positions. + """ n = 2 * self.n_img m = 3 * self.n_img - idx = np.arange(n) rows = np.concatenate([np.arange(1, n, 2), np.arange(0, n, 2)]) cols = np.concatenate([np.arange(0, n, 2), np.arange(1, n, 2)]) data = np.concatenate([(np.sqrt(2) / 2) * y[n:m], (np.sqrt(2) / 2) * y[n:m]]) # Combine diagonal elements diag_data = y[:n] - diag_idx = idx + diag_idx = np.arange(n) # Construct the full matrix data = np.concatenate([data, diag_data]) @@ -466,30 +481,36 @@ def _lud_prep(self): def _restructure_Gram(self, G): """ - Restructures the input Gram matrix into a block structure based on odd and even - indexed rows and columns. - - The new structure is: - New G = [[Top Left Block, Top Right Block], - [Bottom Left Block, Bottom Right Block]] - - Blocks: - - Top Left Block: Rows and columns with odd indices. - - Top Right Block: Odd rows and even columns. - - Bottom Left Block: Even rows and odd columns. - - Bottom Right Block: Even rows and columns. + Restructures the input Gram matrix into a block structure based on the following + format: + + .. math:: + + G = + \\begin{pmatrix} + G^{11} & G^{12} \\\\ + G^{21} & G^{22} + \\end{pmatrix} + = + \\begin{pmatrix} + {R^1}^T R^1 & {R^1}^T R^2 \\\\ + {R^2}^T R^1 & {R^2}^T R^2 + \\end{pmatrix} + + :param G: Gram matrix from ADMM method. + :return: Restructured Gram matrix. """ # Get odd and even indices odd_indices = np.arange(0, G.shape[0], 2) even_indices = np.arange(1, G.shape[0], 2) # Extract blocks - top_left = G[np.ix_(odd_indices, odd_indices)] - top_right = G[np.ix_(odd_indices, even_indices)] - bottom_left = G[np.ix_(even_indices, odd_indices)] - bottom_right = G[np.ix_(even_indices, even_indices)] + G_11 = G[np.ix_(odd_indices, odd_indices)] + G_12 = G[np.ix_(odd_indices, even_indices)] + G_21 = G[np.ix_(even_indices, odd_indices)] + G_22 = G[np.ix_(even_indices, even_indices)] # Combine blocks into the new structure - restructured_G = np.block([[top_left, top_right], [bottom_left, bottom_right]]) + restructured_G = np.block([[G_11, G_12], [G_21, G_22]]) return restructured_G From 099336a19255a88bd7b98e7734825a0ed6ce2d1d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 14 Jan 2025 08:41:24 -0500 Subject: [PATCH 015/216] tox --- src/aspire/abinitio/commonline_lud.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 62fe67fbc4..2a3d4de361 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -486,7 +486,7 @@ def _restructure_Gram(self, G): .. math:: - G = + G = \\begin{pmatrix} G^{11} & G^{12} \\\\ G^{21} & G^{22} From dcb771c7802854003a509e2ed35874ee2cf43b84 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 14 Jan 2025 10:35:20 -0500 Subject: [PATCH 016/216] staticmethod and docs for compute_AX/ATy --- src/aspire/abinitio/commonline_lud.py | 63 ++++++++++++++++++++------- 1 file changed, 48 insertions(+), 15 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 2a3d4de361..7d282e842b 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -156,7 +156,6 @@ def _compute_Gram(self, C): # Compute initial values S, theta = self._Q_theta(Phi, C, self.mu) - S = (S + S.T) / 2 AS = self._compute_AX(S) resi = self._compute_AX(G) - b @@ -182,7 +181,6 @@ def _compute_Gram(self, C): if self.spectral_norm_constraint: Phi -= Z S, theta = self._Q_theta(Phi, C, self.mu) - S = (S + S.T) / 2 ############# # Compute Z # @@ -370,6 +368,9 @@ def _Q_theta(self, phi, C, mu): S[2 * i + k, 2 * j] = theta[i, j, k] * C[j, i, 0] S[2 * i + k, 2 * j + 1] = theta[i, j, k] * C[j, i, 1] + # Ensure S is symmetric + S = (S + S.T) / 2 + return S, theta def _cl_to_C(self, clmatrix): @@ -394,16 +395,27 @@ def _cl_to_C(self, clmatrix): return C - def _compute_AX(self, X): + @staticmethod + def _compute_AX(X): """ - Compute the application of the linear operator A to the input matrix X. + Compute the application of the linear operator A to the input matrix X, + where A(X) is defined as: + + A(X) = [ + X_ii^(11), + X_ii^(22), + sqrt(2) X_ii^(12) + sqrt(2) X_ii^(21) + ] - The operator A extracts diagonal elements of X and computes a scaled - subset of off-diagonal elements, combining them into a single vector. + i = 1, 2, ..., K + + where X_{ii}^{pq} denotes the (p,q)-th element in the 2x2 subblock X_{ii}. + + :param X: 2D square array. + :return: A(X) """ - n = 2 * self.n_img - rows = np.arange(1, n, 2) - cols = np.arange(0, n, 2) + rows = np.arange(1, X.shape[0], 2) + cols = np.arange(0, X.shape[0], 2) # Create diagonal matrix with X on the main diagonal diags = np.diag(X) @@ -416,15 +428,36 @@ def _compute_AX(self, X): return AX - def _compute_ATy(self, y): + @staticmethod + def _compute_ATy(y): """ - Compute the application of the adjoint operator A^T to the input vector y. + Compute the application of the adjoint operator A^T to the input vector y, + where - The adjoint operator reconstructs a sparse matrix from the input vector, - placing the values onto the diagonal and selected off-diagonal positions. + y = [ + y_i^1, + y_i^2, + y_i^3 + ] for i = 1, 2, ..., K, + + and the adjoint of the operator A is defined as: + + AT(y) = Y = [ + [Y_ii^(11), Y_ii^(12)], + [Y_ii^(21), Y_ii^(22)] + ], + + where for i = 1, 2, ..., K: + + Y_ii^(11) = y_i^1, + Y_ii^(22) = y_i^2, + Y_ii^(12) = Y_ii^(21) = y_i^3 / sqrt(2). + + :param y: 1D array of length 3 * n_img. + :return: Sparse matrix AT(y) """ - n = 2 * self.n_img - m = 3 * self.n_img + n = 2 * len(y) // 3 + m = len(y) rows = np.concatenate([np.arange(1, n, 2), np.arange(0, n, 2)]) cols = np.concatenate([np.arange(0, n, 2), np.arange(1, n, 2)]) data = np.concatenate([(np.sqrt(2) / 2) * y[n:m], (np.sqrt(2) / 2) * y[n:m]]) From c53090490844338cf6bc19ad661a8b841b84e790 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 14 Jan 2025 14:51:21 -0500 Subject: [PATCH 017/216] dtype pass through --- src/aspire/abinitio/commonline_base.py | 3 +- src/aspire/abinitio/commonline_lud.py | 54 ++++++++++++++++---------- tests/test_orient_lud.py | 12 ++++-- 3 files changed, 44 insertions(+), 25 deletions(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index f9a221b382..9585ea2736 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -830,7 +830,6 @@ def _deterministic_rounding(self, gram): :return: An n_img x 3 x 3 stack of rotation matrices. """ logger.info("Recovering rotations from Gram matrix.") - # Obtain top eigenvectors from Gram matrix. d, v = stable_eigsh(gram, 5) sort_idx = np.argsort(-d) @@ -899,7 +898,7 @@ def _ATA_solver(v1, v2): k += 1 # b = [1 1 0 1 1 0 ...]' is the right hand side vector - b = np.ones(3 * n_img) + b = np.ones(3 * n_img, dtype=v1.dtype) b[2::3] = 0 # Find the least squares approximation of A'*A in vector form diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 7d282e842b..507e0395e7 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -136,9 +136,13 @@ def _compute_Gram(self, C): line between image i and image j. :return: The gram matrix G. """ + logger.info("Performing ADMM to compute Gram matrix.") + # Initialize problem parameters n = 2 * self.n_img - b = np.concatenate([np.ones(n), np.zeros(self.n_img)]) + b = np.concatenate( + [np.ones(n, dtype=self.dtype), np.zeros(self.n_img, dtype=self.dtype)] + ) # Adjust rank limits self.max_rankW = self.max_rankW or max(6, self.n_img // 2) @@ -205,7 +209,12 @@ def _compute_Gram(self, C): else: if nev > 0: drops = dH[:-1] / dH[1:] - dmx, imx = max((val, idx) for idx, val in enumerate(drops)) + + # Find max drop + imx = np.argmax(drops) + dmx = drops[imx] + + # Relative drop rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) if rel_drp > 50: @@ -287,9 +296,12 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): rel_drp = np.inf elif len(zz) > 2: drops = zz[:-1] / zz[1:] - dmx, imx = max( - (val, idx) for idx, val in enumerate(drops) - ) # Find max drop and its index + + # Find max drop + imx = np.argmax(drops) + dmx = drops[imx] + + # Relative drop rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) # Update `kk` based on relative drop @@ -345,7 +357,7 @@ def _Q_theta(self, phi, C, mu): line between image i and image j. """ # Initialize outputs - S = np.zeros((2 * self.n_img, 2 * self.n_img)) + S = np.zeros((2 * self.n_img, 2 * self.n_img), dtype=self.dtype) theta = np.zeros_like(C) # Main routine @@ -421,7 +433,7 @@ def _compute_AX(X): diags = np.diag(X) # Compute the second part of AX - sqrt_2_X_col = np.sqrt(2) * X[rows, cols] + sqrt_2_X_col = np.sqrt(2, dtype=X.dtype) * X[rows, cols] # Concatenate results vertically AX = np.concatenate((diags, sqrt_2_X_col)) @@ -460,7 +472,12 @@ def _compute_ATy(y): m = len(y) rows = np.concatenate([np.arange(1, n, 2), np.arange(0, n, 2)]) cols = np.concatenate([np.arange(0, n, 2), np.arange(1, n, 2)]) - data = np.concatenate([(np.sqrt(2) / 2) * y[n:m], (np.sqrt(2) / 2) * y[n:m]]) + data = np.concatenate( + [ + (np.sqrt(2, dtype=y.dtype) / 2) * y[n:m], + (np.sqrt(2, dtype=y.dtype) / 2) * y[n:m], + ] + ) # Combine diagonal elements diag_data = y[:n] @@ -517,18 +534,15 @@ def _restructure_Gram(self, G): Restructures the input Gram matrix into a block structure based on the following format: - .. math:: - - G = - \\begin{pmatrix} - G^{11} & G^{12} \\\\ - G^{21} & G^{22} - \\end{pmatrix} - = - \\begin{pmatrix} - {R^1}^T R^1 & {R^1}^T R^2 \\\\ - {R^2}^T R^1 & {R^2}^T R^2 - \\end{pmatrix} + G = + [ G^(11) G^(12) ] + [ G^(21) G^(22) ] + + = + [ (R^1)^T R^1 (R^1)^T R^2 ] + [ (R^2)^T R^1 (R^2)^T R^2 ] + + where R^i is the concatenation of all i'th columns of the rotations R. :param G: Gram matrix from ADMM method. :return: Restructured Gram matrix. diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index c735773231..43ef4a67f6 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -54,10 +54,13 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): src = Simulation( n=60, L=resolution, - vols=AsymmetricVolume(L=resolution, C=1, K=100, seed=10).generate(), + vols=AsymmetricVolume( + L=resolution, C=1, K=100, seed=10, dtype=dtype + ).generate(), offsets=offsets, amplitudes=1, seed=0, + dtype=dtype, ) # Increase max_shift and set shift_step to be sub-pixel when using @@ -87,9 +90,12 @@ def test_estimate_rotations(src_orient_est_fixture): if backend_available("cufinufft") and src.dtype == np.float32: pytest.skip("CI on GPU fails for singles.") - orient_est.estimate_rotations() + est_rots = orient_est.estimate_rotations() # Register estimates to ground truth rotations and compute the # angular distance between them (in degrees). # Assert that mean aligned angular distance is less than 3 degrees. - mean_aligned_angular_distance(orient_est.rotations, src.rotations, degree_tol=3) + mean_aligned_angular_distance(est_rots, src.rotations, degree_tol=3) + + # Check dtype pass-through + np.testing.assert_equal(src.dtype, est_rots.dtype) From 2862e8895d8c791b12cc27dd9da43f5ab3532a80 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 15 Jan 2025 09:42:16 -0500 Subject: [PATCH 018/216] try upcast for eigs --- src/aspire/abinitio/commonline_lud.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 507e0395e7..b1be1dc66c 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -103,7 +103,7 @@ def __init__(self, *args, **kwargs): self.delta_mu_l = kwargs.pop("delta_mu_l", 0.1) self.delta_mu_u = kwargs.pop("delta_mu_u", 10) - # Call the parent class initializer + # Initialize commonline base class super().__init__(*args, **kwargs) def estimate_rotations(self): @@ -311,7 +311,7 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): kk = min(kk, 2 * self.n_img) pi, U = eigs( - B, k=kk, which="LM" + B.astype(np.float64, copy=False), k=kk, which="LM" ) # Compute top `kk` eigenvalues and eigenvectors # Sort by eigenvalue magnitude. From 4a0603c4ba524c68504af4fb0d73f8991c11dac3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 15 Jan 2025 10:37:38 -0500 Subject: [PATCH 019/216] one more upcast --- src/aspire/abinitio/commonline_lud.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index b1be1dc66c..ac0aa1bd09 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -227,13 +227,13 @@ def _compute_Gram(self, C): n_eigs = nev if self.spectral_norm_constraint: n_eigs = min(nev, n) - dH, V = eigs(-H, k=n_eigs, which="LR") + dH, V = eigs(-H.astype(np.float64), k=n_eigs, which="LR") # Sort by eigenvalue magnitude. - dH = dH.real + dH = dH.real.astype(self.dtype, copy=False) idx = np.argsort(dH)[::-1] dH = dH[idx] - V = V[:, idx].real + V = V[:, idx].real.astype(self.dtype, copy=False) nD = dH > self.EPS dH = dH[nD] nev = np.count_nonzero(nD) From e4c977aab4faa279a2ac352dda304608ffec8398 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 16 Jan 2025 08:58:13 -0500 Subject: [PATCH 020/216] init parameters --- src/aspire/abinitio/commonline_lud.py | 62 ++++++++++++++++++--------- 1 file changed, 42 insertions(+), 20 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index ac0aa1bd09..7ae8128cda 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -17,7 +17,29 @@ class CommonlineLUD(CLOrient3D): Least Unsquared Deviations, SIAM J. Imaging Sciences, 6, 2450-2483 (2013). """ - def __init__(self, *args, **kwargs): + def __init__( + self, + src, + alpha=None, + tol=1e-3, + mu=1, + gam=1.618, + EPS=1e-12, + maxit=1000, + adp_proj=True, + max_rankZ=None, + max_rankW=None, + adp_mu=True, + dec_mu=0.5, + inc_mu=2, + mu_min=1e-4, + mu_max=1e4, + min_mu_itr=5, + max_mu_itr=20, + delta_mu_l=0.1, + delta_mu_u=10, + **kwargs, + ): """ Initialize a class for estimating 3D orientations using a Least Unsquared Deviations algorithm. @@ -66,7 +88,7 @@ def __init__(self, *args, **kwargs): """ # Handle parameters specific to CommonlineLUD - self.alpha = kwargs.pop("alpha", None) # Spectral norm constraint bound + self.alpha = alpha # Spectral norm constraint bound if self.alpha is not None: if not (2 / 3 <= self.alpha < 1): raise ValueError( @@ -83,28 +105,28 @@ def __init__(self, *args, **kwargs): ) self.spectral_norm_constraint = False - self.tol = kwargs.pop("tol", 1e-3) - self.mu = kwargs.pop("mu", 1) - self.gam = kwargs.pop("gam", 1.618) - self.EPS = kwargs.pop("EPS", 1e-12) - self.maxit = kwargs.pop("maxit", 1000) - self.adp_proj = kwargs.pop("adp_proj", True) - self.max_rankZ = kwargs.pop("max_rankZ", None) - self.max_rankW = kwargs.pop("max_rankW", None) + self.tol = tol + self.mu = mu + self.gam = gam + self.EPS = EPS + self.maxit = maxit + self.adp_proj = adp_proj + self.max_rankZ = max_rankZ + self.max_rankW = max_rankW # Parameters for adjusting mu - self.adp_mu = kwargs.pop("adp_mu", True) - self.dec_mu = kwargs.pop("dec_mu", 0.5) - self.inc_mu = kwargs.pop("inc_mu", 2) - self.mu_min = kwargs.pop("mu_min", 1e-4) - self.mu_max = kwargs.pop("mu_max", 1e4) - self.min_mu_itr = kwargs.pop("min_mu_itr", 5) - self.max_mu_itr = kwargs.pop("max_mu_itr", 20) - self.delta_mu_l = kwargs.pop("delta_mu_l", 0.1) - self.delta_mu_u = kwargs.pop("delta_mu_u", 10) + self.adp_mu = adp_mu + self.dec_mu = dec_mu + self.inc_mu = inc_mu + self.mu_min = mu_min + self.mu_max = mu_max + self.min_mu_itr = min_mu_itr + self.max_mu_itr = max_mu_itr + self.delta_mu_l = delta_mu_l + self.delta_mu_u = delta_mu_u # Initialize commonline base class - super().__init__(*args, **kwargs) + super().__init__(src, **kwargs) def estimate_rotations(self): """ From 1d8098e29552289412188f33678b6ad2b50bd6d5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 22 Jan 2025 15:05:14 -0500 Subject: [PATCH 021/216] Add 50S ribosome to downloader --- src/aspire/downloader/__init__.py | 1 + src/aspire/downloader/data_fetcher.py | 15 +++++++++++++++ src/aspire/downloader/registry.py | 3 +++ src/aspire/utils/__init__.py | 1 + src/aspire/utils/coor_trans.py | 19 +++++++++++++++++++ 5 files changed, 39 insertions(+) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index be0d375878..f6bffe85a3 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -17,6 +17,7 @@ emdb_8511, emdb_10835, emdb_14621, + emdb_51751, remove_downloads, simulated_channelspin, ) diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index f655d06ab3..bd18ec881c 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -90,6 +90,21 @@ def emdb_2660(): return vol +def emdb_51751(): + """ + Downloads the EMDB-51751 volume map and returns a `Volume` instance. + + Cryo-EM map of shaped pulse revitrified 50S ribosomal subunit. Resolution + of 2.9 Angstrom. + + :return: A `Volume` instance. + """ + file_path = fetch_data("emdb_51751.map") + vol = Volume.load(file_path, symmetry_group="C1") + + return vol + + def emdb_8012(): """ Downloads the EMDB-8012 volume map and returns the file path. diff --git a/src/aspire/downloader/registry.py b/src/aspire/downloader/registry.py index 467ad3c772..5188995d15 100644 --- a/src/aspire/downloader/registry.py +++ b/src/aspire/downloader/registry.py @@ -1,6 +1,7 @@ # dataset registry for ASPIRE example data. registry = { "emdb_2660.map": "49aecfd4efce09afc937d1786bbed6f18c2a353c73a4e16a643a304342d0660e", + "emdb_51751.map": "811cd15ab6414903bbe462a3bbbe8bd9060e07405b451cc9f239c0c060390888", "emdb_8012.map": "85a1c9ab958b1dd051d011515212d58decaf2537351b9c016acd3e5852e30d63", "emdb_2984.map": "0194f44cb28f2a8daa5d477d25852de9cc81ed093487de181ee4b30b0d77ef90", "emdb_8511.map": "1f03ec4a0cadb407b6b972c803ffe1e97ff5087d4c2ce9fec2c404747a7fb3fe", @@ -18,6 +19,7 @@ registry_urls = { "emdb_2660.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-2660/map/emd_2660.map.gz", + "emdb_51751.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-51751/map/emd_51751.map.gz", "emdb_8012.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-8012/map/emd_8012.map.gz", "emdb_2984.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-2984/map/emd_2984.map.gz", "emdb_8511.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-8511/map/emd_8511.map.gz", @@ -35,6 +37,7 @@ file_to_method_map = { "emdb_2660.map": "emdb_2660", + "emdb_51751.map": "emdb_51751", "emdb_8012.map": "emdb_8012", "emdb_2984.map": "emdb_2984", "emdb_8511.map": "emdb_8511", diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index 74e13c8b82..1f5e52cdad 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -1,5 +1,6 @@ from .types import complex_type, real_type, utest_tolerance # isort:skip from .coor_trans import ( # isort:skip + aligned_mse, common_line_from_rots, mean_aligned_angular_distance, crop_pad_2d, diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index ce5e5e8a49..1d48cdf861 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -317,6 +317,25 @@ def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): return mean_ang_dist +def aligned_mse(rots_est, rots_gt): + """ + Register estimates to ground truth rotations and compute the + mean angular distance between them (in degrees). + + :param rots_est: A set of estimated rotations of size nx3x3. + :param rots_gt: A set of ground truth rotations of size nx3x3. + + :return: The mean squared error between ground truth rotations + and globally aligned estimated rotations. + """ + Q_mat, flag = register_rotations(rots_est, rots_gt) + logger.debug(f"Registration Q_mat: {Q_mat}\nflag: {flag}") + regrot = get_aligned_rotations(rots_est, Q_mat, flag) + mse = get_rots_mse(regrot, rots_gt) + + return mse + + def common_line_from_rots(r1, r2, ell): """ Compute the common line induced by rotation matrices r1 and r2. From 1e23c38fad50a674f4a8f7800a8f83c15f5140cc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 10 Feb 2025 12:38:16 -0500 Subject: [PATCH 022/216] add logger message for ADMM iterations. Change ADMM tol for tests (cut test time in half). --- src/aspire/abinitio/commonline_lud.py | 3 +++ tests/test_orient_lud.py | 1 + 2 files changed, 4 insertions(+) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 7ae8128cda..339f36cc2c 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -275,6 +275,9 @@ def _compute_Gram(self, C): dinf_term -= Z dinf = np.linalg.norm(dinf_term, "fro") / max(np.linalg.norm(S, np.inf), 1) + logger.info( + f"Iteration: {itr}, residual: {max(pinf, dinf)}, target: {self.tol}" + ) if max(pinf, dinf) <= self.tol: return G diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 43ef4a67f6..9b1605c00e 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -79,6 +79,7 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): max_shift=max_shift, shift_step=shift_step, mask=False, + tol=0.005, ) return src, orient_est From cabd5131d62b204e15a7be397eee97b1c7c89810 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 25 Feb 2025 09:44:13 -0500 Subject: [PATCH 023/216] vectorize _Q_theta --- src/aspire/abinitio/commonline_lud.py | 64 +++++++++++++-------------- 1 file changed, 30 insertions(+), 34 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 339f36cc2c..1d2e26c83c 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -128,6 +128,11 @@ def __init__( # Initialize commonline base class super().__init__(src, **kwargs) + # Upper-triangular mask + ut_mask = np.zeros((self.n_img, self.n_img), dtype=bool) + ut_mask[np.triu_indices(self.n_img, k=1)] = True + self.ut_mask = ut_mask + def estimate_rotations(self): """ Estimate rotation matrices using the common lines method with LUD optimization. @@ -135,14 +140,14 @@ def estimate_rotations(self): logger.info("Computing the common lines matrix.") self.build_clmatrix() - C = self._cl_to_C(self.clmatrix) - gram = self._compute_Gram(C) + self._cl_to_C(self.clmatrix) + gram = self._compute_Gram() gram = self._restructure_Gram(gram) self.rotations = self._deterministic_rounding(gram) return self.rotations - def _compute_Gram(self, C): + def _compute_Gram(self): """ Perform the alternating direction method of multipliers (ADMM) for the SDP problem: @@ -181,7 +186,7 @@ def _compute_Gram(self, C): Phi = W + G / self.mu # Compute initial values - S, theta = self._Q_theta(Phi, C, self.mu) + S, theta = self._Q_theta(Phi) AS = self._compute_AX(S) resi = self._compute_AX(G) - b @@ -206,7 +211,7 @@ def _compute_Gram(self, C): Phi = W + ATy + G / self.mu if self.spectral_norm_constraint: Phi -= Z - S, theta = self._Q_theta(Phi, C, self.mu) + S, theta = self._Q_theta(Phi) ############# # Compute Z # @@ -357,7 +362,7 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): return Z, kk, zz - def _Q_theta(self, phi, C, mu): + def _Q_theta(self, phi): """ Compute the matrix S and auxiliary variables theta for the optimization problem. @@ -369,11 +374,6 @@ def _Q_theta(self, phi, C, mu): Python equivalent of matlab Qtheta MEX function. :param phi: ndarray, A 2*n_img x 2*n_img scaled dual variable matrix (Phi) used in the ADMM iterations. - :param C: ndarray, A 3D array (n_img x n_img x 2) containing commonline coordinates (in Cartesian form) - between pairs of images. Each C[i, j] stores the x and y coordinates of the common - line between image i and image j. - :param mu: float, The penalty parameter in the augmented Lagrangian. It controls the scaling of the - dual variable contribution in the ADMM updates. :returns: - S, A 2*n_img x 2*n_img matrix representing the fidelity term. It is a symmetric matrix derived from the commonline constraints, normalized by theta. @@ -381,29 +381,24 @@ def _Q_theta(self, phi, C, mu): used to compute S. Each theta[i, j] stores the normalized adjustments for the common line between image i and image j. """ - # Initialize outputs - S = np.zeros((2 * self.n_img, 2 * self.n_img), dtype=self.dtype) - theta = np.zeros_like(C) + # Initialize theta, shape = (n_img, n_img, 2). + theta = np.zeros_like(self.C) - # Main routine - for i in range(self.n_img - 1): - for j in range(i + 1, self.n_img): - t = 0 - for k in range(2): - theta[i, j, k] = C[i, j, k] - mu * ( - phi[2 * i + k, 2 * j] * C[j, i, 0] - + phi[2 * i + k, 2 * j + 1] * C[j, i, 1] - ) - t += theta[i, j, k] ** 2 - - t = np.sqrt(t) - for k in range(2): - if self.spectral_norm_constraint: - theta[i, j, k] /= t - else: - theta[i, j, k] /= max(t, self.mu) - S[2 * i + k, 2 * j] = theta[i, j, k] * C[j, i, 0] - S[2 * i + k, 2 * j + 1] = theta[i, j, k] * C[j, i, 1] + # Compute theta + phi = phi.reshape(self.n_img, 2, self.n_img, 2).transpose(0, 2, 1, 3) + sum_prod = (phi[self.ut_mask] * self.C_t[self.ut_mask, None]).sum(axis=2) + theta[self.ut_mask] = self.C[self.ut_mask] - self.mu * sum_prod + + # Normalize theta + theta_norm = np.linalg.norm(theta[self.ut_mask], axis=-1)[..., None] + if self.spectral_norm_constraint: + theta[self.ut_mask] /= theta_norm + else: + theta[self.ut_mask] /= np.maximum(theta_norm, self.mu) + + # Construct S + S = theta[..., None] * self.C_t[:, :, None] + S = S.transpose(0, 2, 1, 3).reshape(2 * self.n_img, 2 * self.n_img) # Ensure S is symmetric S = (S + S.T) / 2 @@ -430,7 +425,8 @@ def _cl_to_C(self, clmatrix): C[j, i, 0] = np.cos(2 * np.pi * cl_ji / self.n_theta) C[j, i, 1] = np.sin(2 * np.pi * cl_ji / self.n_theta) - return C + self.C = C + self.C_t = np.ascontiguousarray(C.transpose(1, 0, 2)) @staticmethod def _compute_AX(X): From ea9ff04012480c4b823c5a6d8f387d3a13b1b175 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 25 Feb 2025 14:45:38 -0500 Subject: [PATCH 024/216] small cleanup --- src/aspire/abinitio/commonline_lud.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 1d2e26c83c..933970a81d 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -128,7 +128,7 @@ def __init__( # Initialize commonline base class super().__init__(src, **kwargs) - # Upper-triangular mask + # Upper-triangular mask used in `_Q_theta` ut_mask = np.zeros((self.n_img, self.n_img), dtype=bool) ut_mask[np.triu_indices(self.n_img, k=1)] = True self.ut_mask = ut_mask @@ -251,10 +251,9 @@ def _compute_Gram(self): else: nev = 6 - n_eigs = nev if self.spectral_norm_constraint: - n_eigs = min(nev, n) - dH, V = eigs(-H.astype(np.float64), k=n_eigs, which="LR") + nev = min(nev, n) + dH, V = eigs(-H.astype(np.float64), k=nev, which="LR") # Sort by eigenvalue magnitude. dH = dH.real.astype(self.dtype, copy=False) From 45244a425c96d7400f97b03c263a2cdee88232ee Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 25 Feb 2025 14:46:25 -0500 Subject: [PATCH 025/216] add gallery experiment --- .../commonline_lud_simulated_data.py | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 gallery/experiments/commonline_lud_simulated_data.py diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py new file mode 100644 index 0000000000..dd9316bd6f --- /dev/null +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -0,0 +1,96 @@ +""" +Commonlines Method using LUD +============================ + +This tutorial demonstrates using the Least Unsquared Deviations (LUD) +commonlines method for estimating particle orientations. This tutorial +reproduces the "Experiments on simulated images" found in the publication: + +Orientation Determination of Cryo-EM Images Using Least Unsquared Deviations, +L. Wang, A. Singer, and Z. Wen, SIAM J. Imaging Sciences, 6, 2450-2483 (2013). +""" + +# %% +# Imports +# ------- + +import logging + +import numpy as np + +from aspire.abinitio import CommonlineLUD +from aspire.noise import WhiteNoiseAdder +from aspire.source import OrientedSource, Simulation +from aspire.utils import aligned_mse, mean_aligned_angular_distance +from aspire.volume import Volume + +logger = logging.getLogger(__name__) + + +# %% +# Generate Simulated Data +# ----------------------- +# We generate simulated noisy images from a low res volume map of +# the 50S ribosomal subunit of E. coli. + +n_imgs = 500 # Number of images in our source +snr = 1 / 8 # Signal-to-noise ratio +dtype = np.float64 +res = 129 + +# We download the volume map and zero-pad to the indicated resolution. +vol = ( + Volume.load("../tutorials/data/clean70SRibosome_vol_65p.mrc") + .astype(dtype) + .downsample(res) +) +logger.info("Volume map data" f" shape: {vol.shape} dtype:{vol.dtype}") + +# We generate a white noise adder with specifid SNR. +noise_adder = WhiteNoiseAdder.from_snr(snr=snr) + +# Now we initialize a Simulation source to generate noisy, centered images. +src = Simulation( + n=n_imgs, + vols=vol, + offsets=0, + noise_adder=noise_adder, + dtype=dtype, +) + +# We can view the noisy images. +src.images[:5].show() + +# %% +# Estimate Orientations +# --------------------- +# We use the LUD commonline algorithm to estimate the orientation of the noisy images. + +logger.info("Begin Orientation Estimation") + +# Create a custom orientation estimation object which uses the LUD algorithm. +# By default, we use the algortihm without spectral norm constraint. +orient_est = CommonlineLUD(src, n_theta=360) + +# Initialize an ``OrientedSource`` class instance that performs orientation +# estimation in a lazy fashion upon request of images or rotations. +oriented_src = OrientedSource(src, orient_est) + +# %% +# Results +# ------- +# We measure our results by finding the mean angular distance between the +# ground truth rotations and the estimated rotations adjusted by the best +# global alignment. +mean_ang_dist = mean_aligned_angular_distance(oriented_src.rotations, src.rotations) +logger.info( + f"Mean angular distance between globally aligned estimates and ground truth rotations: {mean_ang_dist}\n" +) + +# Additionally, we can compute the mean squared error between the ground truth rotations +# and the globally aligned estimated rotations. +mse = aligned_mse(oriented_src.rotations, src.rotations) +logger.info( + f"Mean squared error between globally aligned estimates and ground truth rotations: {mse}\n" +) + From 19243b63cc492af5b599ff3b7eba935db38c3638 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 25 Feb 2025 15:36:08 -0500 Subject: [PATCH 026/216] blank line --- gallery/experiments/commonline_lud_simulated_data.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index dd9316bd6f..412cdccf29 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -93,4 +93,3 @@ logger.info( f"Mean squared error between globally aligned estimates and ground truth rotations: {mse}\n" ) - From a713bf72adc3f1446a1a07c52517dd8f4217833c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Feb 2025 09:25:52 -0500 Subject: [PATCH 027/216] remove aligned_mse. --- .../commonline_lud_simulated_data.py | 9 +-------- src/aspire/utils/coor_trans.py | 19 ------------------- 2 files changed, 1 insertion(+), 27 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index 412cdccf29..804f83d8c5 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -21,7 +21,7 @@ from aspire.abinitio import CommonlineLUD from aspire.noise import WhiteNoiseAdder from aspire.source import OrientedSource, Simulation -from aspire.utils import aligned_mse, mean_aligned_angular_distance +from aspire.utils import mean_aligned_angular_distance from aspire.volume import Volume logger = logging.getLogger(__name__) @@ -86,10 +86,3 @@ logger.info( f"Mean angular distance between globally aligned estimates and ground truth rotations: {mean_ang_dist}\n" ) - -# Additionally, we can compute the mean squared error between the ground truth rotations -# and the globally aligned estimated rotations. -mse = aligned_mse(oriented_src.rotations, src.rotations) -logger.info( - f"Mean squared error between globally aligned estimates and ground truth rotations: {mse}\n" -) diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index 1d48cdf861..ce5e5e8a49 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -317,25 +317,6 @@ def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): return mean_ang_dist -def aligned_mse(rots_est, rots_gt): - """ - Register estimates to ground truth rotations and compute the - mean angular distance between them (in degrees). - - :param rots_est: A set of estimated rotations of size nx3x3. - :param rots_gt: A set of ground truth rotations of size nx3x3. - - :return: The mean squared error between ground truth rotations - and globally aligned estimated rotations. - """ - Q_mat, flag = register_rotations(rots_est, rots_gt) - logger.debug(f"Registration Q_mat: {Q_mat}\nflag: {flag}") - regrot = get_aligned_rotations(rots_est, Q_mat, flag) - mse = get_rots_mse(regrot, rots_gt) - - return mse - - def common_line_from_rots(r1, r2, ell): """ Compute the common line induced by rotation matrices r1 and r2. From 2ed1f0a455ca0428f453ee19f01fb5dc08744f61 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Feb 2025 09:47:06 -0500 Subject: [PATCH 028/216] remove import --- src/aspire/utils/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index 1f5e52cdad..74e13c8b82 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -1,6 +1,5 @@ from .types import complex_type, real_type, utest_tolerance # isort:skip from .coor_trans import ( # isort:skip - aligned_mse, common_line_from_rots, mean_aligned_angular_distance, crop_pad_2d, From 7f9451f519952476de504aaed806fa2e86b45e25 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Feb 2025 10:48:12 -0500 Subject: [PATCH 029/216] add docstring params --- src/aspire/abinitio/commonline_lud.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 933970a81d..1f66130792 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -305,6 +305,20 @@ def _compute_Gram(self): def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): """ Update ADMM subproblem for enforcing the spectral norm constraint. + + :param S: A 2*n_img x 2*n_img symmetric matrix representing the fidelity term. + :param W: A 2*n_img x 2*n_img array, primary ADMM subproblem matrix. + :param ATy: A 2*n_img x 2*n_img array. + :param G: Current value of the 2*n_img x 2*n_img optimization solution matrix. + :param zz: eigenvalues from previous iteration. + :param itr: ADMM loop iteration. + :param kk: Number of eigenvalues of Z to use to enforce spectral norm constraint. + :param nev: Number of eigenvalues of W used in previous iteration of ADMM. + + :returns: + - Z, Updated 2*n_img x 2*n_img matrix for spectral norm constraint ADMM subproblem. + - kk, Number of eigenvalues of Z to use to enforce spectral norm constraint in next iteration. + - nev, Number of eigenvalues of W to use in this iteration of ADMM. """ lambda_ = self.alpha * self.n_img # Spectral norm bound B = S + W + ATy + G / self.mu From 50c19c2af63f3f7b8fd3efe15bbef9a22178d603 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 6 Mar 2025 08:49:16 -0500 Subject: [PATCH 030/216] src caching for test. remove gpu skip --- tests/test_orient_lud.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 9b1605c00e..8e4ec15434 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -63,12 +63,16 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): dtype=dtype, ) + # Cache source to prevent regenerating images. + src = src.cache() + # Increase max_shift and set shift_step to be sub-pixel when using # random offsets in the Simulation. This improves common-line detection. max_shift = 0.20 shift_step = 0.25 # Set max_shift 1 pixel and shift_step to 1 pixel when using 0 offsets. + # This reduces the search space for commonline detection and improves test speed. if np.all(src.offsets == 0.0): max_shift = 1 / src.L shift_step = 1 @@ -79,7 +83,7 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): max_shift=max_shift, shift_step=shift_step, mask=False, - tol=0.005, + tol=0.005, # Improves test speed ) return src, orient_est @@ -88,9 +92,7 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): def test_estimate_rotations(src_orient_est_fixture): src, orient_est = src_orient_est_fixture - if backend_available("cufinufft") and src.dtype == np.float32: - pytest.skip("CI on GPU fails for singles.") - + # Estimate rotations est_rots = orient_est.estimate_rotations() # Register estimates to ground truth rotations and compute the From 78c47601151389d23c38a18c0909ce31fe69be79 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 6 Mar 2025 08:53:57 -0500 Subject: [PATCH 031/216] unused import --- tests/test_orient_lud.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 8e4ec15434..02e6818ad6 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -2,7 +2,6 @@ import pytest from aspire.abinitio import CommonlineLUD -from aspire.nufft import backend_available from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance from aspire.volume import AsymmetricVolume From 4e16d87ae0bf89b5c997ed3558f6b80b9a6dda65 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 6 Mar 2025 11:06:37 -0500 Subject: [PATCH 032/216] use default max_shift/shift_step. run doubles always. --- tests/test_orient_lud.py | 25 ++++++++++--------------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 02e6818ad6..ab3bbf5ef9 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -18,7 +18,7 @@ DTYPES = [ np.float32, - pytest.param(np.float64, marks=pytest.mark.expensive), + np.float64, ] SPECTRAL_NORM_CONSTRAINT = [ @@ -65,22 +65,10 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): # Cache source to prevent regenerating images. src = src.cache() - # Increase max_shift and set shift_step to be sub-pixel when using - # random offsets in the Simulation. This improves common-line detection. - max_shift = 0.20 - shift_step = 0.25 - - # Set max_shift 1 pixel and shift_step to 1 pixel when using 0 offsets. - # This reduces the search space for commonline detection and improves test speed. - if np.all(src.offsets == 0.0): - max_shift = 1 / src.L - shift_step = 1 - + # Generate LUD orientation estimation object. orient_est = CommonlineLUD( src, alpha=alpha, - max_shift=max_shift, - shift_step=shift_step, mask=False, tol=0.005, # Improves test speed ) @@ -97,7 +85,14 @@ def test_estimate_rotations(src_orient_est_fixture): # Register estimates to ground truth rotations and compute the # angular distance between them (in degrees). # Assert that mean aligned angular distance is less than 3 degrees. - mean_aligned_angular_distance(est_rots, src.rotations, degree_tol=3) + tol = 3 + + # Using LUD without spectral norm constraint, ie. alpha=None, + # on shifted images reduces estimated rotations accuracy. + # This can be improved by using subpixel shift_step in CommonlineLUD. + if orient_est.alpha is None and src.offsets.all() != 0: + tol = 9 + mean_aligned_angular_distance(est_rots, src.rotations, degree_tol=tol) # Check dtype pass-through np.testing.assert_equal(src.dtype, est_rots.dtype) From def285391d57b22bc2a5f1c715d10956309df3d6 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 6 Mar 2025 11:23:11 -0500 Subject: [PATCH 033/216] separate fixture for orient_est --- tests/test_orient_lud.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index ab3bbf5ef9..9fc5cabff5 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -48,8 +48,8 @@ def alpha(request): @pytest.fixture -def src_orient_est_fixture(resolution, offsets, dtype, alpha): - """Fixture for simulation source and orientation estimation object.""" +def source(resolution, offsets, dtype): + """Fixture for simulation source object.""" src = Simulation( n=60, L=resolution, @@ -65,20 +65,24 @@ def src_orient_est_fixture(resolution, offsets, dtype, alpha): # Cache source to prevent regenerating images. src = src.cache() + return src + + +@pytest.fixture +def orient_est(source, alpha): + """Fixture for LUD orientation estimation object.""" # Generate LUD orientation estimation object. orient_est = CommonlineLUD( - src, + source, alpha=alpha, mask=False, tol=0.005, # Improves test speed ) - return src, orient_est - + return orient_est -def test_estimate_rotations(src_orient_est_fixture): - src, orient_est = src_orient_est_fixture +def test_estimate_rotations(source, orient_est): # Estimate rotations est_rots = orient_est.estimate_rotations() @@ -90,9 +94,9 @@ def test_estimate_rotations(src_orient_est_fixture): # Using LUD without spectral norm constraint, ie. alpha=None, # on shifted images reduces estimated rotations accuracy. # This can be improved by using subpixel shift_step in CommonlineLUD. - if orient_est.alpha is None and src.offsets.all() != 0: + if orient_est.alpha is None and source.offsets.all() != 0: tol = 9 - mean_aligned_angular_distance(est_rots, src.rotations, degree_tol=tol) + mean_aligned_angular_distance(est_rots, source.rotations, degree_tol=tol) # Check dtype pass-through - np.testing.assert_equal(src.dtype, est_rots.dtype) + np.testing.assert_equal(source.dtype, est_rots.dtype) From 3f4a31be73d640c2c7dce643970f50bf3a6fa6ff Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 6 Mar 2025 15:55:42 -0500 Subject: [PATCH 034/216] add loop and results table to gallery --- .../commonline_lud_simulated_data.py | 110 +++++++++++------- 1 file changed, 66 insertions(+), 44 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index 804f83d8c5..4fc471b7ec 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -20,7 +20,7 @@ from aspire.abinitio import CommonlineLUD from aspire.noise import WhiteNoiseAdder -from aspire.source import OrientedSource, Simulation +from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance from aspire.volume import Volume @@ -28,61 +28,83 @@ # %% -# Generate Simulated Data -# ----------------------- -# We generate simulated noisy images from a low res volume map of -# the 50S ribosomal subunit of E. coli. +# Parameters +# ---------- +# Set up some initializing parameters. We will run the LUD algorithm +# for various levels of noise and output a table of results. -n_imgs = 500 # Number of images in our source -snr = 1 / 8 # Signal-to-noise ratio +SNR = [1 / 8, 1 / 16, 1 / 32] # Signal-to-noise ratio +n_imgs = 50 # Number of images in our source dtype = np.float64 -res = 129 +pad_size = 129 +results = { + "SNR": ["1/8", "1/16", "1/32"], + "Mean Angular Distance": [], +} # Dictionary to store results -# We download the volume map and zero-pad to the indicated resolution. +# %% +# Load Volume Map +# --------------- +# We will generate simulated noisy images from a low res volume +# map available in our data folder. This volume map is a 65 x 65 x 65 +# voxel volume which we intend to upsample to 129 x 129 x 129. +# To do this we use our ``downsample`` method which, when provided a voxel +# size larger than the input volume, internally zero-pads in Fourier +# space to increase the overall shape of the volume. vol = ( Volume.load("../tutorials/data/clean70SRibosome_vol_65p.mrc") .astype(dtype) - .downsample(res) + .downsample(pad_size) ) logger.info("Volume map data" f" shape: {vol.shape} dtype:{vol.dtype}") -# We generate a white noise adder with specifid SNR. -noise_adder = WhiteNoiseAdder.from_snr(snr=snr) - -# Now we initialize a Simulation source to generate noisy, centered images. -src = Simulation( - n=n_imgs, - vols=vol, - offsets=0, - noise_adder=noise_adder, - dtype=dtype, -) - -# We can view the noisy images. -src.images[:5].show() +# %% +# Generate Noisy Images and Estimate Rotations +# -------------------------------------------- +# A ``Simulation`` object is used to generate simulated data at various +# noise levels. Then rotations are estimated using ``CommonlineLUD` algorithm. +# Results are measured by computing the mean aligned angular distance between +# the ground truth rotations and the globally aligned estimated rotations. +for snr in SNR: + # Generate a white noise adder with specifid SNR. + noise_adder = WhiteNoiseAdder.from_snr(snr=snr) + + # Initialize a Simulation source to generate noisy, centered images. + src = Simulation( + n=n_imgs, + vols=vol, + offsets=0, + amplitudes=1, + noise_adder=noise_adder, + dtype=dtype, + ).cache() + + # Estimate rotations using the LUD algorithm. + orient_est = CommonlineLUD(src) + est_rotations = orient_est.estimate_rotations() + + # Find the mean aligned angular distance between estimates and ground truth rotations. + mean_ang_dist = mean_aligned_angular_distance(est_rotations, src.rotations) + + # Store results. + results["Mean Angular Distance"].append(mean_ang_dist) # %% -# Estimate Orientations -# --------------------- -# We use the LUD commonline algorithm to estimate the orientation of the noisy images. +# Display Results +# --------------- +# Display table of results for various noise levels. -logger.info("Begin Orientation Estimation") +# Column widths +col1_width = 10 +col2_width = 22 -# Create a custom orientation estimation object which uses the LUD algorithm. -# By default, we use the algortihm without spectral norm constraint. -orient_est = CommonlineLUD(src, n_theta=360) +# Create table as a string +table = [] +table.append(f"{'SNR':<{col1_width}} {'Mean Angular Distance':<{col2_width}}") +table.append("-" * (col1_width + col2_width)) -# Initialize an ``OrientedSource`` class instance that performs orientation -# estimation in a lazy fashion upon request of images or rotations. -oriented_src = OrientedSource(src, orient_est) +for snr, angle in zip(results["SNR"], results["Mean Angular Distance"]): + table.append(f"{snr:<{col1_width}} {angle:<{col2_width}}") -# %% -# Results -# ------- -# We measure our results by finding the mean angular distance between the -# ground truth rotations and the estimated rotations adjusted by the best -# global alignment. -mean_ang_dist = mean_aligned_angular_distance(oriented_src.rotations, src.rotations) -logger.info( - f"Mean angular distance between globally aligned estimates and ground truth rotations: {mean_ang_dist}\n" -) +# Log the table +logger.info("\n" + "\n".join(table)) From 298aab592bfa84cf9d4a0f8936d252aad0dbeb73 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 7 Mar 2025 14:30:49 -0500 Subject: [PATCH 035/216] remove emdb_51751 --- src/aspire/downloader/__init__.py | 1 - src/aspire/downloader/data_fetcher.py | 15 --------------- src/aspire/downloader/registry.py | 3 --- 3 files changed, 19 deletions(-) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index f6bffe85a3..be0d375878 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -17,7 +17,6 @@ emdb_8511, emdb_10835, emdb_14621, - emdb_51751, remove_downloads, simulated_channelspin, ) diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index bd18ec881c..f655d06ab3 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -90,21 +90,6 @@ def emdb_2660(): return vol -def emdb_51751(): - """ - Downloads the EMDB-51751 volume map and returns a `Volume` instance. - - Cryo-EM map of shaped pulse revitrified 50S ribosomal subunit. Resolution - of 2.9 Angstrom. - - :return: A `Volume` instance. - """ - file_path = fetch_data("emdb_51751.map") - vol = Volume.load(file_path, symmetry_group="C1") - - return vol - - def emdb_8012(): """ Downloads the EMDB-8012 volume map and returns the file path. diff --git a/src/aspire/downloader/registry.py b/src/aspire/downloader/registry.py index 5188995d15..467ad3c772 100644 --- a/src/aspire/downloader/registry.py +++ b/src/aspire/downloader/registry.py @@ -1,7 +1,6 @@ # dataset registry for ASPIRE example data. registry = { "emdb_2660.map": "49aecfd4efce09afc937d1786bbed6f18c2a353c73a4e16a643a304342d0660e", - "emdb_51751.map": "811cd15ab6414903bbe462a3bbbe8bd9060e07405b451cc9f239c0c060390888", "emdb_8012.map": "85a1c9ab958b1dd051d011515212d58decaf2537351b9c016acd3e5852e30d63", "emdb_2984.map": "0194f44cb28f2a8daa5d477d25852de9cc81ed093487de181ee4b30b0d77ef90", "emdb_8511.map": "1f03ec4a0cadb407b6b972c803ffe1e97ff5087d4c2ce9fec2c404747a7fb3fe", @@ -19,7 +18,6 @@ registry_urls = { "emdb_2660.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-2660/map/emd_2660.map.gz", - "emdb_51751.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-51751/map/emd_51751.map.gz", "emdb_8012.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-8012/map/emd_8012.map.gz", "emdb_2984.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-2984/map/emd_2984.map.gz", "emdb_8511.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-8511/map/emd_8511.map.gz", @@ -37,7 +35,6 @@ file_to_method_map = { "emdb_2660.map": "emdb_2660", - "emdb_51751.map": "emdb_51751", "emdb_8012.map": "emdb_8012", "emdb_2984.map": "emdb_2984", "emdb_8511.map": "emdb_8511", From 982f0aa312e7fd497c32db189c45ff9b5d60bc69 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 7 Mar 2025 14:40:52 -0500 Subject: [PATCH 036/216] doubles expensive --- tests/test_orient_lud.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 9fc5cabff5..46b2c44afb 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -18,7 +18,7 @@ DTYPES = [ np.float32, - np.float64, + pytest.param(np.float64, marks=pytest.mark.expensive), ] SPECTRAL_NORM_CONSTRAINT = [ From 41bad44ba46ae0470fc1e3a284885a14cedd3732 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 7 Mar 2025 14:50:00 -0500 Subject: [PATCH 037/216] remove unused method lud_prep --- src/aspire/abinitio/commonline_lud.py | 38 --------------------------- 1 file changed, 38 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 1f66130792..d5e5761f55 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -525,44 +525,6 @@ def _compute_ATy(y): ATy = csr_array((data, (rows, cols)), shape=(n, n)) return ATy - def _lud_prep(self): - """ - Prepare optimization problem constraints. - - The constraints for the LUD optimization, max tr(SG), performed in `_compute_gram_matrix()` - as min tr(-SG), are that the Gram matrix, G, is semidefinite positive and G11_ii = G22_ii = 1, - G12_ii = G21_ii = 0, i=1,2,...,N, for the block representation of G = [[G11, G12], [G21, G22]]. - - We build a corresponding constraint in the form of tr(A_j @ G) = b_j, j = 1,...,p. - For the constraint G11_ii = G22_ii = 1, we have A_j[i, i] = 1 (zeros elsewhere) and b_j = 1. - For the constraint G12_ii = G21_ii = 0, we have A_j[i, i] = 1 (zeros elsewhere) and b_j = 0. - - :returns: Constraint data A, b. - """ - logger.info("Preparing LUD optimization constraints.") - - n = 2 * self.n_img - A = [] - b = [] - data = np.ones(1, dtype=self.dtype) - for i in range(n): - row_ind = np.array([i]) - col_ind = np.array([i]) - A_i = csr_array((data, (row_ind, col_ind)), shape=(n, n), dtype=self.dtype) - A.append(A_i) - b.append(1) - - for i in range(self.n_img): - row_ind = np.array([i]) - col_ind = np.array([self.n_img + i]) - A_i = csr_array((data, (row_ind, col_ind)), shape=(n, n), dtype=self.dtype) - A.append(A_i) - b.append(0) - - b = np.array(b, dtype=self.dtype) - - return A, b - def _restructure_Gram(self, G): """ Restructures the input Gram matrix into a block structure based on the following From 7f3622d1ee714d2bb0d0677c5e55962b01017bb5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 7 Mar 2025 15:32:01 -0500 Subject: [PATCH 038/216] remove spectral_norm_constraint attribute. Just use alpha --- src/aspire/abinitio/commonline_lud.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index d5e5761f55..b040c10f8e 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -98,12 +98,10 @@ def __init__( logger.info( f"Initializing LUD algorithm using ADMM with spectral norm constraint {self.alpha}." ) - self.spectral_norm_constraint = True else: logger.info( "Initializing LUD algorithm using ADMM without spectral norm constraint." ) - self.spectral_norm_constraint = False self.tol = tol self.mu = mu @@ -173,13 +171,13 @@ def _compute_Gram(self): # Adjust rank limits self.max_rankW = self.max_rankW or max(6, self.n_img // 2) - if self.spectral_norm_constraint: + if self.alpha is not None: self.max_rankZ = self.max_rankZ or max(6, self.n_img // 2) # Initialize variables G = np.eye(n, dtype=self.dtype) W = np.eye(n, dtype=self.dtype) - if self.spectral_norm_constraint: + if self.alpha is not None: Z = W Phi = G / self.mu else: @@ -201,7 +199,7 @@ def _compute_Gram(self): # Compute y # ############# y = -(AS + self._compute_AX(W)) - resi / self.mu - if self.spectral_norm_constraint: + if self.alpha is not None: y += self._compute_AX(Z) ################# @@ -209,21 +207,21 @@ def _compute_Gram(self): ################# ATy = self._compute_ATy(y) Phi = W + ATy + G / self.mu - if self.spectral_norm_constraint: + if self.alpha is not None: Phi -= Z S, theta = self._Q_theta(Phi) ############# # Compute Z # ############# - if self.spectral_norm_constraint: + if self.alpha is not None: Z, kk, zz = self._compute_Z(S, W, ATy, G, zz, itr, kk, nev) ############# # Compute W # ############# H = -S - ATy - G / self.mu - if self.spectral_norm_constraint: + if self.alpha is not None: H += Z H = (H + H.T) / 2 @@ -251,7 +249,7 @@ def _compute_Gram(self): else: nev = 6 - if self.spectral_norm_constraint: + if self.alpha is not None: nev = min(nev, n) dH, V = eigs(-H.astype(np.float64), k=nev, which="LR") @@ -275,7 +273,7 @@ def _compute_Gram(self): pinf = np.linalg.norm(resi) / max(np.linalg.norm(b), 1) dinf_term = S + W + ATy - if self.spectral_norm_constraint: + if self.alpha is not None: dinf_term -= Z dinf = np.linalg.norm(dinf_term, "fro") / max(np.linalg.norm(S, np.inf), 1) @@ -404,7 +402,7 @@ def _Q_theta(self, phi): # Normalize theta theta_norm = np.linalg.norm(theta[self.ut_mask], axis=-1)[..., None] - if self.spectral_norm_constraint: + if self.alpha is not None: theta[self.ut_mask] /= theta_norm else: theta[self.ut_mask] /= np.maximum(theta_norm, self.mu) From 52d30a0900a4ae291f6bd3e44bccba9083c6ac38 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 7 Mar 2025 15:52:51 -0500 Subject: [PATCH 039/216] cleanup logger message logic --- src/aspire/abinitio/commonline_lud.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index b040c10f8e..8fec58f669 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -89,19 +89,17 @@ def __init__( # Handle parameters specific to CommonlineLUD self.alpha = alpha # Spectral norm constraint bound - if self.alpha is not None: - if not (2 / 3 <= self.alpha < 1): - raise ValueError( - "Spectral norm constraint, alpha, must be in [2/3, 1)." - ) - else: - logger.info( - f"Initializing LUD algorithm using ADMM with spectral norm constraint {self.alpha}." - ) - else: - logger.info( - "Initializing LUD algorithm using ADMM without spectral norm constraint." + if self.alpha is not None and not (2 / 3 <= self.alpha < 1): + raise ValueError("Spectral norm constraint, alpha, must be in [2/3, 1).") + + logger.info( + f"Initializing LUD algorithm using ADMM" + + ( + f" with spectral norm constraint: {self.alpha}." + if self.alpha is not None + else " without spectral norm constraint." ) + ) self.tol = tol self.mu = mu From d47365b5c051e6f7776b7e97e02dbb5be7a0bf04 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 7 Mar 2025 15:56:53 -0500 Subject: [PATCH 040/216] tox --- src/aspire/abinitio/commonline_lud.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 8fec58f669..f4f610901d 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -93,7 +93,7 @@ def __init__( raise ValueError("Spectral norm constraint, alpha, must be in [2/3, 1).") logger.info( - f"Initializing LUD algorithm using ADMM" + "Initializing LUD algorithm using ADMM" + ( f" with spectral norm constraint: {self.alpha}." if self.alpha is not None From deb5e366cd0a2f22d40cae37aec54e11837fd010 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 10 Mar 2025 10:10:50 -0400 Subject: [PATCH 041/216] Test branches missed in codecov --- tests/test_orient_lud.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 46b2c44afb..3cd105f651 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -26,6 +26,11 @@ pytest.param(None, marks=pytest.mark.expensive), ] +ADAPTIVE_PROJECTION = [ + True, + False, +] + @pytest.fixture(params=RESOLUTION, ids=lambda x: f"resolution={x}") def resolution(request): @@ -47,6 +52,11 @@ def alpha(request): return request.param +@pytest.fixture(params=ADAPTIVE_PROJECTION, ids=lambda x: f"adp_proj={x}") +def adp_proj(request): + return request.param + + @pytest.fixture def source(resolution, offsets, dtype): """Fixture for simulation source object.""" @@ -69,12 +79,14 @@ def source(resolution, offsets, dtype): @pytest.fixture -def orient_est(source, alpha): +def orient_est(source, alpha, adp_proj): """Fixture for LUD orientation estimation object.""" # Generate LUD orientation estimation object. orient_est = CommonlineLUD( source, alpha=alpha, + adp_proj=adp_proj, + delta_mu_l=0.4, # Ensures branch is tested mask=False, tol=0.005, # Improves test speed ) From b7d1a5b5bc7a292611a1631f34bffb9c24cb00cc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 10 Mar 2025 10:29:50 -0400 Subject: [PATCH 042/216] eigenvalue mask --- src/aspire/abinitio/commonline_lud.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index f4f610901d..27fd8ddf13 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -225,7 +225,9 @@ def _compute_Gram(self): if not self.adp_proj: D, V = np.linalg.eigh(H) - W = V[:, D > self.EPS] @ np.diag(D[D > self.EPS]) @ V[:, D > self.EPS].T + eigs_mask = D > self.EPS + V = V[:, eigs_mask] + W = V @ np.diag(D[eigs_mask]) @ V.T else: if itr == 0: nev = self.max_rankW From 579dd8f2e27f0cac9834ad4a5fc52c0dc1e807ef Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 11 Mar 2025 10:23:52 -0400 Subject: [PATCH 043/216] use eigsh (instead of eigs) on symmetric matrices --- src/aspire/abinitio/commonline_lud.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 27fd8ddf13..541e71cd29 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -2,7 +2,7 @@ import numpy as np from scipy.sparse import csr_array -from scipy.sparse.linalg import eigs +from scipy.sparse.linalg import eigsh from aspire.abinitio import CLOrient3D @@ -251,13 +251,12 @@ def _compute_Gram(self): if self.alpha is not None: nev = min(nev, n) - dH, V = eigs(-H.astype(np.float64), k=nev, which="LR") - # Sort by eigenvalue magnitude. - dH = dH.real.astype(self.dtype, copy=False) - idx = np.argsort(dH)[::-1] - dH = dH[idx] - V = V[:, idx].real.astype(self.dtype, copy=False) + # Compute Eigenvectors and sort by largest algebraic eigenvalue + dH, V = eigsh(-H.astype(np.float64), k=nev, which="LA") + dH = dH[::-1].astype(self.dtype, copy=False) + V = V[:, ::-1].astype(self.dtype, copy=False) + nD = dH > self.EPS dH = dH[nD] nev = np.count_nonzero(nD) @@ -351,15 +350,15 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): kk = 6 kk = min(kk, 2 * self.n_img) - pi, U = eigs( + pi, U = eigsh( B.astype(np.float64, copy=False), k=kk, which="LM" ) # Compute top `kk` eigenvalues and eigenvectors - # Sort by eigenvalue magnitude. + # Sort by eigenvalue magnitude. Note, eigsh does not return + # ordered eigenvalues/vectors. idx = np.argsort(np.abs(pi))[::-1] pi = pi[idx] - U = U[:, idx].real - pi = pi.real # Ensure real eigenvalues for subsequent calculations + U = U[:, idx] # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. zz = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) From 08fe0dd4b8aff7e77aa185c43b34290780485d23 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 12 Mar 2025 15:25:04 -0400 Subject: [PATCH 044/216] compute_num_eigs method --- src/aspire/abinitio/commonline_lud.py | 86 +++++++++++++++------------ 1 file changed, 48 insertions(+), 38 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 541e71cd29..9ca71fff18 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -229,26 +229,13 @@ def _compute_Gram(self): V = V[:, eigs_mask] W = V @ np.diag(D[eigs_mask]) @ V.T else: + # Determine number of eigenvalues to compute for adaptive projection if itr == 0: nev = self.max_rankW else: - if nev > 0: - drops = dH[:-1] / dH[1:] - - # Find max drop - imx = np.argmax(drops) - dmx = drops[imx] - - # Relative drop - rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) - - if rel_drp > 50: - nev = max(imx + 1, 6) - else: - nev = nev + 5 - else: - nev = 6 + nev = self._compute_num_eigs(nev, dH, nev, 50, 5) + # If using a spectral norm constraint cap num_eigs at 2*n_img if self.alpha is not None: nev = min(nev, n) @@ -324,30 +311,11 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): if not self.adp_proj: U, pi = np.linalg.eigh(B) else: + # Determine number of eigenvalues to compute for adaptive projection if itr == 0: kk = self.max_rankZ else: - if kk > 0: - # Initialize relative drop - rel_drp = 0 - imx = 0 - # Calculate relative drop based on `zz` - if len(zz) == 2: - rel_drp = np.inf - elif len(zz) > 2: - drops = zz[:-1] / zz[1:] - - # Find max drop - imx = np.argmax(drops) - dmx = drops[imx] - - # Relative drop - rel_drp = (nev - 1) * dmx / (np.sum(drops) - dmx) - - # Update `kk` based on relative drop - kk = max(imx, 6) if rel_drp > 10 else kk + 3 - else: - kk = 6 + kk = self._compute_num_eigs(kk, zz, nev, 10, 3) kk = min(kk, 2 * self.n_img) pi, U = eigsh( @@ -355,7 +323,7 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): ) # Compute top `kk` eigenvalues and eigenvectors # Sort by eigenvalue magnitude. Note, eigsh does not return - # ordered eigenvalues/vectors. + # ordered eigenvalues/vectors for which="LM". idx = np.argsort(np.abs(pi))[::-1] pi = pi[idx] U = U[:, idx] @@ -438,6 +406,48 @@ def _cl_to_C(self, clmatrix): self.C = C self.C_t = np.ascontiguousarray(C.transpose(1, 0, 2)) + @staticmethod + def _compute_num_eigs(num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_inc): + """ + Compute number of eigenvalues to use when implementing adaptive projection. + + :param num_eigs_prev: Number of eigenvalues used in previous iteration of ADMM. + :param eig_vec: Eigenvector result from previous iteration of ADMM. + :param num_eigs_W: Number of eigenvalues used in previous computation of + ADMM subproblem for solving W. + :param rel_drp_thresh: Relative drop threshold for determining number of + eigenvalues to use. + :param eigs_inc: Number of eigenvalues to increase by if relative drop + threshold is not met. + + :return: Number of eigenvalues to use in current iteration. + """ + if num_eigs_prev > 0: + # Initialize relative drop + rel_drp = 0 + imx = 0 + # Calculate relative drop based on `eig_vec` + if len(eig_vec) == 2: + rel_drp = np.inf + elif len(eig_vec) > 2: + drops = eig_vec[:-1] / eig_vec[1:] + + # Find max drop + imx = np.argmax(drops) + dmx = drops[imx] + + # Relative drop + rel_drp = (num_eigs_W - 1) * dmx / (np.sum(drops) - dmx) + + # Update `num_eigs_prev` based on relative drop + num_eigs = ( + max(imx, 6) if rel_drp > rel_drp_thresh else num_eigs_prev + eigs_inc + ) + else: + num_eigs = 6 + + return num_eigs + @staticmethod def _compute_AX(X): """ From d98f59282330473b8868c0d7ef56d7631b60b48e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Mar 2025 09:27:31 -0400 Subject: [PATCH 045/216] Clarifying comments and unit test for compute_AX --- src/aspire/abinitio/commonline_lud.py | 21 +++++++++++++-------- tests/test_orient_lud.py | 15 +++++++++++++++ 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 9ca71fff18..4ba1da0e8b 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -451,13 +451,13 @@ def _compute_num_eigs(num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_i @staticmethod def _compute_AX(X): """ - Compute the application of the linear operator A to the input matrix X, - where A(X) is defined as: + Compute the application of the linear operator A to the symmetric input + matrix X, where A(X) is defined as: A(X) = [ X_ii^(11), X_ii^(22), - sqrt(2) X_ii^(12) + sqrt(2) X_ii^(21) + sqrt(2)/2 * X_ii^(12) + sqrt(2)/2 * X_ii^(21) ] i = 1, 2, ..., K @@ -467,16 +467,21 @@ def _compute_AX(X): :param X: 2D square array. :return: A(X) """ + # Extract the diagonal entries of X. + diags = np.diag(X) + + # Get row/column indices of the lower left entry of the 2x2 blocks + # along the diagonal of X. rows = np.arange(1, X.shape[0], 2) cols = np.arange(0, X.shape[0], 2) - # Create diagonal matrix with X on the main diagonal - diags = np.diag(X) - - # Compute the second part of AX + # Compute the second part of AX, which is sqrt(2)/2 times the sum of + # the off-diagonal entries of each 2x2 sub-block on the diagonal of X. + # Since each sub-block is symmetric, we take just one entry and multiply + # by sqrt(2). sqrt_2_X_col = np.sqrt(2, dtype=X.dtype) * X[rows, cols] - # Concatenate results vertically + # Form AX by concatenating the results. AX = np.concatenate((diags, sqrt_2_X_col)) return AX diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 3cd105f651..9f2f12b9d2 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -112,3 +112,18 @@ def test_estimate_rotations(source, orient_est): # Check dtype pass-through np.testing.assert_equal(source.dtype, est_rots.dtype) + + +def test_compute_AX(): + n = 3 + X = np.ones((2 * n, 2 * n)) + + # Create symmetric 2x2 blocks to go along the diagonal of X + X[:2, :2] = np.array([[1.0, 2.0], [2.0, 1.0]]) + X[2:4, 2:4] = np.array([[3.0, 4.0], [4.0, 3.0]]) + X[4:, 4:] = np.array([[5.0, 6.0], [6.0, 5.0]]) + + # Check the result. We should have: + AX = np.array([1, 1, 3, 3, 5, 5, np.sqrt(2) * 2, np.sqrt(2) * 4, np.sqrt(2) * 6]) + breakpoint() + np.testing.assert_allclose(CommonlineLUD._compute_AX(X), AX) From dc33d61c988e57baf7a69161bebe9a76afee4cc3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Mar 2025 09:44:15 -0400 Subject: [PATCH 046/216] Cleanup compute_AX. add docstring to test. --- src/aspire/abinitio/commonline_lud.py | 16 +++++++--------- tests/test_orient_lud.py | 19 ++++++++++++++++--- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 4ba1da0e8b..5f71f898b4 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -462,24 +462,22 @@ def _compute_AX(X): i = 1, 2, ..., K - where X_{ii}^{pq} denotes the (p,q)-th element in the 2x2 subblock X_{ii}. + where X_{ii}^{pq} denotes the (p,q)-th element in the 2x2 sub-block X_{ii}. - :param X: 2D square array. - :return: A(X) + :param X: 2D square array of shape (2K, 2K).. + :return: Flattened array representing A(X) """ - # Extract the diagonal entries of X. + # Extract the diagonal elements of (X_ii^(11) and X_ii^(22)) diags = np.diag(X) - # Get row/column indices of the lower left entry of the 2x2 blocks - # along the diagonal of X. - rows = np.arange(1, X.shape[0], 2) - cols = np.arange(0, X.shape[0], 2) + # Extract the off-diagonal elements from each 2x2 sub-block + off_diag_vals = np.diag(X, k=1)[::2] # Every other superdiagonal element # Compute the second part of AX, which is sqrt(2)/2 times the sum of # the off-diagonal entries of each 2x2 sub-block on the diagonal of X. # Since each sub-block is symmetric, we take just one entry and multiply # by sqrt(2). - sqrt_2_X_col = np.sqrt(2, dtype=X.dtype) * X[rows, cols] + sqrt_2_X_col = np.sqrt(2, dtype=X.dtype) * off_diag_vals # Form AX by concatenating the results. AX = np.concatenate((diags, sqrt_2_X_col)) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 9f2f12b9d2..9ce42a602a 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -115,8 +115,22 @@ def test_estimate_rotations(source, orient_est): def test_compute_AX(): - n = 3 - X = np.ones((2 * n, 2 * n)) + """ + Test we get the intended result for `_compute_AX()`, where A(X) is defined as: + + A(X) = [ + X_ii^(11), + X_ii^(22), + sqrt(2)/2 * X_ii^(12) + sqrt(2)/2 * X_ii^(21) + ] + + i = 1, 2, ..., K + + where X_{ii}^{pq} denotes the (p,q)-th element in the 2x2 sub-block X_{ii}. + """ + # We create a symmetric 2k x 2k matrix with specific 2x2 blocks along the diagonal + k = 3 + X = np.ones((2 * k, 2 * k)) # Create symmetric 2x2 blocks to go along the diagonal of X X[:2, :2] = np.array([[1.0, 2.0], [2.0, 1.0]]) @@ -125,5 +139,4 @@ def test_compute_AX(): # Check the result. We should have: AX = np.array([1, 1, 3, 3, 5, 5, np.sqrt(2) * 2, np.sqrt(2) * 4, np.sqrt(2) * 6]) - breakpoint() np.testing.assert_allclose(CommonlineLUD._compute_AX(X), AX) From b51f1784358da922c3df3e73e2ffc4d23d454571 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Mar 2025 15:16:44 -0400 Subject: [PATCH 047/216] refactor compute_num_eigs. Test compute_num_eigs. --- src/aspire/abinitio/commonline_lud.py | 34 +++++++++----------- tests/test_orient_lud.py | 45 +++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 20 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 5f71f898b4..760eb1da6e 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -422,27 +422,21 @@ def _compute_num_eigs(num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_i :return: Number of eigenvalues to use in current iteration. """ + if len(eig_vec) == 1: + # Handles the case were `drops` will be empty + return num_eigs_prev + eigs_inc + if num_eigs_prev > 0: - # Initialize relative drop - rel_drp = 0 - imx = 0 - # Calculate relative drop based on `eig_vec` - if len(eig_vec) == 2: - rel_drp = np.inf - elif len(eig_vec) > 2: - drops = eig_vec[:-1] / eig_vec[1:] - - # Find max drop - imx = np.argmax(drops) - dmx = drops[imx] - - # Relative drop - rel_drp = (num_eigs_W - 1) * dmx / (np.sum(drops) - dmx) - - # Update `num_eigs_prev` based on relative drop - num_eigs = ( - max(imx, 6) if rel_drp > rel_drp_thresh else num_eigs_prev + eigs_inc - ) + drops = eig_vec[:-1] / eig_vec[1:] + imx = np.argmax(drops) + dmx = drops[imx] + + # Relative drop + rel_drp = (num_eigs_W - 1) * dmx / (np.sum(drops) - dmx) + if rel_drp > rel_drp_thresh: + num_eigs = max(imx, 6) + else: + num_eigs = num_eigs_prev + eigs_inc else: num_eigs = 6 diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 9ce42a602a..b0bf36a3a5 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -114,6 +114,51 @@ def test_estimate_rotations(source, orient_est): np.testing.assert_equal(source.dtype, est_rots.dtype) +def test_compute_num_eigs(): + eig_vec = np.array([2, 2, 2, 2, 2, 2, 2, 2, 1]) # max drop = 2/1 at idx=7 + rel_drp_thresh = 5 + eigs_inc = 2 + num_eigs_W = 15 + + # Case for num_eigs_prev = 0: + num_eigs_prev = 0 + num_eigs = CommonlineLUD._compute_num_eigs( + num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_inc + ) + np.testing.assert_equal(num_eigs, 6) + + # Cases for num_eigs_prev > 0: + num_eigs_prev = 7 + + # For len(eig_vec) = 1 , output = num_eigs_prev + eigs_inc = 7 + 2 + ev = eig_vec[:1] + num_eigs = CommonlineLUD._compute_num_eigs( + num_eigs_prev, ev, num_eigs_W, rel_drp_thresh, eigs_inc + ) + np.testing.assert_equal(num_eigs, 9) + + # For len(eig_vec) = 2 , output = 6 + ev = eig_vec[:2] + num_eigs = CommonlineLUD._compute_num_eigs( + num_eigs_prev, ev, num_eigs_W, rel_drp_thresh, eigs_inc + ) + np.testing.assert_equal(num_eigs, 6) + + # For len(eig_vec) > 2 and rel_drp > rel_drp_thresh, output = num_eigs_prev + eigs_inc = 7 + 2 + # where rel_drp = (num_eigs_W - 1) * max_drop / (np.sum(drops) - max_drop) = 14 * 2 / (9 - 2) = 4 + num_eigs = CommonlineLUD._compute_num_eigs( + num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_inc + ) + np.testing.assert_equal(num_eigs, 9) + + # For len(eig_vec) > 2 and rel_drp <= rel_drp_thresh, output = index of max drop in elements of eig_vec = 7 + # Setting rel_drp_thresh = 3 (ie. < 4) + num_eigs = CommonlineLUD._compute_num_eigs( + num_eigs_prev, eig_vec, num_eigs_W, 3, eigs_inc + ) + np.testing.assert_equal(num_eigs, 7) + + def test_compute_AX(): """ Test we get the intended result for `_compute_AX()`, where A(X) is defined as: From f05c84cc9bc425edf7680bc53ccfad1342a22015 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 11:31:51 -0400 Subject: [PATCH 048/216] num images in gallery: 500 --- gallery/experiments/commonline_lud_simulated_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index 4fc471b7ec..fdea64bc2a 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -34,7 +34,7 @@ # for various levels of noise and output a table of results. SNR = [1 / 8, 1 / 16, 1 / 32] # Signal-to-noise ratio -n_imgs = 50 # Number of images in our source +n_imgs = 500 # Number of images in our source dtype = np.float64 pad_size = 129 results = { From 8313bf465aa0415b3c4e1e8e66afb64d0131bb3c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 13:09:23 -0400 Subject: [PATCH 049/216] simppler log message --- src/aspire/abinitio/commonline_lud.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 760eb1da6e..ac093d71b7 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -93,12 +93,7 @@ def __init__( raise ValueError("Spectral norm constraint, alpha, must be in [2/3, 1).") logger.info( - "Initializing LUD algorithm using ADMM" - + ( - f" with spectral norm constraint: {self.alpha}." - if self.alpha is not None - else " without spectral norm constraint." - ) + f"Initializing LUD algorithm using ADMM with spectral norm constraint: {self.alpha}" ) self.tol = tol From f44e7da67d60b3546c4d4c014f31172911e32b79 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 14:11:39 -0400 Subject: [PATCH 050/216] move deterministic_rounding back to SDP. make standalone --- src/aspire/abinitio/__init__.py | 6 +- src/aspire/abinitio/commonline_base.py | 107 +------------------------ src/aspire/abinitio/commonline_lud.py | 4 +- src/aspire/abinitio/commonline_sdp.py | 103 ++++++++++++++++++++++++ 4 files changed, 109 insertions(+), 111 deletions(-) diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index 1848f2b572..964c54bc37 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -1,9 +1,9 @@ from .commonline_base import CLOrient3D -from .commonline_lud import CommonlineLUD -from .commonline_sdp import CommonlineSDP -from .sync_voting import SyncVotingMixin # isort: off +from .commonline_sdp import CommonlineSDP +from .commonline_lud import CommonlineLUD +from .sync_voting import SyncVotingMixin from .commonline_sync import CLSyncVoting from .commonline_sync3n import CLSync3N from .commonline_c3_c4 import CLSymmetryC3C4 diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index 9585ea2736..fe456c645e 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -7,14 +7,7 @@ from aspire.image import Image from aspire.operators import PolarFT -from aspire.utils import ( - common_line_from_rots, - complex_type, - fuzzy_mask, - nearest_rotations, - tqdm, -) -from aspire.utils.matlab_compat import stable_eigsh +from aspire.utils import common_line_from_rots, complex_type, fuzzy_mask, tqdm from aspire.utils.random import choice logger = logging.getLogger(__name__) @@ -816,101 +809,3 @@ def __init_cupy_module(): backend="nvcc", options=("-O3", "--use_fast_math", "--extra-device-vectorization"), ) - - def _deterministic_rounding(self, gram): - """ - Deterministic rounding procedure to recover the rotations from the Gram matrix. - - The Gram matrix contains information about the first two columns of every rotation - matrix. These columns are extracted and used to form the remaining column of every - rotation matrix. - - :param gram: A 2n_img x 2n_img Gram matrix. - - :return: An n_img x 3 x 3 stack of rotation matrices. - """ - logger.info("Recovering rotations from Gram matrix.") - # Obtain top eigenvectors from Gram matrix. - d, v = stable_eigsh(gram, 5) - sort_idx = np.argsort(-d) - logger.info(f"Top 5 eigenvalues from (rank-3) Gram matrix: {d[sort_idx]}") - - # Only need the top 3 eigen-vectors. - v = v[:, sort_idx[:3]] - - # According to the structure of the Gram matrix, the first `n_img` rows, denoted v1, - # correspond to the linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of - # column 1 of all rotation matrices. Similarly, the second `n_img` rows of v, - # denoted v2, are linear combinations of R_{i}^{2}, i=1,...,K, that is, the second - # column of all rotation matrices. - v1 = v[: self.n_img].T - v2 = v[self.n_img : 2 * self.n_img].T - - # Use a least-squares method to get A.T*A and a Cholesky decomposition to find A. - A = self._ATA_solver(v1, v2) - - # Recover the rotations. The first two columns of all rotation - # matrices are given by unmixing V1 and V2 using A. The third - # column is the cross product of the first two. - r1 = np.dot(A.T, v1) - r2 = np.dot(A.T, v2) - r3 = np.cross(r1, r2, axis=0) - rotations = np.stack((r1.T, r2.T, r3.T), axis=-1) - - # Make sure that we got rotations by enforcing R to be - # a rotation (in case the error is large) - rotations = nearest_rotations(rotations) - - return rotations - - @staticmethod - def _ATA_solver(v1, v2): - """ - Uses a least squares method to solve for the linear transformation A - such that A*v1=R1 and A*v2=R2 correspond to the first and second columns - of a sequence of rotation matrices. - - :param v1: 3 x n_img array corresponding to linear combinations of the first - columns of all rotation matrices. - :param v2: 3 x n_img array corresponding to linear combinations of the second - columns of all rotation matrices. - - :return: 3x3 linear transformation mapping v1, v2 to first two columns of rotations. - """ - # We look for a linear transformation (3 x 3 matrix) A such that - # A*v1'=R1 and A*v2=R2 are the columns of the rotations matrices. - # Therefore: - # v1 * A'*A v1' = 1 - # v2 * A'*A v2' = 1 - # v1 * A'*A v2' = 0 - # These are 3*K linear equations for 9 matrix entries of A'*A - # Actually, there are only 6 unknown variables, because A'*A is symmetric. - # So we will truncate from 9 variables to 6 variables corresponding - # to the upper half of the matrix A'*A - n_img = v1.shape[-1] - truncated_equations = np.zeros((3 * n_img, 9), dtype=v1.dtype) - k = 0 - for i in range(3): - for j in range(3): - truncated_equations[0::3, k] = v1[i] * v1[j] - truncated_equations[1::3, k] = v2[i] * v2[j] - truncated_equations[2::3, k] = v1[i] * v2[j] - k += 1 - - # b = [1 1 0 1 1 0 ...]' is the right hand side vector - b = np.ones(3 * n_img, dtype=v1.dtype) - b[2::3] = 0 - - # Find the least squares approximation of A'*A in vector form - ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] - - # Construct the matrix A'*A from the vectorized matrix. - # Note, this is only the lower triangle of A'*A. - ATA = ATA_vec.reshape(3, 3) - - # The Cholesky decomposition of A'*A gives A (lower triangle). - # Note, that `np.linalg.cholesky()` only uses the lower-triangular - # and diagonal elements of ATA. - A = np.linalg.cholesky(ATA) - - return A diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index ac093d71b7..2664e7c2be 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -4,7 +4,7 @@ from scipy.sparse import csr_array from scipy.sparse.linalg import eigsh -from aspire.abinitio import CLOrient3D +from aspire.abinitio import CLOrient3D, CommonlineSDP logger = logging.getLogger(__name__) @@ -134,7 +134,7 @@ def estimate_rotations(self): self._cl_to_C(self.clmatrix) gram = self._compute_Gram() gram = self._restructure_Gram(gram) - self.rotations = self._deterministic_rounding(gram) + self.rotations = CommonlineSDP._deterministic_rounding(gram) return self.rotations diff --git a/src/aspire/abinitio/commonline_sdp.py b/src/aspire/abinitio/commonline_sdp.py index b006036be1..b2604f1e17 100644 --- a/src/aspire/abinitio/commonline_sdp.py +++ b/src/aspire/abinitio/commonline_sdp.py @@ -5,6 +5,8 @@ from scipy.sparse import csr_array from aspire.abinitio import CLOrient3D +from aspire.utils import nearest_rotations +from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) @@ -148,3 +150,104 @@ def _compute_gram_matrix(self, S, A, b): prob.solve() return G.value + + @staticmethod + def _deterministic_rounding(gram): + """ + Deterministic rounding procedure to recover the rotations from the Gram matrix. + + The Gram matrix contains information about the first two columns of every rotation + matrix. These columns are extracted and used to form the remaining column of every + rotation matrix. + + :param gram: A 2K x 2K Gram matrix. + + :return: An K x 3 x 3 stack of rotation matrices. + """ + K = gram.shape[0] // 2 + + logger.info("Recovering rotations from Gram matrix.") + # Obtain top eigenvectors from Gram matrix. + d, v = stable_eigsh(gram, 5) + sort_idx = np.argsort(-d) + logger.info(f"Top 5 eigenvalues from (rank-3) Gram matrix: {d[sort_idx]}") + + # Only need the top 3 eigen-vectors. + v = v[:, sort_idx[:3]] + + # According to the structure of the Gram matrix, the first `K` rows, denoted v1, + # correspond to the linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of + # column 1 of all rotation matrices. Similarly, the second `K` rows of v, + # denoted v2, are linear combinations of R_{i}^{2}, i=1,...,K, that is, the second + # column of all rotation matrices. + v1 = v[:K].T + v2 = v[K : 2 * K].T + + # Use a least-squares method to get A.T*A and a Cholesky decomposition to find A. + A = CommonlineSDP._ATA_solver(v1, v2) + + # Recover the rotations. The first two columns of all rotation + # matrices are given by unmixing V1 and V2 using A. The third + # column is the cross product of the first two. + r1 = np.dot(A.T, v1) + r2 = np.dot(A.T, v2) + r3 = np.cross(r1, r2, axis=0) + rotations = np.stack((r1.T, r2.T, r3.T), axis=-1) + + # Make sure that we got rotations by enforcing R to be + # a rotation (in case the error is large) + rotations = nearest_rotations(rotations) + + return rotations + + @staticmethod + def _ATA_solver(v1, v2): + """ + Uses a least squares method to solve for the linear transformation A + such that A*v1=R1 and A*v2=R2 correspond to the first and second columns + of a sequence of rotation matrices. + + :param v1: 3 x n_img array corresponding to linear combinations of the first + columns of all rotation matrices. + :param v2: 3 x n_img array corresponding to linear combinations of the second + columns of all rotation matrices. + + :return: 3x3 linear transformation mapping v1, v2 to first two columns of rotations. + """ + # We look for a linear transformation (3 x 3 matrix) A such that + # A*v1'=R1 and A*v2=R2 are the columns of the rotations matrices. + # Therefore: + # v1 * A'*A v1' = 1 + # v2 * A'*A v2' = 1 + # v1 * A'*A v2' = 0 + # These are 3*K linear equations for 9 matrix entries of A'*A + # Actually, there are only 6 unknown variables, because A'*A is symmetric. + # So we will truncate from 9 variables to 6 variables corresponding + # to the upper half of the matrix A'*A + n_img = v1.shape[-1] + truncated_equations = np.zeros((3 * n_img, 9), dtype=v1.dtype) + k = 0 + for i in range(3): + for j in range(3): + truncated_equations[0::3, k] = v1[i] * v1[j] + truncated_equations[1::3, k] = v2[i] * v2[j] + truncated_equations[2::3, k] = v1[i] * v2[j] + k += 1 + + # b = [1 1 0 1 1 0 ...]' is the right hand side vector + b = np.ones(3 * n_img, dtype=v1.dtype) + b[2::3] = 0 + + # Find the least squares approximation of A'*A in vector form + ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] + + # Construct the matrix A'*A from the vectorized matrix. + # Note, this is only the lower triangle of A'*A. + ATA = ATA_vec.reshape(3, 3) + + # The Cholesky decomposition of A'*A gives A (lower triangle). + # Note, that `np.linalg.cholesky()` only uses the lower-triangular + # and diagonal elements of ATA. + A = np.linalg.cholesky(ATA) + + return A From 8b7d82666e43c152192a8225b7247c64fad439a4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 14:26:29 -0400 Subject: [PATCH 051/216] collapse logic branches --- src/aspire/abinitio/commonline_lud.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 2664e7c2be..7f25e892f6 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -170,11 +170,10 @@ def _compute_Gram(self): # Initialize variables G = np.eye(n, dtype=self.dtype) W = np.eye(n, dtype=self.dtype) + Phi = W + G / self.mu if self.alpha is not None: Z = W - Phi = G / self.mu - else: - Phi = W + G / self.mu + Phi -= Z # Compute initial values S, theta = self._Q_theta(Phi) From 9c096478714a39a1ee93bfdeab704dc44cd8b85b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 14:40:15 -0400 Subject: [PATCH 052/216] clean up gallery result dict and table --- gallery/experiments/commonline_lud_simulated_data.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index fdea64bc2a..d3848e9b2a 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -34,13 +34,10 @@ # for various levels of noise and output a table of results. SNR = [1 / 8, 1 / 16, 1 / 32] # Signal-to-noise ratio -n_imgs = 500 # Number of images in our source +n_imgs = 50 # Number of images in our source dtype = np.float64 pad_size = 129 -results = { - "SNR": ["1/8", "1/16", "1/32"], - "Mean Angular Distance": [], -} # Dictionary to store results +results = {} # Dictionary to store results # %% # Load Volume Map @@ -87,7 +84,7 @@ mean_ang_dist = mean_aligned_angular_distance(est_rotations, src.rotations) # Store results. - results["Mean Angular Distance"].append(mean_ang_dist) + results[snr] = mean_ang_dist # %% # Display Results @@ -103,7 +100,7 @@ table.append(f"{'SNR':<{col1_width}} {'Mean Angular Distance':<{col2_width}}") table.append("-" * (col1_width + col2_width)) -for snr, angle in zip(results["SNR"], results["Mean Angular Distance"]): +for snr, angle in results.items(): table.append(f"{snr:<{col1_width}} {angle:<{col2_width}}") # Log the table From 00b7d8121cd60977461ee4a9c85c29756c5d8373 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 15:19:23 -0400 Subject: [PATCH 053/216] clean up ATy computation --- src/aspire/abinitio/commonline_lud.py | 31 ++++++++++----------------- 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 7f25e892f6..a8e3179155 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -497,30 +497,21 @@ def _compute_ATy(y): Y_ii^(22) = y_i^2, Y_ii^(12) = Y_ii^(21) = y_i^3 / sqrt(2). - :param y: 1D array of length 3 * n_img. - :return: Sparse matrix AT(y) + :param y: 1D NumPy array of length 3K. + :return: 2D NumPy array of shape (2K, 2K). """ - n = 2 * len(y) // 3 - m = len(y) - rows = np.concatenate([np.arange(1, n, 2), np.arange(0, n, 2)]) - cols = np.concatenate([np.arange(0, n, 2), np.arange(1, n, 2)]) - data = np.concatenate( - [ - (np.sqrt(2, dtype=y.dtype) / 2) * y[n:m], - (np.sqrt(2, dtype=y.dtype) / 2) * y[n:m], - ] - ) + K = len(y) // 3 + n = 2 * K # Size of the output matrix + ATy = np.zeros((n, n), dtype=y.dtype) - # Combine diagonal elements - diag_data = y[:n] - diag_idx = np.arange(n) + # Assign diagonal elements + ATy[::1, ::1] = np.diag(y[:n]) - # Construct the full matrix - data = np.concatenate([data, diag_data]) - rows = np.concatenate([rows, diag_idx]) - cols = np.concatenate([cols, diag_idx]) + # Assign symmetric off-diagonal elements + off_diag_vals = np.sqrt(2, dtype=y.dtype) / 2 * y[n:] + ATy[::2, 1::2] = np.diag(off_diag_vals) # Y_ii^(12) + ATy[1::2, ::2] = np.diag(off_diag_vals) # Y_ii^(21) - ATy = csr_array((data, (rows, cols)), shape=(n, n)) return ATy def _restructure_Gram(self, G): From ea5f45856e659fc43a04971af25522211d80636d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 15:27:04 -0400 Subject: [PATCH 054/216] Add ATy to unit test. --- tests/test_orient_lud.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index b0bf36a3a5..68adb1a352 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -159,7 +159,7 @@ def test_compute_num_eigs(): np.testing.assert_equal(num_eigs, 7) -def test_compute_AX(): +def test_compute_AX_ATy(): """ Test we get the intended result for `_compute_AX()`, where A(X) is defined as: @@ -172,10 +172,12 @@ def test_compute_AX(): i = 1, 2, ..., K where X_{ii}^{pq} denotes the (p,q)-th element in the 2x2 sub-block X_{ii}. + + Also test ATy. For a 2x2 block diagonal matrix, X, if A(X) = y, AT(y) = X. """ # We create a symmetric 2k x 2k matrix with specific 2x2 blocks along the diagonal k = 3 - X = np.ones((2 * k, 2 * k)) + X = np.zeros((2 * k, 2 * k)) # Create symmetric 2x2 blocks to go along the diagonal of X X[:2, :2] = np.array([[1.0, 2.0], [2.0, 1.0]]) @@ -183,5 +185,9 @@ def test_compute_AX(): X[4:, 4:] = np.array([[5.0, 6.0], [6.0, 5.0]]) # Check the result. We should have: - AX = np.array([1, 1, 3, 3, 5, 5, np.sqrt(2) * 2, np.sqrt(2) * 4, np.sqrt(2) * 6]) - np.testing.assert_allclose(CommonlineLUD._compute_AX(X), AX) + y_gt = np.array([1, 1, 3, 3, 5, 5, np.sqrt(2) * 2, np.sqrt(2) * 4, np.sqrt(2) * 6]) + y = CommonlineLUD._compute_AX(X) + np.testing.assert_allclose(y, y_gt) + + # Check ATy(y) = X + np.testing.assert_allclose(CommonlineLUD._compute_ATy(y), X) From bcf8a1d28850e77edd81f7e17fbec4dcca49f49d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 15:52:18 -0400 Subject: [PATCH 055/216] Fix warning --- src/aspire/abinitio/commonline_lud.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index a8e3179155..3235915f9c 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -426,7 +426,11 @@ def _compute_num_eigs(num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_i dmx = drops[imx] # Relative drop - rel_drp = (num_eigs_W - 1) * dmx / (np.sum(drops) - dmx) + rel_drp = ( + (num_eigs_W - 1) * dmx / (np.sum(drops) - dmx) + if len(drops) > 1 + else np.inf + ) if rel_drp > rel_drp_thresh: num_eigs = max(imx, 6) else: From 018f186566f5e1e19dd228c92aa09ab714598b31 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 16:02:08 -0400 Subject: [PATCH 056/216] use Fraction in gallery --- gallery/experiments/commonline_lud_simulated_data.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index d3848e9b2a..a5b27379b7 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -15,6 +15,7 @@ # ------- import logging +from fractions import Fraction import numpy as np @@ -33,8 +34,8 @@ # Set up some initializing parameters. We will run the LUD algorithm # for various levels of noise and output a table of results. -SNR = [1 / 8, 1 / 16, 1 / 32] # Signal-to-noise ratio -n_imgs = 50 # Number of images in our source +SNR = ["1/8", "1/16", "1/32"] # Signal-to-noise ratio +n_imgs = 500 # Number of images in our source dtype = np.float64 pad_size = 129 results = {} # Dictionary to store results @@ -64,7 +65,7 @@ # the ground truth rotations and the globally aligned estimated rotations. for snr in SNR: # Generate a white noise adder with specifid SNR. - noise_adder = WhiteNoiseAdder.from_snr(snr=snr) + noise_adder = WhiteNoiseAdder.from_snr(snr=Fraction(snr)) # Initialize a Simulation source to generate noisy, centered images. src = Simulation( From ff0d1a15f453d385db351151390756d1047b2a15 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Mar 2025 16:04:49 -0400 Subject: [PATCH 057/216] unused import --- src/aspire/abinitio/commonline_lud.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 3235915f9c..0c084d6c1e 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -1,7 +1,6 @@ import logging import numpy as np -from scipy.sparse import csr_array from scipy.sparse.linalg import eigsh from aspire.abinitio import CLOrient3D, CommonlineSDP From 5ddd34a26a391859dad22fd98362e954cffe6bba Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 17 Mar 2025 15:20:52 -0400 Subject: [PATCH 058/216] Fix test tol. Cleanup branch logic. --- src/aspire/abinitio/commonline_lud.py | 8 ++++---- tests/test_orient_lud.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 0c084d6c1e..61509461b4 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -169,10 +169,10 @@ def _compute_Gram(self): # Initialize variables G = np.eye(n, dtype=self.dtype) W = np.eye(n, dtype=self.dtype) - Phi = W + G / self.mu - if self.alpha is not None: - Z = W - Phi -= Z + Z = np.eye(n, dtype=self.dtype) + Phi = G / self.mu + if self.alpha is None: + Phi += W # Compute initial values S, theta = self._Q_theta(Phi) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 68adb1a352..e498badffe 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -100,8 +100,8 @@ def test_estimate_rotations(source, orient_est): # Register estimates to ground truth rotations and compute the # angular distance between them (in degrees). - # Assert that mean aligned angular distance is less than 3 degrees. - tol = 3 + # Assert that mean aligned angular distance is less than 5 degrees. + tol = 5 # Using LUD without spectral norm constraint, ie. alpha=None, # on shifted images reduces estimated rotations accuracy. From fd7b6e42f0875263bbef875d32987d4cb2953c85 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Mar 2025 14:23:10 -0400 Subject: [PATCH 059/216] gallery typos --- gallery/experiments/commonline_lud_simulated_data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index a5b27379b7..4883c575ae 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -60,11 +60,11 @@ # Generate Noisy Images and Estimate Rotations # -------------------------------------------- # A ``Simulation`` object is used to generate simulated data at various -# noise levels. Then rotations are estimated using ``CommonlineLUD` algorithm. +# noise levels. Then rotations are estimated using the ``CommonlineLUD`` algorithm. # Results are measured by computing the mean aligned angular distance between # the ground truth rotations and the globally aligned estimated rotations. for snr in SNR: - # Generate a white noise adder with specifid SNR. + # Generate a white noise adder with specified SNR. noise_adder = WhiteNoiseAdder.from_snr(snr=Fraction(snr)) # Initialize a Simulation source to generate noisy, centered images. From aee6441bd98a69cae669ddac00989976a0bd8ef1 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Mar 2025 08:45:59 -0400 Subject: [PATCH 060/216] commonline_lud review cleanup --- src/aspire/abinitio/commonline_lud.py | 99 ++++++++++++++------------- 1 file changed, 50 insertions(+), 49 deletions(-) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 61509461b4..1898df4792 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -23,7 +23,7 @@ def __init__( tol=1e-3, mu=1, gam=1.618, - EPS=1e-12, + eps=1e-12, maxit=1000, adp_proj=True, max_rankZ=None, @@ -54,7 +54,7 @@ def __init__( Default is 1. :param gam: Relaxation factor for updating variables in the algorithm (typically between 1 and 2). Default is 1.618. - :param EPS: Small positive value used to filter out negligible eigenvalues. + :param eps: Small positive value used to filter out negligible eigenvalues. Default is 1e-12. :param maxit: Maximum number of iterations allowed for the algorithm. Default is 1000. @@ -62,9 +62,9 @@ def __init__( - True: Adaptive rank selection (Default). - False: Full eigenvalue decomposition. :param max_rankZ: Maximum rank used for projecting the Z matrix (for adaptive projection). - Default is None (will be computed based on `n_img`). + If None, defaults to max(6, n_img // 2). :param max_rankW: Maximum rank used for projecting the W matrix (for adaptive projection). - Default is None (will be computed based on `n_img`). + If None, defaults to max(6, n_img // 2). :param adp_mu: Adaptive adjustment of the penalty parameter `mu`: - True: Enabled (Default). - False: Disabled. @@ -98,11 +98,9 @@ def __init__( self.tol = tol self.mu = mu self.gam = gam - self.EPS = EPS + self.eps = eps self.maxit = maxit self.adp_proj = adp_proj - self.max_rankZ = max_rankZ - self.max_rankW = max_rankW # Parameters for adjusting mu self.adp_mu = adp_mu @@ -118,6 +116,10 @@ def __init__( # Initialize commonline base class super().__init__(src, **kwargs) + # Adjust rank limits + self.max_rankZ = max_rankW or max(6, self.n_img // 2) + self.max_rankW = max_rankW or max(6, self.n_img // 2) + # Upper-triangular mask used in `_Q_theta` ut_mask = np.zeros((self.n_img, self.n_img), dtype=bool) ut_mask[np.triu_indices(self.n_img, k=1)] = True @@ -161,11 +163,6 @@ def _compute_Gram(self): [np.ones(n, dtype=self.dtype), np.zeros(self.n_img, dtype=self.dtype)] ) - # Adjust rank limits - self.max_rankW = self.max_rankW or max(6, self.n_img // 2) - if self.alpha is not None: - self.max_rankZ = self.max_rankZ or max(6, self.n_img // 2) - # Initialize variables G = np.eye(n, dtype=self.dtype) W = np.eye(n, dtype=self.dtype) @@ -179,12 +176,12 @@ def _compute_Gram(self): AS = self._compute_AX(S) resi = self._compute_AX(G) - b - nev = 0 + num_eigs = 0 itmu_pinf = 0 itmu_dinf = 0 - zz = 0 - kk = 0 - dH = 0 + eigs_Z = 0 + num_eigs_Z = None + eigs_H = 0 for itr in range(self.maxit): ############# # Compute y # @@ -206,7 +203,9 @@ def _compute_Gram(self): # Compute Z # ############# if self.alpha is not None: - Z, kk, zz = self._compute_Z(S, W, ATy, G, zz, itr, kk, nev) + Z, num_eigs_Z, eigs_Z = self._compute_Z( + S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs + ) ############# # Compute W # @@ -218,29 +217,29 @@ def _compute_Gram(self): if not self.adp_proj: D, V = np.linalg.eigh(H) - eigs_mask = D > self.EPS + eigs_mask = D > self.eps V = V[:, eigs_mask] W = V @ np.diag(D[eigs_mask]) @ V.T else: # Determine number of eigenvalues to compute for adaptive projection if itr == 0: - nev = self.max_rankW + num_eigs = self.max_rankW else: - nev = self._compute_num_eigs(nev, dH, nev, 50, 5) + num_eigs = self._compute_num_eigs(num_eigs, eigs_H, num_eigs, 50, 5) # If using a spectral norm constraint cap num_eigs at 2*n_img if self.alpha is not None: - nev = min(nev, n) + num_eigs = min(num_eigs, n) # Compute Eigenvectors and sort by largest algebraic eigenvalue - dH, V = eigsh(-H.astype(np.float64), k=nev, which="LA") - dH = dH[::-1].astype(self.dtype, copy=False) + eigs_H, V = eigsh(-H.astype(np.float64), k=num_eigs, which="LA") + eigs_H = eigs_H[::-1].astype(self.dtype, copy=False) V = V[:, ::-1].astype(self.dtype, copy=False) - nD = dH > self.EPS - dH = dH[nD] - nev = np.count_nonzero(nD) - W = V[:, nD] @ np.diag(dH) @ V[:, nD].T + H if nD.any() else H + nD = eigs_H > self.eps + eigs_H = eigs_H[nD] + num_eigs = np.count_nonzero(nD) + W = (V[:, nD] @ np.diag(eigs_H) @ V[:, nD].T + H) if nD.any() else H ############ # Update G # @@ -279,7 +278,7 @@ def _compute_Gram(self): return G - def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): + def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): """ Update ADMM subproblem for enforcing the spectral norm constraint. @@ -287,15 +286,14 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): :param W: A 2*n_img x 2*n_img array, primary ADMM subproblem matrix. :param ATy: A 2*n_img x 2*n_img array. :param G: Current value of the 2*n_img x 2*n_img optimization solution matrix. - :param zz: eigenvalues from previous iteration. - :param itr: ADMM loop iteration. - :param kk: Number of eigenvalues of Z to use to enforce spectral norm constraint. - :param nev: Number of eigenvalues of W used in previous iteration of ADMM. + :param eigs_Z: eigenvalues from previous iteration. + :param num_eigs_Z: Number of eigenvalues of Z to use to enforce spectral norm constraint. + :param num_eigs: Number of eigenvalues of W used in previous iteration of ADMM. :returns: - Z, Updated 2*n_img x 2*n_img matrix for spectral norm constraint ADMM subproblem. - - kk, Number of eigenvalues of Z to use to enforce spectral norm constraint in next iteration. - - nev, Number of eigenvalues of W to use in this iteration of ADMM. + - num_eigs_Z, Number of eigenvalues of Z to use to enforce spectral norm constraint in next iteration. + - num_eigs, Number of eigenvalues of W to use in this iteration of ADMM. """ lambda_ = self.alpha * self.n_img # Spectral norm bound B = S + W + ATy + G / self.mu @@ -305,15 +303,15 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): U, pi = np.linalg.eigh(B) else: # Determine number of eigenvalues to compute for adaptive projection - if itr == 0: - kk = self.max_rankZ + if num_eigs_Z is None: + num_eigs_Z = self.max_rankZ else: - kk = self._compute_num_eigs(kk, zz, nev, 10, 3) + num_eigs_Z = self._compute_num_eigs(num_eigs_Z, eigs_Z, num_eigs, 10, 3) - kk = min(kk, 2 * self.n_img) + num_eigs_Z = min(num_eigs_Z, 2 * self.n_img) pi, U = eigsh( - B.astype(np.float64, copy=False), k=kk, which="LM" - ) # Compute top `kk` eigenvalues and eigenvectors + B.astype(np.float64, copy=False), k=num_eigs_Z, which="LM" + ) # Compute top `num_eigs_Z` eigenvalues and eigenvectors # Sort by eigenvalue magnitude. Note, eigsh does not return # ordered eigenvalues/vectors for which="LM". @@ -322,16 +320,16 @@ def _compute_Z(self, S, W, ATy, G, zz, itr, kk, nev): U = U[:, idx] # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. - zz = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) - nD = zz > 0 - kk = np.count_nonzero(nD) - if kk > 0: - zz = zz[nD] - Z = U[:, nD] @ np.diag(zz) @ U[:, nD].T + eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) + nD = eigs_Z > 0 + num_eigs_Z = np.count_nonzero(nD) + if num_eigs_Z > 0: + eigs_Z = eigs_Z[nD] + Z = U[:, nD] @ np.diag(eigs_Z) @ U[:, nD].T else: Z = np.zeros_like(B) - return Z, kk, zz + return Z, num_eigs_Z, eigs_Z def _Q_theta(self, phi): """ @@ -379,10 +377,13 @@ def _Q_theta(self, phi): def _cl_to_C(self, clmatrix): """ For each pair of commonline indices cl[i, j] and cl[j, i], convert - from polar commonline indices to cartesion coordinates. + from polar commonline indices to Cartesion coordinates. + + This method sets the attribute `self.C` and its transpose `self.C_t`. + `self.C` is an n_img x n_img x 2 array where `self.C[i, j]` gives the + (x, y)-coordinates for cl[i, j]. :param clmatrix: n_img x n_img commonline matrix. - :return: n_img x n_img x 2 array of commonline cartesian coordinates. """ C = np.zeros((self.n_img, self.n_img, 2), dtype=self.dtype) for i in range(self.n_img): From db22ef8bb8d6d1f6eb22e1ffa3241ab11d1fca09 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Mar 2025 09:24:07 -0400 Subject: [PATCH 061/216] test adjoint property --- tests/test_orient_lud.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index e498badffe..3d33fb3cc2 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -191,3 +191,22 @@ def test_compute_AX_ATy(): # Check ATy(y) = X np.testing.assert_allclose(CommonlineLUD._compute_ATy(y), X) + + +def test_adjoint_property_A(dtype): + """ + Test = for random symmetric matrix `u` and + random vector `v`. + """ + n = 10 + u = np.random.rand(2 * n, 2 * n).astype(dtype, copy=False) + u = (u + u.T) / 2 + v = np.random.rand(3 * n).astype(dtype, copy=False) + + Au = CommonlineLUD._compute_AX(u) + ATv = CommonlineLUD._compute_ATy(v) + + lhs = np.dot(Au, v) + rhs = np.dot(u.flatten(), ATv.flatten()) + + np.testing.assert_allclose(lhs, rhs) From c62851149889f576b498588b1d3774ee39c110b4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Mar 2025 14:41:43 -0400 Subject: [PATCH 062/216] numpy allclose defaults --- tests/test_orient_lud.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_orient_lud.py b/tests/test_orient_lud.py index 3d33fb3cc2..2ecb7fc3b6 100644 --- a/tests/test_orient_lud.py +++ b/tests/test_orient_lud.py @@ -209,4 +209,4 @@ def test_adjoint_property_A(dtype): lhs = np.dot(Au, v) rhs = np.dot(u.flatten(), ATv.flatten()) - np.testing.assert_allclose(lhs, rhs) + np.testing.assert_allclose(lhs, rhs, rtol=1e-05, atol=1e-08) From 79a4618ef9eed36f28bd23f410a143cc2e385ba9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Apr 2025 12:01:58 -0400 Subject: [PATCH 063/216] remove global declaration for read only use of module level variable --- src/aspire/nufft/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/aspire/nufft/__init__.py b/src/aspire/nufft/__init__.py index 23ebe2c115..2a98e5365d 100644 --- a/src/aspire/nufft/__init__.py +++ b/src/aspire/nufft/__init__.py @@ -115,8 +115,6 @@ def all_backends(): :return: A list of strings representing available NFFT backends """ - global backends - if backends is None: check_backends(raise_errors=False) return [k for k, v in backends.items() if v is not None] From 4b85dd9485d33c59560c80294a3e706d3376a4e9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 22 Apr 2025 13:54:20 -0400 Subject: [PATCH 064/216] bump ubuntu version in CI --- .github/workflows/workflow.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml index a29105c790..bf908b655f 100644 --- a/.github/workflows/workflow.yml +++ b/.github/workflows/workflow.yml @@ -156,7 +156,7 @@ jobs: shell: bash -el {0} strategy: matrix: - os: [ubuntu-latest, ubuntu-20.04, macOS-latest, macOS-13] + os: [ubuntu-latest, ubuntu-22.04, macOS-latest, macOS-13] backend: [default, openblas] python-version: ['3.9'] include: From 90c0bb44f62e84f67595a7440d4d1a75f287bcc1 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Apr 2025 11:01:27 -0400 Subject: [PATCH 065/216] private Image._load for raw data. Allow rectangular micrographs. Test rectangular micrographs. --- src/aspire/image/image.py | 23 +++++++++++++++++++-- src/aspire/source/coordinates.py | 4 ++-- tests/test_coordinate_source.py | 35 ++++++++++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 4 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 92f43caac6..62db77a76c 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -612,7 +612,7 @@ def save(self, mrcs_filepath, overwrite=None): mrc.voxel_size = self.pixel_size @staticmethod - def load(filepath, dtype=None): + def _load(filepath, dtype=None): """ Load raw data from supported files. @@ -621,7 +621,9 @@ def load(filepath, dtype=None): :param filepath: File path (string). :param dtype: Optionally force cast to `dtype`. Default dtype is inferred from the file contents. - :return: numpy array of image data. + :returns: + - numpy array of image data. + - pixel size """ # Get the file extension @@ -640,6 +642,23 @@ def load(filepath, dtype=None): if dtype is not None: im = im.astype(dtype, copy=False) + return im, pixel_size + + @staticmethod + def load(filepath, dtype=None): + """ + Load raw data from supported files. + + Currently MRC and TIFF are supported. + + :param filepath: File path (string). + :param dtype: Optionally force cast to `dtype`. + Default dtype is inferred from the file contents. + :return: Image instance + """ + # Load raw data from filepath with pixel size + im, pixel_size = Image._load(filepath, dtype=dtype) + # Return as Image instance return Image(im, pixel_size=pixel_size) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index 21fec57a22..c74edef13b 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -308,7 +308,7 @@ def _get_mrc_shapes(self): mrc_shapes = np.zeros((self.num_micrographs, 2), dtype=int) for i, mrc in enumerate(self.mrc_paths): - mrc_shapes[i, :] = Image.load(mrc).resolution + mrc_shapes[i, :] = Image._load(mrc)[0].shape return mrc_shapes @@ -469,7 +469,7 @@ def _images(self, indices): # their origin micrograph for mrc_index, coord_list in grouped.items(): # Load file as 2D numpy array. - arr = Image.load(self.mrc_paths[mrc_index]).asnumpy()[0] + arr = Image._load(self.mrc_paths[mrc_index])[0] # create iterable of the coordinates in this mrc # we don't need to worry about exhausting this iter diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index c38feff939..48f2a1e31f 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -696,3 +696,38 @@ def testCommand(self): self.assertTrue(result_coord.exit_code == 0) self.assertTrue(result_star.exit_code == 0) self.assertTrue(result_preprocess.exit_code == 0) + + +def create_test_rectangular_micrograph_and_star(tmp_path): + # Create a rectangular micrograph (e.g., 128x256) + data = np.random.rand(128, 256).astype(np.float32) + mrc_path = tmp_path / "test_micrograph.mrc" + + with mrcfile.new(mrc_path, overwrite=True) as mrc: + mrc.set_data(data) + + # Two sample coordinates + coordinates = [(50.0, 30.0), (200.0, 100.0)] + + # Write a simple STAR file + star_path = tmp_path / "test_coordinates.star" + with open(star_path, "w") as f: + f.write("data_particles\n\n") + f.write("loop_\n") + f.write("_rlnCoordinateX #1\n") + f.write("_rlnCoordinateY #2\n") + for x, y in coordinates: + f.write(f"{x:.1f} {y:.1f}\n") + + return mrc_path, star_path + + +def test_restangular_coordinate_source(tmp_path): + mrc_file, star_file = create_test_rectangular_micrograph_and_star(tmp_path) + file_list = [(mrc_file, star_file)] + + # Check we can instantiate a CoordinateSource with a rectangular micrograph. + coord_src = CentersCoordinateSource(file_list, particle_size=32) + + # Check we can access images. + _ = coord_src.images[:] From 6d9a73b3cfcb42aafea31fe74941b386efdf407a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Apr 2025 14:30:22 -0400 Subject: [PATCH 066/216] Switch from raise to log warning for voxel size mismatch --- src/aspire/image/image.py | 2 +- tests/test_coordinate_source.py | 26 +++++++++++++++++++++----- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 62db77a76c..7cdf5c29bd 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -886,7 +886,7 @@ def _vx_array_to_size(vx): # checks uniformity. if isinstance(vx, np.recarray): if vx.x != vx.y: - raise ValueError(f"Voxel sizes are not uniform: {vx}") + logger.warning(f"Voxel sizes are not uniform: {vx}") vx = vx.x # Convert `0` to `None` diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index 48f2a1e31f..ea337185c3 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -1,3 +1,4 @@ +import logging import os import random import shutil @@ -17,6 +18,8 @@ from aspire.storage import StarFile from aspire.utils import RelionStarFile, importlib_path +logger = logging.getLogger(__name__) + class CoordinateSourceTestCase(TestCase): def setUp(self): @@ -698,13 +701,14 @@ def testCommand(self): self.assertTrue(result_preprocess.exit_code == 0) -def create_test_rectangular_micrograph_and_star(tmp_path): +def create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=(2.0, 2.0, 1.0)): # Create a rectangular micrograph (e.g., 128x256) data = np.random.rand(128, 256).astype(np.float32) mrc_path = tmp_path / "test_micrograph.mrc" with mrcfile.new(mrc_path, overwrite=True) as mrc: mrc.set_data(data) + mrc.voxel_size = voxel_size # Two sample coordinates coordinates = [(50.0, 30.0), (200.0, 100.0)] @@ -719,15 +723,27 @@ def create_test_rectangular_micrograph_and_star(tmp_path): for x, y in coordinates: f.write(f"{x:.1f} {y:.1f}\n") - return mrc_path, star_path + # Pack files into a list of tuples for consumption by CoordinatSource + file_list = [(mrc_path, star_path)] + + return file_list -def test_restangular_coordinate_source(tmp_path): - mrc_file, star_file = create_test_rectangular_micrograph_and_star(tmp_path) - file_list = [(mrc_file, star_file)] +def test_rectangular_coordinate_source(tmp_path): + file_list = create_test_rectangular_micrograph_and_star(tmp_path) # Check we can instantiate a CoordinateSource with a rectangular micrograph. coord_src = CentersCoordinateSource(file_list, particle_size=32) # Check we can access images. _ = coord_src.images[:] + + +def test_coordinate_source_pixel_warning(tmp_path, caplog): + # Create micrograph with mismatched pixel dimensions. + vx = (2.3, 2.1, 1.0) + file_list = create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=vx) + with caplog.at_level(logging.WARNING): + caplog.clear() + _ = CentersCoordinateSource(file_list, particle_size=32) + assert f"Voxel sizes are not uniform: {vx}" in caplog.text From 57d085ff0e9dead7ad4d46e274db47ff52ec9fc3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Apr 2025 16:12:41 -0400 Subject: [PATCH 067/216] remove stack axis for singleton Image.save --- src/aspire/image/image.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 7cdf5c29bd..bb07410c51 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -598,6 +598,10 @@ def save(self, mrcs_filepath, overwrite=None): if self.stack_ndim > 1: raise NotImplementedError("`save` is currently limited to 1D image stacks.") + data = self._data.astype(np.float32) + if self.n_images == 1: + data = data[0] + if overwrite is None and os.path.exists(mrcs_filepath): # If the file exists, append a timestamp to the old file and rename it _ = rename_with_timestamp(mrcs_filepath) @@ -606,7 +610,7 @@ def save(self, mrcs_filepath, overwrite=None): with mrcfile.new(mrcs_filepath, overwrite=overwrite) as mrc: # original input format (the image index first) - mrc.set_data(self._data.astype(np.float32)) + mrc.set_data(data) # Note assigning voxel_size must come after `set_data` if self.pixel_size is not None: mrc.voxel_size = self.pixel_size From d4ae0c07b17f8b7509df218e9c06939706fec470 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 2 Apr 2025 08:45:34 -0400 Subject: [PATCH 068/216] Cleanup warning in caplog --- tests/test_coordinate_source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index ea337185c3..58353c3f16 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -744,6 +744,6 @@ def test_coordinate_source_pixel_warning(tmp_path, caplog): vx = (2.3, 2.1, 1.0) file_list = create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=vx) with caplog.at_level(logging.WARNING): - caplog.clear() _ = CentersCoordinateSource(file_list, particle_size=32) assert f"Voxel sizes are not uniform: {vx}" in caplog.text + caplog.clear() From 6cf1c24183128517c796380804974e34dbd101c9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 2 Apr 2025 09:14:53 -0400 Subject: [PATCH 069/216] remove floating points from assert message --- tests/test_coordinate_source.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index 58353c3f16..9296777832 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -745,5 +745,4 @@ def test_coordinate_source_pixel_warning(tmp_path, caplog): file_list = create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=vx) with caplog.at_level(logging.WARNING): _ = CentersCoordinateSource(file_list, particle_size=32) - assert f"Voxel sizes are not uniform: {vx}" in caplog.text - caplog.clear() + assert f"Voxel sizes are not uniform" in caplog.text From ce051378f33da2b6254dd302d9649318b3cf4fa0 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 2 Apr 2025 09:16:16 -0400 Subject: [PATCH 070/216] remove f-string --- tests/test_coordinate_source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index 9296777832..5304ef0b74 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -745,4 +745,4 @@ def test_coordinate_source_pixel_warning(tmp_path, caplog): file_list = create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=vx) with caplog.at_level(logging.WARNING): _ = CentersCoordinateSource(file_list, particle_size=32) - assert f"Voxel sizes are not uniform" in caplog.text + assert "Voxel sizes are not uniform" in caplog.text From 9613a2bcd8c85b525d59d53fa7a82aafd7050449 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 22 Apr 2025 08:31:41 -0400 Subject: [PATCH 071/216] _load_raw --- src/aspire/image/image.py | 4 ++-- src/aspire/source/coordinates.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index bb07410c51..0395126bb0 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -616,7 +616,7 @@ def save(self, mrcs_filepath, overwrite=None): mrc.voxel_size = self.pixel_size @staticmethod - def _load(filepath, dtype=None): + def _load_raw(filepath, dtype=None): """ Load raw data from supported files. @@ -661,7 +661,7 @@ def load(filepath, dtype=None): :return: Image instance """ # Load raw data from filepath with pixel size - im, pixel_size = Image._load(filepath, dtype=dtype) + im, pixel_size = Image._load_raw(filepath, dtype=dtype) # Return as Image instance return Image(im, pixel_size=pixel_size) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index c74edef13b..495f7a87f6 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -308,7 +308,7 @@ def _get_mrc_shapes(self): mrc_shapes = np.zeros((self.num_micrographs, 2), dtype=int) for i, mrc in enumerate(self.mrc_paths): - mrc_shapes[i, :] = Image._load(mrc)[0].shape + mrc_shapes[i, :] = Image._load_raw(mrc)[0].shape return mrc_shapes @@ -469,7 +469,7 @@ def _images(self, indices): # their origin micrograph for mrc_index, coord_list in grouped.items(): # Load file as 2D numpy array. - arr = Image._load(self.mrc_paths[mrc_index])[0] + arr = Image._load_raw(self.mrc_paths[mrc_index])[0] # create iterable of the coordinates in this mrc # we don't need to worry about exhausting this iter From 1ee71fe97a1189f1cdcb95203210f10dccf74df2 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 28 Mar 2025 09:29:09 -0400 Subject: [PATCH 072/216] legacy flag for image/volume downsample --- src/aspire/image/image.py | 21 +++++++++++++++------ src/aspire/volume/volume.py | 17 +++++++++++++---- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 0395126bb0..dd3c5719a5 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -510,7 +510,7 @@ def legacy_whiten(self, psd, delta): return Image(res) - def downsample(self, ds_res, zero_nyquist=True): + def downsample(self, ds_res, zero_nyquist=True, legacy=False): """ Downsample Image to a specific resolution. This method returns a new Image. @@ -518,6 +518,8 @@ def downsample(self, ds_res, zero_nyquist=True): of this Image :param zero_nyquist: Option to keep or remove Nyquist frequency for even resolution. Defaults to zero_nyquist=True, removing the Nyquist frequency. + :param legacy: Option to match legacy Matlab downsample method. + Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions. :return: The downsampled Image object. """ @@ -528,19 +530,26 @@ def downsample(self, ds_res, zero_nyquist=True): # because all of the subsequent calls until `asnumpy` are GPU # when xp and fft in `cupy` mode. - # compute FT with centered 0-frequency - fx = fft.centered_fft2(xp.asarray(im._data)) + if legacy: + fx = fft.fftshift(fft.fft2(xp.asarray(im._data))) + else: + # compute FT with centered 0-frequency + fx = fft.centered_fft2(xp.asarray(im._data)) + # crop 2D Fourier transform for each image crop_fx = crop_pad_2d(fx, ds_res) # If downsampled resolution is even, optionally zero out the nyquist frequency. - if ds_res % 2 == 0 and zero_nyquist is True: + if ds_res % 2 == 0 and zero_nyquist and not legacy: crop_fx[:, 0, :] = 0 crop_fx[:, :, 0] = 0 # take back to real space, discard complex part, and scale - out = fft.centered_ifft2(crop_fx).real * (ds_res**2 / self.resolution**2) - out = xp.asnumpy(out) + if legacy: + out = fft.ifft2(fft.ifftshift(crop_fx)) + else: + out = fft.centered_ifft2(crop_fx) + out = xp.asnumpy(out.real * ds_res**2 / self.resolution**2) # Optionally scale pixel size ds_pixel_size = self.pixel_size diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py index 013642bba3..09373fc9e1 100644 --- a/src/aspire/volume/volume.py +++ b/src/aspire/volume/volume.py @@ -517,7 +517,7 @@ def flip(self, axis=-3): symmetry_group=symmetry, ) - def downsample(self, ds_res, mask=None, zero_nyquist=True): + def downsample(self, ds_res, mask=None, zero_nyquist=True, legacy=False): """ Downsample each volume to a desired resolution (only cubic supported). @@ -525,19 +525,25 @@ def downsample(self, ds_res, mask=None, zero_nyquist=True): :param zero_nyquist: Option to keep or remove Nyquist frequency for even resolution. Defaults to zero_nyquist=True, removing the Nyquist frequency. :param mask: Optional NumPy array mask to multiply in Fourier space. + :param legacy: Option to match legacy Matlab downsample method. + Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions. + :return: The downsampled Volume object. """ original_stack_shape = self.stack_shape v = self.stack_reshape(-1) # take 3D Fourier transform of each volume in the stack - fx = fft.centered_fftn(xp.asarray(v._data)) + if legacy: + fx = fft.fftshift(fft.fftn(xp.asarray(v._data))) + else: + fx = fft.centered_fftn(xp.asarray(v._data)) # crop each volume to the desired resolution in frequency space fx = crop_pad_3d(fx, ds_res) # If downsample resolution is even, optionally zero out the nyquist frequency. - if ds_res % 2 == 0 and zero_nyquist is True: + if ds_res % 2 == 0 and zero_nyquist and not legacy: fx[:, 0, :, :] = 0 fx[:, :, 0, :] = 0 fx[:, :, :, 0] = 0 @@ -547,7 +553,10 @@ def downsample(self, ds_res, mask=None, zero_nyquist=True): fx = fx * xp.asarray(mask) # inverse Fourier transform of each volume - out = fft.centered_ifftn(fx) + if legacy: + out = fft.ifftn(fft.ifftshift(fx)) + else: + out = fft.centered_ifftn(fx) out = out.real * (ds_res**3 / self.resolution**3) # Optionally scale pixel size From c6a75c4da258d45b0ff6400e39c5795d2a817d7a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 28 Mar 2025 10:13:31 -0400 Subject: [PATCH 073/216] test project/dowsample property with legacy downsample. --- tests/test_downsample.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/tests/test_downsample.py b/tests/test_downsample.py index 6c1cac82dd..284ee741a2 100644 --- a/tests/test_downsample.py +++ b/tests/test_downsample.py @@ -124,6 +124,7 @@ def test_integer_offsets(): DTYPES = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] RES = [65, 66] RES_DS = [32, 33] +LEGACY = [True, False] @pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module") @@ -141,6 +142,11 @@ def res_ds(request): return request.param +@pytest.fixture(params=LEGACY, ids=lambda x: f"legacy={x}", scope="module") +def legacy(request): + return request.param + + @pytest.fixture(scope="module") def emdb_vol(): return emdb_2660() @@ -153,17 +159,20 @@ def volume(emdb_vol, res, dtype): return vol -def test_downsample_project(volume, res_ds): +def test_downsample_project(volume, res_ds, legacy): """ Test that vol.downsample.project == vol.project.downsample. """ rot = np.eye(3, dtype=volume.dtype) # project along z-axis - im_ds_proj = volume.downsample(res_ds).project(rot) - im_proj_ds = volume.project(rot).downsample(res_ds) + im_ds_proj = volume.downsample(res_ds, legacy=legacy).project(rot) + im_proj_ds = volume.project(rot).downsample(res_ds, legacy=legacy) + + tol = 1e-09 + if volume.dtype == np.float32: + tol = 1e-07 + if legacy: + tol = 1e-03 - tol = 1e-07 - if volume.dtype == np.float64: - tol = 1e-09 np.testing.assert_allclose(im_ds_proj, im_proj_ds, atol=tol) From 105aa8cc151da648ce879a436ab38a7812d3134b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 28 Mar 2025 11:02:14 -0400 Subject: [PATCH 074/216] test legacy downsample --- tests/test_downsample.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/test_downsample.py b/tests/test_downsample.py index 284ee741a2..178274da47 100644 --- a/tests/test_downsample.py +++ b/tests/test_downsample.py @@ -176,6 +176,40 @@ def test_downsample_project(volume, res_ds, legacy): np.testing.assert_allclose(im_ds_proj, im_proj_ds, atol=tol) +def test_downsample_legacy(volume, res_ds): + """ + The legay Matlab downsample method differs from ASPIRE-Python + downsample in that is uses a different centering convention, + off by a half pixel for odd images, and does not zero out + the nyquist frequency. By making these alterations to the + ASPIRE-Python downsampled images we can match legacy downsample + upt to `allclose`. + """ + n_img = 10 + dtype = volume.dtype + src = Simulation( + n=n_img, + vols=volume, + amplitudes=1, + dtype=dtype, + seed=1980, + ) + ims = src.images[:] + + # Legacy downsampled images. + ims_ds_legacy = ims.downsample(res_ds, legacy=True) + + # ASPIRE-Python downsample with centering adjustments for odd resolution images. + shifts = 0.5 * np.ones((n_img, 2), dtype=dtype) + if src.L % 2 == 1: + ims = ims.shift(shifts) + ims_ds_py = ims.downsample(res_ds, zero_nyquist=False) + if res_ds % 2 == 1: + ims_ds_py = ims_ds_py.shift(-shifts) + + np.testing.assert_allclose(ims_ds_legacy, ims_ds_py, atol=1e-08) + + def test_simulation_relion_downsample(): """ Test that Simulation.downsample corresponds to RelionSource.downsample From 1d3067552ba6a0e6e8a9cf2cbcee4b060c71735f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 2 Apr 2025 10:12:11 -0400 Subject: [PATCH 075/216] legacy flag smoke test. docstring cleanup. --- src/aspire/image/image.py | 6 +++--- src/aspire/volume/volume.py | 6 +++--- tests/test_downsample.py | 2 ++ tests/test_volume.py | 30 ++++++++++++++++++++++++++++++ 4 files changed, 38 insertions(+), 6 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index dd3c5719a5..27d5ced568 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -516,9 +516,9 @@ def downsample(self, ds_res, zero_nyquist=True, legacy=False): :param ds_res: int - new resolution, should be <= the current resolution of this Image - :param zero_nyquist: Option to keep or remove Nyquist frequency for even resolution. - Defaults to zero_nyquist=True, removing the Nyquist frequency. - :param legacy: Option to match legacy Matlab downsample method. + :param zero_nyquist: Option to keep or remove Nyquist frequency for even + resolution (boolean). Defaults to zero_nyquist=True, removing the Nyquist frequency. + :param legacy: Option to match legacy Matlab downsample method (boolean). Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions. :return: The downsampled Image object. """ diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py index 09373fc9e1..8339d8c28a 100644 --- a/src/aspire/volume/volume.py +++ b/src/aspire/volume/volume.py @@ -522,10 +522,10 @@ def downsample(self, ds_res, mask=None, zero_nyquist=True, legacy=False): Downsample each volume to a desired resolution (only cubic supported). :param ds_res: Desired resolution. - :param zero_nyquist: Option to keep or remove Nyquist frequency for even resolution. - Defaults to zero_nyquist=True, removing the Nyquist frequency. + :param zero_nyquist: Option to keep or remove Nyquist frequency for even + resolution (boolean). Defaults to zero_nyquist=True, removing the Nyquist frequency. :param mask: Optional NumPy array mask to multiply in Fourier space. - :param legacy: Option to match legacy Matlab downsample method. + :param legacy: Option to match legacy Matlab downsample method (boolean). Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions. :return: The downsampled Volume object. """ diff --git a/tests/test_downsample.py b/tests/test_downsample.py index 178274da47..6155d68fd0 100644 --- a/tests/test_downsample.py +++ b/tests/test_downsample.py @@ -171,6 +171,8 @@ def test_downsample_project(volume, res_ds, legacy): if volume.dtype == np.float32: tol = 1e-07 if legacy: + # project does not enforce legacy centering convention, + # so this property will not hold up to allclose tolerance. tol = 1e-03 np.testing.assert_allclose(im_ds_proj, im_proj_ds, atol=tol) diff --git a/tests/test_volume.py b/tests/test_volume.py index 42ed4675f7..db6fae82d5 100644 --- a/tests/test_volume.py +++ b/tests/test_volume.py @@ -11,6 +11,7 @@ from numpy import pi from pytest import raises, skip +from aspire.downloader import emdb_2660 from aspire.source import _LegacySimulation from aspire.utils import Rotation, anorm, grid_2d, powerset, utest_tolerance from aspire.volume import ( @@ -32,6 +33,7 @@ def res_id(params): RES = [42, 43] +RES_DS = [32, 33] TEST_PX_SZ = 4.56 @@ -40,6 +42,11 @@ def res(request): return request.param +@pytest.fixture(params=RES_DS, ids=lambda x: f"resolution_ds={x}", scope="module") +def res_ds(request): + return request.param + + def dtype_id(params): return f"dtype={params}" @@ -122,6 +129,18 @@ def vols_hot_cold(res, dtype): return vols, hot_cold_locs, vol_center +@pytest.fixture(scope="module") +def emdb_vol(): + return emdb_2660() + + +@pytest.fixture(scope="module") +def volume(emdb_vol, res, dtype): + vol = emdb_vol.astype(dtype, copy=False) + vol = vol.downsample(res) + return vol + + @pytest.fixture def random_data(res, dtype): return np.random.randn(res, res, res).astype(dtype) @@ -673,6 +692,17 @@ def test_downsample(res): ) +def test_downsample_legacy(volume, res_ds): + """ + Smoke test for the downsample legacy flag. + """ + # Legacy downsampled images. + vol_ds = volume.downsample(res_ds, legacy=True) + + # Check downsampled volume resolution. + np.testing.assert_equal(vol_ds.resolution, res_ds) + + def test_shape(vols_1, res): assert vols_1.shape == (N, res, res, res) assert vols_1.stack_shape == (N,) From df82ff307226d0d4424ca25596d9fff367fe220b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 4 Apr 2025 10:15:04 -0400 Subject: [PATCH 076/216] add full set of parameters to ImageSource.downsample --- src/aspire/image/xform.py | 18 ++++++++++++++++-- src/aspire/source/image.py | 6 ++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py index 4710f02d56..10fddf1df2 100644 --- a/src/aspire/image/xform.py +++ b/src/aspire/image/xform.py @@ -199,12 +199,26 @@ class Downsample(LinearXform): A Xform that downsamples an Image object to a resolution specified by this Xform's resolution. """ - def __init__(self, resolution): + def __init__(self, resolution, zero_nyquist=True, legacy=False): + """ + Initialize Xform to downsample Image to a specific resolution. + + :param resolution: int - new resolution, should be <= the current resolution + of this Image + :param zero_nyquist: Option to keep or remove Nyquist frequency for even + resolution (boolean). Defaults to zero_nyquist=True, removing the Nyquist frequency. + :param legacy: Option to match legacy Matlab downsample method (boolean). + Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions. + """ self.resolution = resolution + self.zero_nyquist = zero_nyquist + self.legacy = legacy super().__init__() def _forward(self, im, indices): - return im.downsample(self.resolution) + return im.downsample( + self.resolution, zero_nyquist=self.zero_nyquist, legacy=self.legacy + ) def _adjoint(self, im, indices): # TODO: Implement up-sampling with zero-padding diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py index c6fe79aaea..f7bbfa4cd8 100644 --- a/src/aspire/source/image.py +++ b/src/aspire/source/image.py @@ -769,14 +769,16 @@ def _images(self, indices): """ @_as_copy - def downsample(self, L): + def downsample(self, L, zero_nyquist=True, legacy=False): if L > self.L: raise ValueError( "Max desired resolution {L} should be less than the current resolution {self.L}." ) logger.info(f"Setting max. resolution of source = {L}") - self.generation_pipeline.add_xform(Downsample(resolution=L)) + self.generation_pipeline.add_xform( + Downsample(resolution=L, zero_nyquist=zero_nyquist, legacy=legacy) + ) ds_factor = self.L / L self.unique_filters = [f.scale(ds_factor) for f in self.unique_filters] From 5e7413a56aeab011411a9371f721c5e407726b62 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 8 Apr 2025 10:59:28 -0400 Subject: [PATCH 077/216] fix testing logic branch. --- tests/test_downsample.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_downsample.py b/tests/test_downsample.py index 6155d68fd0..525dba1de6 100644 --- a/tests/test_downsample.py +++ b/tests/test_downsample.py @@ -170,10 +170,10 @@ def test_downsample_project(volume, res_ds, legacy): tol = 1e-09 if volume.dtype == np.float32: tol = 1e-07 - if legacy: - # project does not enforce legacy centering convention, - # so this property will not hold up to allclose tolerance. - tol = 1e-03 + if legacy: + # project does not enforce legacy centering convention, + # so this property will not hold up to allclose tolerance. + tol = 1e-03 np.testing.assert_allclose(im_ds_proj, im_proj_ds, atol=tol) From 9fd54fd787eaf74fc2e6c55ab274b750c699921a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 10 Apr 2025 11:40:02 -0400 Subject: [PATCH 078/216] update pyproject license syntax --- pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d07c986643..b78d6ffed1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,8 @@ version = "0.13.2" description = "Algorithms for Single Particle Reconstruction" readme = "README.md" # Optional requires-python = ">=3.9" -license = {file = "LICENSE"} +license = "GPL-3.0-or-later" +license-files = ["LICENSE"] maintainers = [ {name = "ASPIRE Developers", email = "ASPIRE-DEVS@princeton.edu"} ] @@ -23,8 +24,7 @@ authors = [ classifiers = [ "Development Status :: 3 - Alpha", - "Programming Language :: Python", - "License :: OSI Approved :: GNU General Public License v3 (GPLv3)" + "Programming Language :: Python" ] dependencies = [ From dbad10ef34a622f014db396dce7a7e8f50ceec93 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 23 Apr 2025 10:14:05 -0400 Subject: [PATCH 079/216] GPL-3.0-only --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index b78d6ffed1..66950c1770 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ version = "0.13.2" description = "Algorithms for Single Particle Reconstruction" readme = "README.md" # Optional requires-python = ">=3.9" -license = "GPL-3.0-or-later" +license = "GPL-3.0-only" license-files = ["LICENSE"] maintainers = [ {name = "ASPIRE Developers", email = "ASPIRE-DEVS@princeton.edu"} From 28b55e0ba55b244a2100db6cb4b293a9bdbed5cc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Apr 2025 15:11:48 -0400 Subject: [PATCH 080/216] Add seed param to BOT_align --- src/aspire/utils/bot_align.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/aspire/utils/bot_align.py b/src/aspire/utils/bot_align.py index fc8fbe4fdc..d3df896954 100644 --- a/src/aspire/utils/bot_align.py +++ b/src/aspire/utils/bot_align.py @@ -11,7 +11,7 @@ from numpy.linalg import norm from scipy.optimize import minimize -from aspire.utils.rotation import Rotation +from aspire.utils.rotation import Random, Rotation # Store parameters specific to each loss_type. # `lengthscale` is used to scale the modeled covariance @@ -38,6 +38,7 @@ def align_BO( surrogate_min_step=0.1, verbosity=0, dtype=None, + seed=None, ): """ This function returns a rotation matrix R that best aligns vol_ref with the rotated version of vol_given. @@ -61,8 +62,13 @@ def align_BO( :param verbosity: Surrogate problem optimization detail level. integer, defaults 0 (silent). 2 is most verbose. :param dtype: Numeric dtype to perform computations with. Default `None` infers dtype from `vol_ref`. + :param seed: Random seed for reproducible results. Integer, defaults None. :return: Rotation matrix R_init (without refinement) or (R_init, R_est) (with refinement). """ + # Generate sequence of seeds for pymanopt initial points + ss = np.random.SeedSequence(seed) + point_seeds = ss.generate_state(max_iters, dtype=np.uint32) + # Avoid utils/operators/utils circular import from aspire.operators import wemd_embed @@ -167,7 +173,14 @@ def euclidean_grad(new, t=t, q=q): min_step_size=surrogate_min_step, verbosity=verbosity, ) - result = optimizer.run(problem) + + # If provided, use seed to set initial point for optimizer + if seed is not None: + with Random(int(point_seeds[t])): + initial = manifold.random_point() + result = optimizer.run(problem, initial_point=initial) + else: + result = optimizer.run(problem) R_new = result.point.astype(dtype, copy=False) loss[t] = loss_fun(R_new) From 6ad3ad0632a771a9af45478c5e97a9a6c5bb46a8 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Apr 2025 15:22:48 -0400 Subject: [PATCH 081/216] seed initial point --- src/aspire/utils/bot_align.py | 12 ++++-------- tests/test_bot_align.py | 11 +++++++---- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/src/aspire/utils/bot_align.py b/src/aspire/utils/bot_align.py index d3df896954..8abb2596f5 100644 --- a/src/aspire/utils/bot_align.py +++ b/src/aspire/utils/bot_align.py @@ -65,10 +65,6 @@ def align_BO( :param seed: Random seed for reproducible results. Integer, defaults None. :return: Rotation matrix R_init (without refinement) or (R_init, R_est) (with refinement). """ - # Generate sequence of seeds for pymanopt initial points - ss = np.random.SeedSequence(seed) - point_seeds = ss.generate_state(max_iters, dtype=np.uint32) - # Avoid utils/operators/utils circular import from aspire.operators import wemd_embed @@ -175,12 +171,12 @@ def euclidean_grad(new, t=t, q=q): ) # If provided, use seed to set initial point for optimizer + initial = None if seed is not None: - with Random(int(point_seeds[t])): + with Random(seed + t): initial = manifold.random_point() - result = optimizer.run(problem, initial_point=initial) - else: - result = optimizer.run(problem) + + result = optimizer.run(problem, initial_point=initial) R_new = result.point.astype(dtype, copy=False) loss[t] = loss_fun(R_new) diff --git a/tests/test_bot_align.py b/tests/test_bot_align.py index a6758828ff..7740bf8e58 100644 --- a/tests/test_bot_align.py +++ b/tests/test_bot_align.py @@ -6,7 +6,7 @@ from numpy.linalg import norm from numpy.random import normal -from aspire.utils import Rotation, align_BO +from aspire.utils import Random, Rotation, align_BO from aspire.volume import Volume @@ -36,6 +36,7 @@ def _angular_dist_degrees(R1, R2): SNRS = [pytest.param(float("inf"), marks=pytest.mark.expensive), 0.5] DTYPES = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] +SEED = 1234 def algo_params_id(params): @@ -74,10 +75,11 @@ def vol_data_fixture(snr, dtype): L = v.resolution shape = (L, L, L) ns_std = np.sqrt(norm(v) ** 2 / (L**3 * snr)).astype(v.dtype) - reference_vol = v + normal(0, ns_std, shape).astype(dtype, copy=False) - r = Rotation.generate_random_rotations(1, dtype=v.dtype, seed=1234) + r = Rotation.generate_random_rotations(1, dtype=v.dtype, seed=SEED) R_true = r.matrices[0] - test_vol = v.rotate(r) + normal(0, ns_std, shape).astype(dtype, copy=False) + with Random(SEED): + reference_vol = v + normal(0, ns_std, shape).astype(dtype, copy=False) + test_vol = v.rotate(r) + normal(0, ns_std, shape).astype(dtype, copy=False) return reference_vol, test_vol, R_true @@ -97,6 +99,7 @@ def test_bot_align(algo_params, vol_data_fixture): loss_type=algo_params[0], downsampled_size=algo_params[1], max_iters=algo_params[2], + seed=SEED, ) # Recovery without refinement (degrees) From a95be862891af16fcd1297e0a364b1c733de037e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 7 Apr 2025 15:29:45 -0400 Subject: [PATCH 082/216] normalize_background legacy mode --- src/aspire/image/image.py | 33 ++++++++++++++++++++----------- src/aspire/source/image.py | 10 ++++++++-- src/aspire/utils/coor_trans.py | 2 -- tests/test_preprocess_pipeline.py | 19 ++++++++++++++++++ 4 files changed, 48 insertions(+), 16 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 27d5ced568..9cdd9f1a06 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -24,7 +24,7 @@ logger = logging.getLogger(__name__) -def normalize_bg(imgs, bg_radius=1.0, do_ramp=True): +def normalize_bg(imgs, bg_radius=1.0, do_ramp=True, legacy=False): """ Normalize backgrounds and apply to a stack of images @@ -33,16 +33,30 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True): :param do_ramp: When it is `True`, fit a ramping background to the data and subtract. Namely perform normalization based on values from each image. Otherwise, a constant background level from all images is used. + :param legacy: Option to match Matlab legacy normalize_background. Default, False, + uses ASPIRE-Python implementation. When True, ramping is disable, a shifted + 2d grid and alternative `bg_radius` is used to generate the background mask, + and standard deviation is computed using N - 1 degrees of freedom. :return: The modified images """ if imgs.ndim > 3: raise NotImplementedError( "`normalize_bg` is currently limited to 1D image stacks." ) - L = imgs.shape[-1] + + # Make adjustments for legacy mode + shifted = False + ddof = 0 # Degrees of freedom for standard deviation + if legacy: + do_ramp = False + shifted = True # Shifts 2d grid by 1/2 pixel for even resolution + bg_radius = 2 * (L // 2) / L + ddof = 1 + + # Generate background mask input_dtype = imgs.dtype - grid = grid_2d(L, indexing="yx", dtype=input_dtype) + grid = grid_2d(L, shifted=shifted, indexing="yx", dtype=input_dtype) mask = grid["r"] > bg_radius if do_ramp: @@ -71,15 +85,10 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True): imgs = imgs.reshape((-1, L, L)) # Apply mask images and calculate mean and std values of background - imgs_masked = imgs * mask - denominator = np.count_nonzero(mask) # scalar int - first_moment = np.sum(imgs_masked, axis=(1, 2)) / denominator - second_moment = np.sum(imgs_masked**2, axis=(1, 2)) / denominator - mean = first_moment.reshape(-1, 1, 1) - variance = second_moment.reshape(-1, 1, 1) - mean**2 - std = np.sqrt(variance) - - return (imgs - mean) / std + mean = np.mean(imgs[:, mask], axis=1) + std = np.std(imgs[:, mask], ddof=ddof, axis=1) + + return (imgs - mean[:, None, None]) / std[:, None, None] def load_mrc(filepath): diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py index f7bbfa4cd8..f529e1534f 100644 --- a/src/aspire/source/image.py +++ b/src/aspire/source/image.py @@ -945,7 +945,7 @@ def invert_contrast(self, batch_size=512): self.generation_pipeline.add_xform(Multiply(scale_factor)) @_as_copy - def normalize_background(self, bg_radius=1.0, do_ramp=True): + def normalize_background(self, bg_radius=1.0, do_ramp=True, legacy=False): """ Normalize the images by the noise background @@ -958,6 +958,10 @@ def normalize_background(self, bg_radius=1.0, do_ramp=True): :param do_ramp: When it is `True`, fit a ramping background to the data and subtract. Namely perform normalization based on values from each image. Otherwise, a constant background level from all images is used. + :param legacy: Option to match Matlab legacy normalize_background. Default, False, + uses ASPIRE-Python implementation. When True, ramping is disable, a shifted + 2d grid and alternative `bg_radius` is used to generate the background mask, + and standard deviation is computed using N - 1 degrees of freedom. :return: On return, the `ImageSource` object has been modified in place. """ @@ -966,7 +970,9 @@ def normalize_background(self, bg_radius=1.0, do_ramp=True): f"size of {bg_radius} and do_ramp of {do_ramp}" ) self.generation_pipeline.add_xform( - LambdaXform(normalize_bg, bg_radius=bg_radius, do_ramp=do_ramp) + LambdaXform( + normalize_bg, bg_radius=bg_radius, do_ramp=do_ramp, legacy=legacy + ) ) def im_backward(self, im, start, weights=None, symmetry_group=None): diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index ce5e5e8a49..80719d497c 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -75,8 +75,6 @@ def _mgrid_slice(n, shifted, normalized): if normalized: # Compute the denominator for normalization denom = n / 2 - if shifted and n % 2 == 0: - denom -= 1 / 2 # Apply the normalization start /= denom diff --git a/tests/test_preprocess_pipeline.py b/tests/test_preprocess_pipeline.py index e53322a266..2717dcdfda 100644 --- a/tests/test_preprocess_pipeline.py +++ b/tests/test_preprocess_pipeline.py @@ -92,6 +92,25 @@ def testNormBackground(L, dtype): assert dtype == imgs_nb.dtype +@pytest.mark.parametrize("L, dtype", params) +def test_norm_background_legacy(L, dtype): + sim = get_sim_object(L, dtype) + bg_radius = 2 * (L // 2) / L + grid = grid_2d(sim.L, shifted=True, indexing="yx", dtype=dtype) + mask = grid["r"] > bg_radius + sim = sim.normalize_background(legacy=True) + imgs_nb = sim.images[:num_images].asnumpy() + new_mean = np.mean(imgs_nb[:, mask]) + new_variance = np.var(imgs_nb[:, mask], ddof=1) + + # new mean of noise should be close to zero and variance should be close to 1 + assert new_mean < utest_tolerance(dtype) + assert abs(new_variance - 1) < 2e-3 + + # dtype of returned images should be the same + assert dtype == imgs_nb.dtype + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def testWhiten(dtype): # Note this atol holds only for L even. Odd tested in testWhiten2. From 1a02baa712c773630bb69574982706eab1a2ff46 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 9 Apr 2025 09:58:03 -0400 Subject: [PATCH 083/216] custom grid for legacy normalize_bg --- src/aspire/image/image.py | 12 +++++++++++- src/aspire/utils/coor_trans.py | 2 ++ tests/test_grids.py | 6 +++--- 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 9cdd9f1a06..4b44be7836 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -56,7 +56,17 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True, legacy=False): # Generate background mask input_dtype = imgs.dtype - grid = grid_2d(L, shifted=shifted, indexing="yx", dtype=input_dtype) + if legacy and L % 2 == 0: + # ASPIRE-Python grid convention for even/shifted/normalized grids + # differs from legacy grid. Rolling custom grid for this case. + start = (-L // 2 + 1 / 2) / (L / 2) + end = (L // 2 - 1 / 2) / (L / 2) + grid_slice = slice(start, end, L * 1j) + y, x = np.mgrid[grid_slice, grid_slice].astype(input_dtype) + phi, r = np.arctan2(y, x), np.hypot(x, y) + grid = {"x": x, "y": y, "phi": phi, "r": r} + else: + grid = grid_2d(L, shifted=shifted, indexing="yx", dtype=input_dtype) mask = grid["r"] > bg_radius if do_ramp: diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index 80719d497c..ce5e5e8a49 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -75,6 +75,8 @@ def _mgrid_slice(n, shifted, normalized): if normalized: # Compute the denominator for normalization denom = n / 2 + if shifted and n % 2 == 0: + denom -= 1 / 2 # Apply the normalization start /= denom diff --git a/tests/test_grids.py b/tests/test_grids.py index c2a4b657c3..bfdffb2cc0 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -1413,7 +1413,7 @@ def test_grid_2d(self): a = self.legacy_references_2d[k][j] b = grid_2d(*k, indexing="xy")[j] if j != "phi": - assert np.allclose(a, b, atol=utest_tolerance(np.float32)) + np.testing.assert_allclose(a, b, atol=utest_tolerance(np.float32)) def test_grid_3d(self): for k in product(self.ns, self.shifts, self.norms): @@ -1425,6 +1425,6 @@ def test_grid_3d(self): continue elif j == "theta": a, b = np.mod(a, np.pi), np.mod(b, np.pi) - assert np.allclose(a, b, atol=utest_tolerance(np.float32)) + np.testing.assert_allclose(a, b, atol=utest_tolerance(np.float32)) else: - assert np.allclose(a, b, atol=utest_tolerance(np.float32)) + np.testing.assert_allclose(a, b, atol=utest_tolerance(np.float32)) From 81b941de61773bf1cc87bbf15c21dc1656dd545e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 9 Apr 2025 10:19:58 -0400 Subject: [PATCH 084/216] Add custom grid to test --- tests/test_preprocess_pipeline.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/test_preprocess_pipeline.py b/tests/test_preprocess_pipeline.py index 2717dcdfda..baa1ecab3a 100644 --- a/tests/test_preprocess_pipeline.py +++ b/tests/test_preprocess_pipeline.py @@ -96,7 +96,18 @@ def testNormBackground(L, dtype): def test_norm_background_legacy(L, dtype): sim = get_sim_object(L, dtype) bg_radius = 2 * (L // 2) / L - grid = grid_2d(sim.L, shifted=True, indexing="yx", dtype=dtype) + + # ASPIRE-Python grid convention for even/shifted/normalized grids + # differs from legacy grid. Rolling custom grid for this case. + if L % 2 == 0: + start = (-L // 2 + 1 / 2) / (L / 2) + end = (L // 2 - 1 / 2) / (L / 2) + grid_slice = slice(start, end, L * 1j) + y, x = np.mgrid[grid_slice, grid_slice].astype(dtype) + phi, r = np.arctan2(y, x), np.hypot(x, y) + grid = {"x": x, "y": y, "phi": phi, "r": r} + else: + grid = grid_2d(sim.L, shifted=True, indexing="yx", dtype=dtype) mask = grid["r"] > bg_radius sim = sim.normalize_background(legacy=True) imgs_nb = sim.images[:num_images].asnumpy() From ab9998a69fcc5d0cd98e71dacce5a35eac91035c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 25 Apr 2025 10:09:40 -0400 Subject: [PATCH 085/216] remove normalization adjustment for even/shifted/normalized grids. Adapt hard-coded reference values in test. --- src/aspire/utils/coor_trans.py | 10 +-- tests/test_grids.py | 154 +++++++++++++++++---------------- 2 files changed, 80 insertions(+), 84 deletions(-) diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index ce5e5e8a49..c7c4ac81a5 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -73,14 +73,8 @@ def _mgrid_slice(n, shifted, normalized): end -= 1 if normalized: - # Compute the denominator for normalization - denom = n / 2 - if shifted and n % 2 == 0: - denom -= 1 / 2 - - # Apply the normalization - start /= denom - end /= denom + start /= n / 2 + end /= n / 2 return slice(start, end, num_points) diff --git a/tests/test_grids.py b/tests/test_grids.py index bfdffb2cc0..1b156d1179 100644 --- a/tests/test_grids.py +++ b/tests/test_grids.py @@ -16,6 +16,8 @@ def setUp(self): self.shifts = [False, True] self.norms = [False, True] + # Note: Hard-coded references for (4, True, True) changed to reflect alteration + # in normalization factor for even/shifted/normalized grids (see PR #1258). self.legacy_references_2d = { (3, False, False): { "x": np.array( @@ -254,19 +256,19 @@ def setUp(self): (4, True, True): { "x": np.array( [ - [-1.0, -1.0, -1.0, -1.0], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [1.0, 1.0, 1.0, 1.0], + [-0.75, -0.75, -0.75, -0.75], + [-0.25, -0.25, -0.25, -0.25], + [0.25, 0.25, 0.25, 0.25], + [0.75, 0.75, 0.75, 0.75], ], dtype=np.float32, ), "y": np.array( [ - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], ], dtype=np.float32, ), @@ -281,10 +283,10 @@ def setUp(self): ), "r": np.array( [ - [1.4142135, 1.0540925, 1.0540925, 1.4142135], - [1.0540925, 0.47140452, 0.47140452, 1.0540925], - [1.0540925, 0.47140452, 0.47140452, 1.0540925], - [1.4142135, 1.0540925, 1.0540925, 1.4142135], + [1.0606601, 0.7905694, 0.7905694, 1.0606601], + [0.7905694, 0.35355338, 0.35355338, 0.7905694], + [0.7905694, 0.35355338, 0.35355338, 0.7905694], + [1.0606601, 0.7905694, 0.7905694, 1.0606601], ], dtype=np.float32, ), @@ -1232,28 +1234,28 @@ def setUp(self): "x": np.array( [ [ - [-1.0, -1.0, -1.0, -1.0], - [-1.0, -1.0, -1.0, -1.0], - [-1.0, -1.0, -1.0, -1.0], - [-1.0, -1.0, -1.0, -1.0], + [-0.75, -0.75, -0.75, -0.75], + [-0.75, -0.75, -0.75, -0.75], + [-0.75, -0.75, -0.75, -0.75], + [-0.75, -0.75, -0.75, -0.75], ], [ - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], + [-0.25, -0.25, -0.25, -0.25], + [-0.25, -0.25, -0.25, -0.25], + [-0.25, -0.25, -0.25, -0.25], + [-0.25, -0.25, -0.25, -0.25], ], [ - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], + [0.25, 0.25, 0.25, 0.25], + [0.25, 0.25, 0.25, 0.25], + [0.25, 0.25, 0.25, 0.25], + [0.25, 0.25, 0.25, 0.25], ], [ - [1.0, 1.0, 1.0, 1.0], - [1.0, 1.0, 1.0, 1.0], - [1.0, 1.0, 1.0, 1.0], - [1.0, 1.0, 1.0, 1.0], + [0.75, 0.75, 0.75, 0.75], + [0.75, 0.75, 0.75, 0.75], + [0.75, 0.75, 0.75, 0.75], + [0.75, 0.75, 0.75, 0.75], ], ], dtype=np.float32, @@ -1261,28 +1263,28 @@ def setUp(self): "y": np.array( [ [ - [-1.0, -1.0, -1.0, -1.0], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [1.0, 1.0, 1.0, 1.0], + [-0.75, -0.75, -0.75, -0.75], + [-0.25, -0.25, -0.25, -0.25], + [0.25, 0.25, 0.25, 0.25], + [0.75, 0.75, 0.75, 0.75], ], [ - [-1.0, -1.0, -1.0, -1.0], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [1.0, 1.0, 1.0, 1.0], + [-0.75, -0.75, -0.75, -0.75], + [-0.25, -0.25, -0.25, -0.25], + [0.25, 0.25, 0.25, 0.25], + [0.75, 0.75, 0.75, 0.75], ], [ - [-1.0, -1.0, -1.0, -1.0], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [1.0, 1.0, 1.0, 1.0], + [-0.75, -0.75, -0.75, -0.75], + [-0.25, -0.25, -0.25, -0.25], + [0.25, 0.25, 0.25, 0.25], + [0.75, 0.75, 0.75, 0.75], ], [ - [-1.0, -1.0, -1.0, -1.0], - [-0.33333334, -0.33333334, -0.33333334, -0.33333334], - [0.33333334, 0.33333334, 0.33333334, 0.33333334], - [1.0, 1.0, 1.0, 1.0], + [-0.75, -0.75, -0.75, -0.75], + [-0.25, -0.25, -0.25, -0.25], + [0.25, 0.25, 0.25, 0.25], + [0.75, 0.75, 0.75, 0.75], ], ], dtype=np.float32, @@ -1290,28 +1292,28 @@ def setUp(self): "z": np.array( [ [ - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], ], [ - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], ], [ - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], ], [ - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], - [-1.0, -0.33333334, 0.33333334, 1.0], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], + [-0.75, -0.25, 0.25, 0.75], ], ], dtype=np.float32, @@ -1377,28 +1379,28 @@ def setUp(self): "r": np.array( [ [ - [1.7320508, 1.4529663, 1.4529663, 1.7320508], - [1.4529663, 1.1055416, 1.1055416, 1.4529663], - [1.4529663, 1.1055416, 1.1055416, 1.4529663], - [1.7320508, 1.4529663, 1.4529663, 1.7320508], + [1.299038, 1.0897247, 1.0897247, 1.299038], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], + [1.299038, 1.0897247, 1.0897247, 1.299038], ], [ - [1.4529663, 1.1055416, 1.1055416, 1.4529663], - [1.1055416, 0.57735026, 0.57735026, 1.1055416], - [1.1055416, 0.57735026, 0.57735026, 1.1055416], - [1.4529663, 1.1055416, 1.1055416, 1.4529663], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], + [0.8291562, 0.4330127, 0.4330127, 0.8291562], + [0.8291562, 0.4330127, 0.4330127, 0.8291562], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], ], [ - [1.4529663, 1.1055416, 1.1055416, 1.4529663], - [1.1055416, 0.57735026, 0.57735026, 1.1055416], - [1.1055416, 0.57735026, 0.57735026, 1.1055416], - [1.4529663, 1.1055416, 1.1055416, 1.4529663], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], + [0.8291562, 0.4330127, 0.4330127, 0.8291562], + [0.8291562, 0.4330127, 0.4330127, 0.8291562], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], ], [ - [1.7320508, 1.4529663, 1.4529663, 1.7320508], - [1.4529663, 1.1055416, 1.1055416, 1.4529663], - [1.4529663, 1.1055416, 1.1055416, 1.4529663], - [1.7320508, 1.4529663, 1.4529663, 1.7320508], + [1.299038, 1.0897247, 1.0897247, 1.299038], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], + [1.0897248, 0.8291562, 0.8291562, 1.0897248], + [1.299038, 1.0897247, 1.0897247, 1.299038], ], ], dtype=np.float32, From d3cfbfa0428f695157dcadb2eb460c614ee186ed Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 25 Apr 2025 10:57:28 -0400 Subject: [PATCH 086/216] Remove custom grids from Image.normalize_bg and test. --- src/aspire/image/image.py | 12 +----------- tests/test_preprocess_pipeline.py | 24 ++++++++---------------- 2 files changed, 9 insertions(+), 27 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 4b44be7836..9cdd9f1a06 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -56,17 +56,7 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True, legacy=False): # Generate background mask input_dtype = imgs.dtype - if legacy and L % 2 == 0: - # ASPIRE-Python grid convention for even/shifted/normalized grids - # differs from legacy grid. Rolling custom grid for this case. - start = (-L // 2 + 1 / 2) / (L / 2) - end = (L // 2 - 1 / 2) / (L / 2) - grid_slice = slice(start, end, L * 1j) - y, x = np.mgrid[grid_slice, grid_slice].astype(input_dtype) - phi, r = np.arctan2(y, x), np.hypot(x, y) - grid = {"x": x, "y": y, "phi": phi, "r": r} - else: - grid = grid_2d(L, shifted=shifted, indexing="yx", dtype=input_dtype) + grid = grid_2d(L, shifted=shifted, indexing="yx", dtype=input_dtype) mask = grid["r"] > bg_radius if do_ramp: diff --git a/tests/test_preprocess_pipeline.py b/tests/test_preprocess_pipeline.py index baa1ecab3a..94c5c1dae4 100644 --- a/tests/test_preprocess_pipeline.py +++ b/tests/test_preprocess_pipeline.py @@ -94,32 +94,24 @@ def testNormBackground(L, dtype): @pytest.mark.parametrize("L, dtype", params) def test_norm_background_legacy(L, dtype): + # Legacy normalize_background uses a shifted grid, a different + # mask radius, disabled ramping, and N - 1 degrees of freedom + # when computing standard deviation. sim = get_sim_object(L, dtype) bg_radius = 2 * (L // 2) / L - - # ASPIRE-Python grid convention for even/shifted/normalized grids - # differs from legacy grid. Rolling custom grid for this case. - if L % 2 == 0: - start = (-L // 2 + 1 / 2) / (L / 2) - end = (L // 2 - 1 / 2) / (L / 2) - grid_slice = slice(start, end, L * 1j) - y, x = np.mgrid[grid_slice, grid_slice].astype(dtype) - phi, r = np.arctan2(y, x), np.hypot(x, y) - grid = {"x": x, "y": y, "phi": phi, "r": r} - else: - grid = grid_2d(sim.L, shifted=True, indexing="yx", dtype=dtype) + grid = grid_2d(sim.L, shifted=True, indexing="yx", dtype=dtype) mask = grid["r"] > bg_radius sim = sim.normalize_background(legacy=True) - imgs_nb = sim.images[:num_images].asnumpy() + imgs_nb = sim.images[:].asnumpy() new_mean = np.mean(imgs_nb[:, mask]) new_variance = np.var(imgs_nb[:, mask], ddof=1) # new mean of noise should be close to zero and variance should be close to 1 - assert new_mean < utest_tolerance(dtype) - assert abs(new_variance - 1) < 2e-3 + np.testing.assert_array_less(new_mean, utest_tolerance(dtype)) + np.testing.assert_array_less(abs(new_variance - 1), 2e-3) # dtype of returned images should be the same - assert dtype == imgs_nb.dtype + np.testing.assert_equal(dtype, imgs_nb.dtype) @pytest.mark.parametrize("dtype", [np.float32, np.float64]) From 672a4e1bd3b0a073ca28382e37a5f0b8180bf706 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 May 2025 09:42:33 -0400 Subject: [PATCH 087/216] disable -> disabled --- src/aspire/image/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 9cdd9f1a06..8fcf60fa5f 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -34,7 +34,7 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True, legacy=False): and subtract. Namely perform normalization based on values from each image. Otherwise, a constant background level from all images is used. :param legacy: Option to match Matlab legacy normalize_background. Default, False, - uses ASPIRE-Python implementation. When True, ramping is disable, a shifted + uses ASPIRE-Python implementation. When True, ramping is disabled, a shifted 2d grid and alternative `bg_radius` is used to generate the background mask, and standard deviation is computed using N - 1 degrees of freedom. :return: The modified images From 1437a375e63cfdf78dd74df747eedd95a607adc3 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 6 May 2025 15:00:44 -0400 Subject: [PATCH 088/216] Add a test reproducing the oriented source save_mode bug. [skip ci] --- tests/test_oriented_source.py | 50 ++++++++++++++++++++++++++++++++++- 1 file changed, 49 insertions(+), 1 deletion(-) diff --git a/tests/test_oriented_source.py b/tests/test_oriented_source.py index bdc0b3f01e..ebabb818d1 100644 --- a/tests/test_oriented_source.py +++ b/tests/test_oriented_source.py @@ -1,10 +1,13 @@ +import glob import logging +import os +import tempfile import numpy as np import pytest from aspire.abinitio import CLSymmetryC3C4, CLSymmetryCn, CLSync3N, CLSyncVoting -from aspire.source import OrientedSource, Simulation +from aspire.source import OrientedSource, RelionSource, Simulation from aspire.volume import CnSymmetricVolume logger = logging.getLogger(__name__) @@ -16,6 +19,7 @@ (CLSymmetryC3C4, "C4"), pytest.param((CLSymmetryCn, "C6"), marks=pytest.mark.expensive), ] +SAVE_MODES = [None, "single"] def src_fixture_id(params): @@ -46,6 +50,11 @@ def src_fixture(request): return og_src, oriented_src +@pytest.fixture(params=SAVE_MODES, ids=lambda x: f"save_mode={x}") +def save_mode(request): + return request.param + + def test_repr(src_fixture): og_src, oriented_src = src_fixture @@ -122,3 +131,42 @@ def test_lazy_evaluation(src_fixture, caplog): assert msg not in caplog.text _ = oriented_src.rotations assert msg in caplog.text + + +def test_save(src_fixture, save_mode): + """ + Test save function and save_mode. + """ + + src = src_fixture[1] + + # Make a fresh tmp_dir + with tempfile.TemporaryDirectory() as tmp_dir: + # Construct file path + fn = os.path.join(tmp_dir, f"test_oriented_source-save_mode-{save_mode}.star") + + # Sanity check test configuration + assert src.n % 2 == 0, "Test expects even number of source images" + + # Configure save + # Ensure we have a (small) unit test sized `batch_size` + batch_size = src.n // 2 + # Set expected output file counts + if save_mode == "single": + num_batches = 1 + else: + num_batches = src.n // batch_size + + # Save + src.save(fn, save_mode=save_mode, batch_size=batch_size) + + # Load from saved `fn` + reloaded_src = RelionSource(fn) + + # Assert reloaded image data identical + np.testing.assert_allclose(reloaded_src.images[:], src.images[:]) + + # Assert the correct number of mrcs files were created + mrc_filelist = glob.glob(os.path.join(tmp_dir, "*.mrcs")) + num_mrc_files = len(mrc_filelist) + assert num_mrc_files == num_batches, "Incorrect number of mrcs files" From 72aed20532db3603f4ea1001bed1fe084a704be2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 6 May 2025 15:55:43 -0400 Subject: [PATCH 089/216] fix save_mode bug --- src/aspire/source/image.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py index f529e1534f..d084deb28d 100644 --- a/src/aspire/source/image.py +++ b/src/aspire/source/image.py @@ -1695,7 +1695,7 @@ def offsets(self, values): def save_metadata(self, starfile_filepath, batch_size=512, save_mode=None): self._orient() return super().save_metadata( - starfile_filepath, batch_size=batch_size, save_mode=None + starfile_filepath, batch_size=batch_size, save_mode=save_mode ) def get_metadata( From fee9bf4991b8c2d322eb90e1a8dc80f9eabf2e90 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 8 May 2025 11:00:47 -0400 Subject: [PATCH 090/216] Remove Numpy2 warning --- src/aspire/__init__.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py index f74af4dd86..f332f8a6e0 100644 --- a/src/aspire/__init__.py +++ b/src/aspire/__init__.py @@ -87,16 +87,3 @@ def __getattr__(attr): return importlib.import_module(f"aspire.{attr}") else: raise AttributeError(f"module `{__name__}` has no attribute `{attr}`.") - - -if parse_version(np.version.version) >= parse_version("2.0.0"): - # ImportWarnings are generally filtered, but raise this one for now. - with warnings.catch_warnings(): - warnings.simplefilter("default") - warnings.warn( - "ASPIRE's Numpy 2 support is in beta. If you experience a runtime" - " crash relating to mismatched dtypes or a Numpy call please try" - ' `pip install "numpy<2"` and report to ASPIRE developers.', - ImportWarning, - stacklevel=1, - ) From 0c1dab9cf5ba39e31eef1a20ea2b4805547550b7 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 May 2025 13:28:22 -0400 Subject: [PATCH 091/216] patch volume rotate dtype, issue #1274 --- src/aspire/volume/volume.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py index 8339d8c28a..6988cd563b 100644 --- a/src/aspire/volume/volume.py +++ b/src/aspire/volume/volume.py @@ -607,9 +607,13 @@ def rotate(self, rot_matrices, zero_nyquist=True): f"{self.__class__.__name__}" f" rot_matrices.dtype {rot_matrices.dtype}" f" != self.dtype {self.dtype}." - " In the future this will raise an error." + f" rot_matrices will be cast to {self.dtype}." ) + # Enforce `rots_inverted.dtype` is always `vol.dtype` + # On mismatch the above warning should have been emitted. + rots_inverted = rots_inverted.astype(self.dtype, copy=False) + # If K = 1 we broadcast the single Rotation object across each volume. if K == 1: pts_rot = rotated_grids_3d(self.resolution, rots_inverted) From 8eb1a931332ecc159366527835b99eef1ab73b76 Mon Sep 17 00:00:00 2001 From: Itamar Tzafrir Date: Fri, 9 May 2025 12:20:25 +0300 Subject: [PATCH 092/216] added I symmetry and tests for symmetry_groups --- src/aspire/volume/symmetry_groups.py | 36 ++++++++++++++++++++++++++++ tests/test_symmetry_groups.py | 22 ++++++++++++++++- 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py index e239144150..6bbfc5439b 100644 --- a/src/aspire/volume/symmetry_groups.py +++ b/src/aspire/volume/symmetry_groups.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod, abstractproperty import numpy as np +from scipy.spatial.transform import Rotation as R from aspire.utils import Rotation @@ -79,6 +80,7 @@ def parse(symmetry): "D": DnSymmetryGroup, "T": TSymmetryGroup, "O": OSymmetryGroup, + "I": ISymmetryGroup, } if symmetry_type not in map_to_sym_group.keys(): raise ValueError( @@ -303,3 +305,37 @@ def generate_rotations(self): # Return rotations. return Rotation.from_rotvec(rot_vecs, dtype=self.dtype) + + +class ISymmetryGroup(SymmetryGroup): + """ + Icosahedral symmetry group. + """ + + def __init__(self): + """ + `ISymmetryGroup` instance that serves up a `Rotation` object + containing rotation matrices of the symmetry group (including the + Identity) accessed via the `matrices` attribute. Note, this is the + chiral icosahedral symmetry group which does not contain reflections. + """ + super().__init__() + + self._symmetry_group = self.generate_rotations() + + @property + def to_string(self): + return "I" + + def generate_rotations(self): + """ + Icosahedral rotation group I (60 elements): + - 1 identity rotation (order 1) + - 24 rotations of order 5: ±72°, ±144° around 6 face-center axes + - 20 rotations of order 3: ±120° around 10 vertex axes + - 15 rotations of order 2: 180° around 15 edge-midpoint axes + + :return: Rotation object containing the icosahedral symmetry group and the identity. + """ + scipy_rotations = R.create_group("I").as_matrix() + return Rotation(scipy_rotations) diff --git a/tests/test_symmetry_groups.py b/tests/test_symmetry_groups.py index 2be56cbf31..d7f49815b6 100644 --- a/tests/test_symmetry_groups.py +++ b/tests/test_symmetry_groups.py @@ -5,10 +5,11 @@ import pytest from aspire.utils import Rotation -from aspire.volume import ( +from aspire.volume.symmetry_groups import ( CnSymmetryGroup, DnSymmetryGroup, IdentitySymmetryGroup, + ISymmetryGroup, OSymmetryGroup, SymmetryGroup, TSymmetryGroup, @@ -23,6 +24,7 @@ GROUPS_WITHOUT_ORDER = [ (TSymmetryGroup,), (OSymmetryGroup,), + (ISymmetryGroup,), ] ORDERS = [2, 3, 4, 5] PARAMS_ORDER = list(itertools.product(GROUPS_WITH_ORDER, ORDERS)) @@ -96,3 +98,21 @@ def test_parser_error(): ValueError, match=f"Symmetry type {junk_symmetry[0]} not supported.*" ): _ = SymmetryGroup.parse(junk_symmetry) + + +def test_group_order(group_fixture): + """Check the number of elements (order) in each symmetry group.""" + if type(group_fixture) in GROUPS_WITH_ORDER: + expected_orders = { + CnSymmetryGroup: 1 * group_fixture.order, + DnSymmetryGroup: 2 * group_fixture.order, + } + else: + expected_orders = { + IdentitySymmetryGroup: 1, + TSymmetryGroup: 12, + OSymmetryGroup: 24, + ISymmetryGroup: 60, + } + expected_order = expected_orders[type(group_fixture)] + assert len(group_fixture.matrices) == expected_order From 90da8521c2c52c2c07689c2ed63354b332c5dbe6 Mon Sep 17 00:00:00 2001 From: Itamar Tzafrir Date: Sun, 11 May 2025 10:52:08 +0300 Subject: [PATCH 093/216] added ISymmetryGroup to volume __init__.py --- src/aspire/volume/__init__.py | 1 + tests/test_symmetry_groups.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/aspire/volume/__init__.py b/src/aspire/volume/__init__.py index 595a07170a..ab30fa45fd 100644 --- a/src/aspire/volume/__init__.py +++ b/src/aspire/volume/__init__.py @@ -2,6 +2,7 @@ CnSymmetryGroup, DnSymmetryGroup, IdentitySymmetryGroup, + ISymmetryGroup, OSymmetryGroup, SymmetryGroup, TSymmetryGroup, diff --git a/tests/test_symmetry_groups.py b/tests/test_symmetry_groups.py index d7f49815b6..441a7ee2ec 100644 --- a/tests/test_symmetry_groups.py +++ b/tests/test_symmetry_groups.py @@ -5,7 +5,7 @@ import pytest from aspire.utils import Rotation -from aspire.volume.symmetry_groups import ( +from aspire.volume import ( CnSymmetryGroup, DnSymmetryGroup, IdentitySymmetryGroup, From a59d76901908f11402b7fb1975720032bdc5c4bc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 May 2025 10:58:03 -0400 Subject: [PATCH 094/216] Get pixel_size from STAR file for RelionSource --- src/aspire/source/relion.py | 17 +++++++++++-- src/aspire/utils/relion_interop.py | 1 + tests/test_relion_source.py | 38 ++++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/aspire/source/relion.py b/src/aspire/source/relion.py index bd6d660dd3..77a9e19d08 100644 --- a/src/aspire/source/relion.py +++ b/src/aspire/source/relion.py @@ -28,7 +28,7 @@ def __init__( self, filepath, data_folder=None, - pixel_size=1, + pixel_size=None, B=0, n_workers=-1, max_rows=None, @@ -42,7 +42,8 @@ def __init__( :param filepath: Absolute or relative path to STAR file :param data_folder: Path to folder w.r.t which all relative paths to .mrcs files are resolved. If None, the folder corresponding to filepath is used. - :param pixel_size: the pixel size of the images in angstroms (Default 1) + :param pixel_size: The pixel size of the images in angstroms. By default, pixel size is + populated from the STAR file if relevant metadata fields exist, otherwise set to 1. :param B: the envelope decay of the CTF in inverse square angstrom (Default 0) :param n_workers: Number of threads to spawn to read referenced .mrcs files (Default -1 to auto detect) :param max_rows: Maximum number of rows in STAR file to read. If None, all rows are read. @@ -114,6 +115,18 @@ def __init__( pixel_size=pixel_size, ) + # Populate pixel_size with metadata if possible. + if pixel_size is None: + if self.has_metadata(["_rlnImagePixelSize"]): + pixel_size = self.get_metadata(["_rlnImagePixelSize"])[0] + elif self.has_metadata(["_rlnDetectorPixelSize", "_rlnMagnification"]): + detector_pixel_size = self.get_metadata(["_rlnDetectorPixelSize"])[0] + magnification = self.get_metadata(["_rlnMagnification"])[0] + pixel_size = 10000 * detector_pixel_size / magnification + else: + pixel_size = 1.0 + self.pixel_size = float(pixel_size) + # CTF estimation parameters coming from Relion CTF_params = [ "_rlnVoltage", diff --git a/src/aspire/utils/relion_interop.py b/src/aspire/utils/relion_interop.py index 9240e47298..a359a48448 100644 --- a/src/aspire/utils/relion_interop.py +++ b/src/aspire/utils/relion_interop.py @@ -20,6 +20,7 @@ "_rlnDetectorPixelSize": float, "_rlnCtfFigureOfMerit": float, "_rlnMagnification": float, + "_rlnImagePixelSize": float, "_rlnAmplitudeContrast": float, "_rlnImageName": str, "_rlnOriginalName": str, diff --git a/tests/test_relion_source.py b/tests/test_relion_source.py index 6ad6c0721d..7746b39570 100644 --- a/tests/test_relion_source.py +++ b/tests/test_relion_source.py @@ -1,6 +1,7 @@ import logging import os +import numpy as np import pytest from aspire.source import RelionSource @@ -56,3 +57,40 @@ def test_symmetry_group(caplog): assert isinstance(src_override_sym.symmetry_group, SymmetryGroup) assert str(src_override_sym.symmetry_group) == "C6" + + +def test_pixel_size(caplog): + """ + Instantiate RelionSource from starfiles containing the following pixel size + field variations: + - "_rlnImagePixelSize" + - "_rlnDetectorPixelSize" and "_rlnMagnification" + - User provided pixel size + - No pixel size provided + and check src.pixel_size is correct. + """ + starfile_im_pix_size = os.path.join(DATA_DIR, "sample_particles_relion31.star") + starfile_detector_pix_size = os.path.join( + DATA_DIR, "sample_particles_relion30.star" + ) + starfile_no_pix_size = os.path.join(DATA_DIR, "rln_proj_64.star") + + # Check pixel size from _rlnImagePixelSize, set to 1.4000 in starfile. + src = RelionSource(starfile_im_pix_size) + np.testing.assert_equal(src.pixel_size, 1.4) + + # Check pixel size calculated from _rlnDetectorPixelSize and _rlnMagnification + src = RelionSource(starfile_detector_pix_size) + det_pix_size = src.get_metadata(["_rlnDetectorPixelSize"])[0] + mag = src.get_metadata("_rlnMagnification")[0] + pix_size = 10000 * det_pix_size / mag + np.testing.assert_equal(src.pixel_size, pix_size) + + # Check user provided pixel size + pix_size = 1.234 + src = RelionSource(starfile_no_pix_size, pixel_size=pix_size) + np.testing.assert_equal(src.pixel_size, pix_size) + + # Check pixel size defaults to 1 if not provided. + src = RelionSource(starfile_no_pix_size) + np.testing.assert_equal(src.pixel_size, 1.0) From 54463510997ec345a3bc12f9c3bfd0ccabb8b98f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 7 May 2025 09:42:48 -0400 Subject: [PATCH 095/216] use file with pixel size metadata for user provided pixel size --- tests/test_relion_source.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_relion_source.py b/tests/test_relion_source.py index 7746b39570..ee488bf1cb 100644 --- a/tests/test_relion_source.py +++ b/tests/test_relion_source.py @@ -79,16 +79,16 @@ def test_pixel_size(caplog): src = RelionSource(starfile_im_pix_size) np.testing.assert_equal(src.pixel_size, 1.4) - # Check pixel size calculated from _rlnDetectorPixelSize and _rlnMagnification + # Check pixel size calculated from _rlnDetectorPixelSize and _rlnMagnification. src = RelionSource(starfile_detector_pix_size) det_pix_size = src.get_metadata(["_rlnDetectorPixelSize"])[0] mag = src.get_metadata("_rlnMagnification")[0] pix_size = 10000 * det_pix_size / mag np.testing.assert_equal(src.pixel_size, pix_size) - # Check user provided pixel size + # Check user provided pixel size. pix_size = 1.234 - src = RelionSource(starfile_no_pix_size, pixel_size=pix_size) + src = RelionSource(starfile_im_pix_size, pixel_size=pix_size) np.testing.assert_equal(src.pixel_size, pix_size) # Check pixel size defaults to 1 if not provided. From d016d2e585bacbde53ec0a5f114240fbe7eb8242 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 7 May 2025 11:08:51 -0400 Subject: [PATCH 096/216] remove unused caplog --- tests/test_relion_source.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_relion_source.py b/tests/test_relion_source.py index ee488bf1cb..756cd3dfbf 100644 --- a/tests/test_relion_source.py +++ b/tests/test_relion_source.py @@ -59,7 +59,7 @@ def test_symmetry_group(caplog): assert str(src_override_sym.symmetry_group) == "C6" -def test_pixel_size(caplog): +def test_pixel_size(): """ Instantiate RelionSource from starfiles containing the following pixel size field variations: From 845a75029ff52bf99dd8f683e845066feeb08565 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 9 May 2025 08:39:10 -0400 Subject: [PATCH 097/216] add pixel size warning --- src/aspire/source/relion.py | 3 +++ tests/test_relion_source.py | 9 ++++++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/aspire/source/relion.py b/src/aspire/source/relion.py index 77a9e19d08..ffab840b17 100644 --- a/src/aspire/source/relion.py +++ b/src/aspire/source/relion.py @@ -124,6 +124,9 @@ def __init__( magnification = self.get_metadata(["_rlnMagnification"])[0] pixel_size = 10000 * detector_pixel_size / magnification else: + logger.warning( + "No pixel size found in STAR file. Defaulting to 1.0 Angstrom" + ) pixel_size = 1.0 self.pixel_size = float(pixel_size) diff --git a/tests/test_relion_source.py b/tests/test_relion_source.py index 756cd3dfbf..2f789999ef 100644 --- a/tests/test_relion_source.py +++ b/tests/test_relion_source.py @@ -59,7 +59,7 @@ def test_symmetry_group(caplog): assert str(src_override_sym.symmetry_group) == "C6" -def test_pixel_size(): +def test_pixel_size(caplog): """ Instantiate RelionSource from starfiles containing the following pixel size field variations: @@ -92,5 +92,8 @@ def test_pixel_size(): np.testing.assert_equal(src.pixel_size, pix_size) # Check pixel size defaults to 1 if not provided. - src = RelionSource(starfile_no_pix_size) - np.testing.assert_equal(src.pixel_size, 1.0) + with caplog.at_level(logging.WARNING): + msg = "No pixel size found in STAR file. Defaulting to 1.0 Angstrom" + src = RelionSource(starfile_no_pix_size) + assert msg in caplog.text + np.testing.assert_equal(src.pixel_size, 1.0) From 3da7c3632add315f3690f8f9e05e4bfae0580222 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 15 May 2025 15:38:20 -0400 Subject: [PATCH 098/216] proceed with download on hash mismatch, with warning. --- src/aspire/downloader/__init__.py | 1 + src/aspire/downloader/data_fetcher.py | 17 ++++++++++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index be0d375878..714bba0e53 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -4,6 +4,7 @@ # isort: on from .data_fetcher import ( available_downloads, + download_all, emdb_2484, emdb_2660, emdb_2824, diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index f655d06ab3..4b81092e6c 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -1,4 +1,5 @@ import shutil +import warnings import numpy as np import pooch @@ -40,7 +41,21 @@ def fetch_data(dataset_name): :return: The absolute path (including the file name) of the file in local storage. """ - return _data_fetcher.fetch(dataset_name) + try: + return _data_fetcher.fetch(dataset_name) + except ValueError as e: + warnings.warn( + f"Hash mismatch for {dataset_name}, proceeding with download. " + "Source file may have been updated." + ) + # force download without hash check + url = _data_fetcher.get_url(dataset_name) + return pooch.retrieve( + url=url, + known_hash=None, + fname=dataset_name, + path=_data_fetcher.path, + ) def download_all(): From 51e9d7fad2ced039964008b4fc212ef19fb77d4d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 15 May 2025 15:45:17 -0400 Subject: [PATCH 099/216] update hash in registry --- src/aspire/downloader/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/downloader/registry.py b/src/aspire/downloader/registry.py index 467ad3c772..d96536d068 100644 --- a/src/aspire/downloader/registry.py +++ b/src/aspire/downloader/registry.py @@ -10,7 +10,7 @@ "emdb_5778.map": "877cbe37b86561c3dfb255aa2308fefcdd8f51f91928b17c2ef5c8dd3afaaef7", "emdb_6287.map": "81463aa6d024c80efcd19aa9b5ac58f3b3464af56e1ef0f104bd25071acc9204", "emdb_2824.map": "7682e1ef6e5bc9f2de9edcf824a03e454ef9cb1ca33bc12920633559f7f826e4", - "emdb_14621.map": "b45774245c2bd5e1a44e801b8fb1705a44d5850631838d060294be42e34a6900", + "emdb_14621.map": "98363ae950229243131025995b5ba0486857ccb1256b3df8d25c1c282155238c", "emdb_2484.map": "6a324e23352bea101c191d5e854026162a5a9b0b8fc73ac5a085cc22038e1999", "emdb_6458.map": "645208af6d36bbd3d172c549e58d387b81142fd320e064bc66105be0eae540d1", "simulated_channelspin.npz": "c0752674acb85417f6a77a28ac55280c1926c73fda9e25ce0a9940728b1dfcc8", From 3bbe7c8a8d49e13636d3ed48bb08187d510507de Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 May 2025 10:57:59 -0400 Subject: [PATCH 100/216] stack_level. unused variable --- src/aspire/downloader/data_fetcher.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index 4b81092e6c..df194c9432 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -34,7 +34,7 @@ def fetch_data(dataset_name): file in local storage doesn’t match the one in the registry, will download a new copy of the file. This is considered a sign that the file was updated in the remote storage. If the hash of the downloaded file still doesn’t match the - one in the registry, will raise an exception to warn of possible file corruption. + one in the registry, will warn user of possible file corruption. :param dataset_name: The file name (as appears in the registry) to fetch from local storage. @@ -43,11 +43,14 @@ def fetch_data(dataset_name): """ try: return _data_fetcher.fetch(dataset_name) - except ValueError as e: + except ValueError: warnings.warn( f"Hash mismatch for {dataset_name}, proceeding with download. " - "Source file may have been updated." + "Source file may have been updated.", + UserWarning, + stacklevel=1, ) + # force download without hash check url = _data_fetcher.get_url(dataset_name) return pooch.retrieve( From 888fd4321da74456bc5986c8de9a68031013323d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 May 2025 14:27:45 -0400 Subject: [PATCH 101/216] add scheduled_workflow and scheduled downloader test. --- .github/workflows/scheduled_workflow.yml | 38 ++++++++++++++++++++++++ tests/test_downloader.py | 9 ++++++ tox.ini | 3 +- 3 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/scheduled_workflow.yml create mode 100644 tests/test_downloader.py diff --git a/.github/workflows/scheduled_workflow.yml b/.github/workflows/scheduled_workflow.yml new file mode 100644 index 0000000000..ac8bac0976 --- /dev/null +++ b/.github/workflows/scheduled_workflow.yml @@ -0,0 +1,38 @@ +name: ASPIRE Python Scheduled Workflow + +on: + workflow_dispatch: # Manual "Run workflow" button + schedule: + - cron: '0 0 * * 0' # Every Sunday at 00:00 UTC + + +jobs: + check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.9' + - name: Install dependencies + run: | + pip install tox + - name: Run Tox Check + run: tox -e check + + scheduled-tests: + needs: check + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + - name: Install Dependencies + run: | + pip install -e ".[dev]" + - name: Scheduled Tests + run: PYTHONWARNINGS=error python -m pytest -m scheduled diff --git a/tests/test_downloader.py b/tests/test_downloader.py new file mode 100644 index 0000000000..cfb32666c8 --- /dev/null +++ b/tests/test_downloader.py @@ -0,0 +1,9 @@ +import pytest + +from aspire.downloader import download_all + + +@pytest.mark.scheduled +def test_download_all(): + """This test will throw a warning if any hashes have changed""" + _ = download_all() diff --git a/tox.ini b/tox.ini index afa6003cc5..3183b3c0ef 100644 --- a/tox.ini +++ b/tox.ini @@ -92,7 +92,8 @@ line_length = 88 testpaths = tests markers = expensive: mark a test as a long running test. -addopts = -m "not expensive" + scheduled: tests that should only run in the scheduled workflow +addopts = -m "not expensive and not scheduled" [gh-actions] python = From 3da38ed3a0de2c6205e5f42093dc5b63e2179664 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 May 2025 15:21:24 -0400 Subject: [PATCH 102/216] same py version --- .github/workflows/scheduled_workflow.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/scheduled_workflow.yml b/.github/workflows/scheduled_workflow.yml index ac8bac0976..0763244184 100644 --- a/.github/workflows/scheduled_workflow.yml +++ b/.github/workflows/scheduled_workflow.yml @@ -14,7 +14,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: '3.9' + python-version: '3.10' - name: Install dependencies run: | pip install tox From fd8aa2c956b6f44ea9fc19269d352fe114743d84 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 20 May 2025 10:31:04 -0400 Subject: [PATCH 103/216] Use logger warning. Update test to fail on hash mismatch warning. --- src/aspire/downloader/data_fetcher.py | 11 ++++++----- tests/test_downloader.py | 11 ++++++++--- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index df194c9432..146eb97448 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -1,5 +1,5 @@ +import logging import shutil -import warnings import numpy as np import pooch @@ -11,6 +11,9 @@ from aspire.utils import Rotation from aspire.volume import Volume +logger = logging.getLogger(__name__) + + # Initialize pooch data fetcher instance. _data_fetcher = pooch.create( # Set the cache path defined in the config. By default, the cache @@ -44,11 +47,9 @@ def fetch_data(dataset_name): try: return _data_fetcher.fetch(dataset_name) except ValueError: - warnings.warn( + logger.warning( f"Hash mismatch for {dataset_name}, proceeding with download. " - "Source file may have been updated.", - UserWarning, - stacklevel=1, + "Source file may have been updated." ) # force download without hash check diff --git a/tests/test_downloader.py b/tests/test_downloader.py index cfb32666c8..68e20b51fa 100644 --- a/tests/test_downloader.py +++ b/tests/test_downloader.py @@ -4,6 +4,11 @@ @pytest.mark.scheduled -def test_download_all(): - """This test will throw a warning if any hashes have changed""" - _ = download_all() +def test_download_all(caplog): + """Fail if a hash mismatch warning is logged during download.""" + caplog.clear() + with caplog.at_level("WARNING"): + _ = download_all() + + if "Hash mismatch" in caplog.text: + pytest.fail(f"Hash mismatch warning was logged.\nCaptured logs:\n{caplog.text}") From c1795ae02e0dcb22f91c53120fc44a1c8c5c1668 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 20 May 2025 10:39:11 -0400 Subject: [PATCH 104/216] Workflow updates: Run on develop, remove check, remove fail on warnings. --- .github/workflows/scheduled_workflow.yml | 20 +++++--------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/.github/workflows/scheduled_workflow.yml b/.github/workflows/scheduled_workflow.yml index 0763244184..e4a1507a9e 100644 --- a/.github/workflows/scheduled_workflow.yml +++ b/.github/workflows/scheduled_workflow.yml @@ -2,25 +2,15 @@ name: ASPIRE Python Scheduled Workflow on: workflow_dispatch: # Manual "Run workflow" button + branches: + - develop schedule: - cron: '0 0 * * 0' # Every Sunday at 00:00 UTC + branches: + - develop jobs: - check: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: '3.10' - - name: Install dependencies - run: | - pip install tox - - name: Run Tox Check - run: tox -e check - scheduled-tests: needs: check runs-on: ubuntu-latest @@ -35,4 +25,4 @@ jobs: run: | pip install -e ".[dev]" - name: Scheduled Tests - run: PYTHONWARNINGS=error python -m pytest -m scheduled + run: pytest -m scheduled From 7674865d9e9676e9083f127716a97ee5cbe07893 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 20 May 2025 10:59:29 -0400 Subject: [PATCH 105/216] Workflow updates: Checkout develop, set cron off-hour, remove needs field. --- .github/workflows/scheduled_workflow.yml | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/.github/workflows/scheduled_workflow.yml b/.github/workflows/scheduled_workflow.yml index e4a1507a9e..d4d82f80df 100644 --- a/.github/workflows/scheduled_workflow.yml +++ b/.github/workflows/scheduled_workflow.yml @@ -5,14 +5,12 @@ on: branches: - develop schedule: - - cron: '0 0 * * 0' # Every Sunday at 00:00 UTC - branches: - - develop + - cron: '15 0 * * 0' # Every Sunday at 00:15 UTC jobs: scheduled-tests: - needs: check + if: github.ref == 'refs/heads/develop' # Ensure only runs on develop runs-on: ubuntu-latest steps: From bae74515900ce8d74c4c923bd54bdca9f1bf95fc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 20 May 2025 14:38:49 -0400 Subject: [PATCH 106/216] test hash mismatch warning works --- tests/test_downloader.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_downloader.py b/tests/test_downloader.py index 68e20b51fa..dfbc135c61 100644 --- a/tests/test_downloader.py +++ b/tests/test_downloader.py @@ -1,6 +1,9 @@ +from pathlib import Path + import pytest from aspire.downloader import download_all +from aspire.downloader.data_fetcher import _data_fetcher, fetch_data @pytest.mark.scheduled @@ -12,3 +15,31 @@ def test_download_all(caplog): if "Hash mismatch" in caplog.text: pytest.fail(f"Hash mismatch warning was logged.\nCaptured logs:\n{caplog.text}") + + +def test_fetch_data_warning(caplog): + """Test that we get expected warning on hash mismatch.""" + # Use the smallest dataset in the registry + dataset_name = "emdb_3645.map" + + # Remove file from cache if it exists + cached_path = Path(_data_fetcher.path) / dataset_name + if cached_path.exists(): + cached_path.unlink() + + # Save original hash from the registry + original_hash = _data_fetcher.registry.get(dataset_name) + assert original_hash is not None + + # Temporarily override the hash to simulate a mismatch + _data_fetcher.registry[dataset_name] = "md5:invalidhash123" + + try: + caplog.clear() + with caplog.at_level("WARNING"): + path = fetch_data(dataset_name) + assert path # Should return the path to the downloaded file + assert f"Hash mismatch for {dataset_name}" in caplog.text + finally: + # Restore original hash + _data_fetcher.registry[dataset_name] = original_hash From ea1960544ebbfda128de313fdf8b5e389499027c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 21 May 2025 08:48:44 -0400 Subject: [PATCH 107/216] checkout on develop. mark test as scheduled. --- .github/workflows/scheduled_workflow.yml | 3 ++- tests/test_downloader.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/scheduled_workflow.yml b/.github/workflows/scheduled_workflow.yml index d4d82f80df..8f96eaa30e 100644 --- a/.github/workflows/scheduled_workflow.yml +++ b/.github/workflows/scheduled_workflow.yml @@ -10,11 +10,12 @@ on: jobs: scheduled-tests: - if: github.ref == 'refs/heads/develop' # Ensure only runs on develop runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 + with: + ref: develop - name: Set up Python uses: actions/setup-python@v5 with: diff --git a/tests/test_downloader.py b/tests/test_downloader.py index dfbc135c61..cddc1cf177 100644 --- a/tests/test_downloader.py +++ b/tests/test_downloader.py @@ -17,6 +17,7 @@ def test_download_all(caplog): pytest.fail(f"Hash mismatch warning was logged.\nCaptured logs:\n{caplog.text}") +@pytest.mark.scheduled def test_fetch_data_warning(caplog): """Test that we get expected warning on hash mismatch.""" # Use the smallest dataset in the registry From 88b21108477743b0f28f973dc43012931d0b2dea Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 21 May 2025 11:17:57 -0400 Subject: [PATCH 108/216] remove manual run-workflow button --- .github/workflows/scheduled_workflow.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/scheduled_workflow.yml b/.github/workflows/scheduled_workflow.yml index 8f96eaa30e..77a789597e 100644 --- a/.github/workflows/scheduled_workflow.yml +++ b/.github/workflows/scheduled_workflow.yml @@ -1,9 +1,6 @@ name: ASPIRE Python Scheduled Workflow on: - workflow_dispatch: # Manual "Run workflow" button - branches: - - develop schedule: - cron: '15 0 * * 0' # Every Sunday at 00:15 UTC From 486bfbd024331c4cde910999aad925c754819ca3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 3 Jun 2025 14:03:22 -0400 Subject: [PATCH 109/216] update windows runner --- .github/workflows/workflow.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml index bf908b655f..4eae25111b 100644 --- a/.github/workflows/workflow.yml +++ b/.github/workflows/workflow.yml @@ -164,7 +164,7 @@ jobs: backend: intel - os: macOS-latest backend: accelerate - - os: windows-2019 + - os: windows-2022 backend: win64 steps: From 865d3a176f1c1808797322e0df0aed9fab88178a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 9 Jun 2025 15:45:22 -0400 Subject: [PATCH 110/216] fix dtype --- src/aspire/covariance/covar2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py index eebfd5a2c7..d2b8745806 100644 --- a/src/aspire/covariance/covar2d.py +++ b/src/aspire/covariance/covar2d.py @@ -596,7 +596,7 @@ def _calc_op(self): M_covar = BlkDiagMatrix.zeros_like(A_mean) for k in np.unique(ctf_idx): - weight = np.count_nonzero(ctf_idx == k) / src.n + weight = (np.count_nonzero(ctf_idx == k) / src.n).astype(self.dtype) ctf_basis_k = ctf_basis[k] ctf_basis_k_t = ctf_basis_k.T From 3f8dbea836d578b8c07bad186568875296fac39d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 9 Jun 2025 15:57:32 -0400 Subject: [PATCH 111/216] one more --- src/aspire/covariance/covar2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py index d2b8745806..9401aa377f 100644 --- a/src/aspire/covariance/covar2d.py +++ b/src/aspire/covariance/covar2d.py @@ -626,7 +626,7 @@ def _mean_correct_covar_rhs(self, b_covar, b_mean, mean_coef): b_covar = b_covar.copy() for k in np.unique(ctf_idx): - weight = np.count_nonzero(ctf_idx == k) / src.n + weight = (np.count_nonzero(ctf_idx == k) / src.n).astype(self.dtype) ctf_basis_k = ctf_basis[k] ctf_basis_k_t = ctf_basis_k.T From 8e7d26da8c793a760b6dee40dbee80532da25203 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 10 Jun 2025 08:51:57 -0400 Subject: [PATCH 112/216] cast as python float --- src/aspire/covariance/covar2d.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py index 9401aa377f..2350a98f9c 100644 --- a/src/aspire/covariance/covar2d.py +++ b/src/aspire/covariance/covar2d.py @@ -596,8 +596,8 @@ def _calc_op(self): M_covar = BlkDiagMatrix.zeros_like(A_mean) for k in np.unique(ctf_idx): - weight = (np.count_nonzero(ctf_idx == k) / src.n).astype(self.dtype) - + weight = float(np.count_nonzero(ctf_idx == k) / src.n) + breakpoint() ctf_basis_k = ctf_basis[k] ctf_basis_k_t = ctf_basis_k.T @@ -626,7 +626,7 @@ def _mean_correct_covar_rhs(self, b_covar, b_mean, mean_coef): b_covar = b_covar.copy() for k in np.unique(ctf_idx): - weight = (np.count_nonzero(ctf_idx == k) / src.n).astype(self.dtype) + weight = float(np.count_nonzero(ctf_idx == k) / src.n) ctf_basis_k = ctf_basis[k] ctf_basis_k_t = ctf_basis_k.T From c2551193f47d9c7e52cb1d13db551496eef7d03a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 10 Jun 2025 09:38:05 -0400 Subject: [PATCH 113/216] oof --- src/aspire/covariance/covar2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py index 2350a98f9c..cb60c63fa6 100644 --- a/src/aspire/covariance/covar2d.py +++ b/src/aspire/covariance/covar2d.py @@ -597,7 +597,7 @@ def _calc_op(self): for k in np.unique(ctf_idx): weight = float(np.count_nonzero(ctf_idx == k) / src.n) - breakpoint() + ctf_basis_k = ctf_basis[k] ctf_basis_k_t = ctf_basis_k.T From 3be1b85f37cdb7188c08c0f8003e93b5f9a699b4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 May 2025 09:51:30 -0400 Subject: [PATCH 114/216] Cn/Dn create_group. Adapt D2 use of DnSymmetryGroup. --- src/aspire/abinitio/commonline_d2.py | 2 +- src/aspire/volume/symmetry_groups.py | 18 ++++-------------- 2 files changed, 5 insertions(+), 15 deletions(-) diff --git a/src/aspire/abinitio/commonline_d2.py b/src/aspire/abinitio/commonline_d2.py index 5d3d2c6b61..6c930105c7 100644 --- a/src/aspire/abinitio/commonline_d2.py +++ b/src/aspire/abinitio/commonline_d2.py @@ -82,7 +82,7 @@ def __init__( # Rearrange in order Identity, about_x, about_y, about_z. # This ordering is necessary for reproducing MATLAB code results. self.gs = DnSymmetryGroup(order=2).matrices.astype(self.dtype, copy=False)[ - [0, 3, 2, 1] + [0, 2, 3, 1] ] def estimate_rotations(self): diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py index 6bbfc5439b..f5ca304e97 100644 --- a/src/aspire/volume/symmetry_groups.py +++ b/src/aspire/volume/symmetry_groups.py @@ -126,8 +126,8 @@ def generate_rotations(self): :return: Rotation object containing the Cn symmetry group and the identity. """ - angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False) - return Rotation.about_axis("z", angles, dtype=self.dtype) + rots = R.create_group("C" + str(self.order)).as_matrix() + return Rotation(rots.astype(self.dtype, copy=False)) class IdentitySymmetryGroup(CnSymmetryGroup): @@ -174,18 +174,8 @@ def generate_rotations(self): :return: Rotation object containing the Dn symmetry group and the identity. """ - - # Rotations to induce cyclic symmetry - angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False) - rot_z = Rotation.about_axis("z", angles, dtype=self.dtype).matrices - - # Perpendicular rotation to induce dihedral symmetry - rot_perp = Rotation.about_axis("y", np.pi, dtype=self.dtype).matrices - - # Full set of rotations. - rots = np.concatenate((rot_z, rot_z @ rot_perp[0])) - - return Rotation(rots) + rots = R.create_group("D" + str(self.order)).as_matrix() + return Rotation(rots.astype(self.dtype, copy=False)) class TSymmetryGroup(SymmetryGroup): From 23cb7705076cf3cae47ea7dbf099bb84c177268e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 May 2025 13:59:27 -0400 Subject: [PATCH 115/216] T/O create group --- src/aspire/volume/symmetry_groups.py | 68 +++------------------------- 1 file changed, 7 insertions(+), 61 deletions(-) diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py index f5ca304e97..dc75df0d38 100644 --- a/src/aspire/volume/symmetry_groups.py +++ b/src/aspire/volume/symmetry_groups.py @@ -201,30 +201,12 @@ def generate_rotations(self): """ A tetrahedron has C3 symmetry along the 4 axes through each vertex and perpendicular to the opposite face, and C2 symmetry along the axes through - the midpoints of opposite edges. We convert from axis-angle representation of - the symmetry group elements into rotation vectors to generate the rotation - matrices via the `from_rotvec()` method. + the midpoints of opposite edges. :return: Rotation object containing the tetrahedral symmetry group and the identity. """ - # C3 rotation vectors, ie. angle * axis. - axes_C3 = np.array( - [[1.0, 1.0, 1.0], [-1.0, -1.0, 1.0], [1.0, -1.0, -1.0], [-1.0, 1.0, -1.0]], - ) - axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis] - angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3]) - rot_vecs_C3 = np.concatenate([angle * axes_C3 for angle in angles_C3]) - - # C2 rotation vectors. - axes_C2 = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) - rot_vecs_C2 = np.pi * axes_C2 - - # The full set of rotation vectors inducing tetrahedral symmetry. - rot_vec_I = np.zeros((1, 3)) - rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C3, rot_vecs_C2)) - - # Return rotations. - return Rotation.from_rotvec(rot_vecs, dtype=self.dtype) + rots = R.create_group("T").as_matrix() + return Rotation(rots.astype(self.dtype, copy=False)) class OSymmetryGroup(SymmetryGroup): @@ -257,44 +239,8 @@ def generate_rotations(self): :return: Rotation object containing the octahedral symmetry group and the identity. """ - - # C4 rotation vectors, ie angle * axis - axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]]) - angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2]) - rot_vecs_C4 = np.array( - [angle * axes_C4 for angle in angles_C4], - ).reshape((9, 3)) - - # C3 rotation vectors - axes_C3 = np.array( - [[1.0, 1.0, 1.0], [-1.0, 1.0, 1.0], [1.0, -1.0, 1.0], [1.0, 1.0, -1.0]] - ) - axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis] - angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3]) - rot_vecs_C3 = np.array( - [angle * axes_C3 for angle in angles_C3], - ).reshape((8, 3)) - - # C2 rotation vectors - axes_C2 = np.array( - [ - [1.0, 1.0, 0.0], - [-1.0, 1.0, 0.0], - [1.0, 0.0, 1.0], - [-1.0, 0.0, 1.0], - [0.0, 1.0, 1.0], - [0.0, -1.0, 1.0], - ], - ) - axes_C2 /= np.linalg.norm(axes_C2, axis=-1)[..., np.newaxis] - rot_vecs_C2 = np.pi * axes_C2 - - # The full set of rotation vectors inducing octahedral symmetry. - rot_vec_I = np.zeros((1, 3)) - rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C4, rot_vecs_C3, rot_vecs_C2)) - - # Return rotations. - return Rotation.from_rotvec(rot_vecs, dtype=self.dtype) + rots = R.create_group("O").as_matrix() + return Rotation(rots.astype(self.dtype, copy=False)) class ISymmetryGroup(SymmetryGroup): @@ -327,5 +273,5 @@ def generate_rotations(self): :return: Rotation object containing the icosahedral symmetry group and the identity. """ - scipy_rotations = R.create_group("I").as_matrix() - return Rotation(scipy_rotations) + rots = R.create_group("I").as_matrix() + return Rotation(rots.astype(self.dtype, copy=False)) From 3c495e96ca63e55c9ca9041ba54765fa3cf3df53 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 May 2025 14:14:03 -0400 Subject: [PATCH 116/216] use to_string --- src/aspire/volume/symmetry_groups.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py index dc75df0d38..9cfb97920d 100644 --- a/src/aspire/volume/symmetry_groups.py +++ b/src/aspire/volume/symmetry_groups.py @@ -126,7 +126,7 @@ def generate_rotations(self): :return: Rotation object containing the Cn symmetry group and the identity. """ - rots = R.create_group("C" + str(self.order)).as_matrix() + rots = R.create_group(self.to_string).as_matrix() return Rotation(rots.astype(self.dtype, copy=False)) @@ -174,7 +174,7 @@ def generate_rotations(self): :return: Rotation object containing the Dn symmetry group and the identity. """ - rots = R.create_group("D" + str(self.order)).as_matrix() + rots = R.create_group(self.to_string).as_matrix() return Rotation(rots.astype(self.dtype, copy=False)) From 87d2274717d11e3e78021878dd664c6628e3b8d8 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 29 May 2025 09:29:16 -0400 Subject: [PATCH 117/216] Remove redundant code --- src/aspire/volume/symmetry_groups.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py index 9cfb97920d..b1c7e6a80d 100644 --- a/src/aspire/volume/symmetry_groups.py +++ b/src/aspire/volume/symmetry_groups.py @@ -1,5 +1,5 @@ import logging -from abc import ABC, abstractmethod, abstractproperty +from abc import ABC, abstractproperty import numpy as np from scipy.spatial.transform import Rotation as R @@ -21,11 +21,12 @@ def __init__(self): self.dtype = np.float64 self.rotations = self.generate_rotations() - @abstractmethod def generate_rotations(self): """ Method for generating a Rotation object for the symmetry group. """ + rots = R.create_group(self.to_string).as_matrix() + return Rotation(rots.astype(self.dtype, copy=False)) def __eq__(self, other): if isinstance(other, self.__class__): @@ -126,8 +127,7 @@ def generate_rotations(self): :return: Rotation object containing the Cn symmetry group and the identity. """ - rots = R.create_group(self.to_string).as_matrix() - return Rotation(rots.astype(self.dtype, copy=False)) + return super().generate_rotations() class IdentitySymmetryGroup(CnSymmetryGroup): @@ -174,8 +174,7 @@ def generate_rotations(self): :return: Rotation object containing the Dn symmetry group and the identity. """ - rots = R.create_group(self.to_string).as_matrix() - return Rotation(rots.astype(self.dtype, copy=False)) + return super().generate_rotations() class TSymmetryGroup(SymmetryGroup): @@ -205,8 +204,7 @@ def generate_rotations(self): :return: Rotation object containing the tetrahedral symmetry group and the identity. """ - rots = R.create_group("T").as_matrix() - return Rotation(rots.astype(self.dtype, copy=False)) + return super().generate_rotations() class OSymmetryGroup(SymmetryGroup): @@ -239,8 +237,7 @@ def generate_rotations(self): :return: Rotation object containing the octahedral symmetry group and the identity. """ - rots = R.create_group("O").as_matrix() - return Rotation(rots.astype(self.dtype, copy=False)) + return super().generate_rotations() class ISymmetryGroup(SymmetryGroup): @@ -273,5 +270,4 @@ def generate_rotations(self): :return: Rotation object containing the icosahedral symmetry group and the identity. """ - rots = R.create_group("I").as_matrix() - return Rotation(rots.astype(self.dtype, copy=False)) + return super().generate_rotations() From 5c1465224401525a87ddd92166bc9be675bfc703 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 29 May 2025 11:36:33 -0400 Subject: [PATCH 118/216] add ISymmetricVolume --- src/aspire/volume/__init__.py | 1 + src/aspire/volume/volume_synthesis.py | 14 ++++++++++++++ tests/test_synthetic_volume.py | 11 +++++++++-- 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/aspire/volume/__init__.py b/src/aspire/volume/__init__.py index ab30fa45fd..eb08beb2aa 100644 --- a/src/aspire/volume/__init__.py +++ b/src/aspire/volume/__init__.py @@ -10,6 +10,7 @@ from .volume import Volume, qr_vols_forward, rotated_grids, rotated_grids_3d from .volume_synthesis import ( # isort:skip + ISymmetricVolume, TSymmetricVolume, OSymmetricVolume, CnSymmetricVolume, diff --git a/src/aspire/volume/volume_synthesis.py b/src/aspire/volume/volume_synthesis.py index d0d9225f48..9a21bce043 100644 --- a/src/aspire/volume/volume_synthesis.py +++ b/src/aspire/volume/volume_synthesis.py @@ -9,6 +9,7 @@ CnSymmetryGroup, DnSymmetryGroup, IdentitySymmetryGroup, + ISymmetryGroup, OSymmetryGroup, TSymmetryGroup, Volume, @@ -253,6 +254,19 @@ def n_blobs(self): return 24 * self.K +class ISymmetricVolume(GaussianBlobsVolume): + """ + A Volume object with Icosahedral symmetry constructed of random 3D Gaussian blobs. + """ + + def _set_symmetry_group(self): + self._symmetry_group = ISymmetryGroup() + + @property + def n_blobs(self): + return 60 * self.K + + class AsymmetricVolume(CnSymmetricVolume): """ An asymmetric Volume constructed of random 3D Gaussian blobs with compact support in the unit sphere. diff --git a/tests/test_synthetic_volume.py b/tests/test_synthetic_volume.py index 019b99b4a5..1c5d91be5a 100644 --- a/tests/test_synthetic_volume.py +++ b/tests/test_synthetic_volume.py @@ -10,6 +10,7 @@ AsymmetricVolume, CnSymmetricVolume, DnSymmetricVolume, + ISymmetricVolume, LegacyVolume, OSymmetricVolume, TSymmetricVolume, @@ -55,6 +56,12 @@ def dtype_fixture(request): (DnSymmetricVolume, 65, 6), ] +# Parameters for icosahedral volumes need higher resolution to give accurate +# results for test_volume)symmetry. +PARAMS_I = [ + pytest.param((ISymmetricVolume, 70), marks=pytest.mark.expensive), + pytest.param((ISymmetricVolume, 71), marks=pytest.mark.expensive), +] # Parameters for tetrahedral, octahedral, asymmetric, and legacy volumes. # These volumes do not have an `order` parameter. @@ -74,7 +81,7 @@ def vol_fixture_id(params): # Create SyntheticVolume fixture for the set of parameters. -@pytest.fixture(params=PARAMS_Cn_Dn + PARAMS, ids=vol_fixture_id) +@pytest.fixture(params=PARAMS_Cn_Dn + PARAMS_I + PARAMS, ids=vol_fixture_id) def vol_fixture(request, dtype_fixture): params = request.param vol_class = params[0] @@ -156,4 +163,4 @@ def test_volume_symmetry(vol_fixture, dtype_fixture): corr = np.dot(rot_vol[0].flatten(), vol[0].flatten()) / np.dot( vol[0].flatten(), vol[0].flatten() ) - assert abs(corr - 1) < 1.1e-5 + assert abs(corr - 1) < 1e-3 From cc685279d985dbcb574e5784fba050db8cb7eb8b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 29 May 2025 11:51:01 -0400 Subject: [PATCH 119/216] typo --- tests/test_synthetic_volume.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_synthetic_volume.py b/tests/test_synthetic_volume.py index 1c5d91be5a..5611672ada 100644 --- a/tests/test_synthetic_volume.py +++ b/tests/test_synthetic_volume.py @@ -57,7 +57,7 @@ def dtype_fixture(request): ] # Parameters for icosahedral volumes need higher resolution to give accurate -# results for test_volume)symmetry. +# results for test_volume_symmetry. PARAMS_I = [ pytest.param((ISymmetricVolume, 70), marks=pytest.mark.expensive), pytest.param((ISymmetricVolume, 71), marks=pytest.mark.expensive), From 9fa3ccc6404e7ca39db7f97639c95cb41d5e2558 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 4 Jun 2025 15:05:22 -0400 Subject: [PATCH 120/216] g_sync symm group order --- tests/test_orient_d2.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_orient_d2.py b/tests/test_orient_d2.py index 968722e905..cc2013d07a 100644 --- a/tests/test_orient_d2.py +++ b/tests/test_orient_d2.py @@ -405,7 +405,7 @@ def g_sync_d2(rots, rots_gt): n_img = len(rots) dtype = rots.dtype - rots_symm = DnSymmetryGroup(2).matrices.astype(dtype, copy=False) + rots_symm = DnSymmetryGroup(2).matrices.astype(dtype, copy=False)[[0, 2, 1, 3]] order = len(rots_symm) A_g = np.zeros((n_img, n_img), dtype=complex) @@ -438,7 +438,6 @@ def g_sync_d2(rots, rots_gt): # Diagonal elements correspond to exp(-i*0) so put 1. # This is important only for verification purposes that spectrum is (K,0,0,0...,0). A_g += np.conj(A_g).T + np.eye(n_img) - _, eig_vecs = np.linalg.eigh(A_g) leading_eig_vec = eig_vecs[:, -1] From 5bfa1b9b61e30b2bb359313ff0cfb9c156cb01df Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 22 May 2025 11:38:23 -0400 Subject: [PATCH 121/216] remove uniform_random_angles --- .../tutorials/weighted_volume_estimation.py | 4 ++-- src/aspire/source/simulation.py | 16 +++++++------- src/aspire/utils/__init__.py | 1 - src/aspire/utils/coor_trans.py | 21 ------------------- tests/test_coor_trans.py | 4 +--- 5 files changed, 10 insertions(+), 36 deletions(-) diff --git a/gallery/tutorials/tutorials/weighted_volume_estimation.py b/gallery/tutorials/tutorials/weighted_volume_estimation.py index a8e7893292..ade958d932 100644 --- a/gallery/tutorials/tutorials/weighted_volume_estimation.py +++ b/gallery/tutorials/tutorials/weighted_volume_estimation.py @@ -100,13 +100,13 @@ # demonstrate that we are in fact generating spectral volumes that # appear reasonably similar to the input volumes. -from aspire.utils import Rotation, uniform_random_angles +from aspire.utils import Rotation reference_v = 0 # Actual volume under comparison spectral_v = 0 # Estimated spectral volume m = 3 # Number of projections -random_rotations = Rotation.from_euler(uniform_random_angles(m, dtype=src.dtype)) +random_rotations = Rotation.generate_random_rotations(m, dtype=src.dtype) # Estimated volume projections estimated_volume[spectral_v].project(random_rotations).show() diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py index a2b9de712c..d4bf53ce01 100644 --- a/src/aspire/source/simulation.py +++ b/src/aspire/source/simulation.py @@ -9,14 +9,7 @@ from aspire.noise import NoiseAdder from aspire.source import ImageSource from aspire.source.image import _ImageAccessor -from aspire.utils import ( - Rotation, - acorr, - ainner, - anorm, - make_symmat, - uniform_random_angles, -) +from aspire.utils import Rotation, acorr, ainner, anorm, make_symmat from aspire.utils.random import randi, randn, random from aspire.volume import AsymmetricVolume, Volume @@ -202,7 +195,12 @@ def __init__( def _init_angles(self, angles): if angles is None: - angles = uniform_random_angles(self.n, seed=self.seed, dtype=self.dtype) + angles = Rotation.generate_random_rotations( + self.n, + seed=self.seed, + dtype=self.dtype, + ).angles + return angles def _populate_ctf_metadata(self, filter_indices): diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index 74e13c8b82..90bfb67cdc 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -11,7 +11,6 @@ grid_3d, register_rotations, rots_to_clmatrix, - uniform_random_angles, ) from .misc import ( # isort:skip diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index c7c4ac81a5..622f9f2c08 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -12,7 +12,6 @@ from aspire import config from aspire.numeric import xp -from aspire.utils.random import Random from aspire.utils.rotation import Rotation logger = logging.getLogger(__name__) @@ -153,26 +152,6 @@ def grid_3d(n, shifted=False, normalized=True, indexing="zyx", dtype=np.float32) return {"x": x, "y": y, "z": z, "phi": phi, "theta": theta, "r": r} -def uniform_random_angles(n, seed=None, dtype=np.float32): - """ - Generate random 3D rotation angles - - :param n: The number of rotation angles to generate - :param seed: Random integer seed to use. If None, the current random state is used. - :return: A n-by-3 ndarray of rotation angles - """ - # Generate random rotation angles, in radians - with Random(seed): - angles = np.column_stack( - ( - np.random.random(n) * 2 * np.pi, - np.arccos(2 * np.random.random(n) - 1), - np.random.random(n) * 2 * np.pi, - ) - ) - return angles.astype(dtype) - - def register_rotations(rots, rots_ref): """ Register estimated orientations to reference ones. diff --git a/tests/test_coor_trans.py b/tests/test_coor_trans.py index 70d4dc37a4..05f9bc3d18 100644 --- a/tests/test_coor_trans.py +++ b/tests/test_coor_trans.py @@ -12,7 +12,6 @@ grid_3d, mean_aligned_angular_distance, register_rotations, - uniform_random_angles, ) DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data") @@ -70,8 +69,7 @@ def testGrid3d(self): ) def testRegisterRots(self): - angles = uniform_random_angles(32, seed=0) - rots_ref = Rotation.from_euler(angles).matrices + rots_ref = Rotation.generate_random_rotations(32, seed=0).matrices q_ang = [[np.pi / 4, np.pi / 4, np.pi / 4]] q_mat = Rotation.from_euler(q_ang).matrices[0] From 9c771b19501c0c508bde0eeba258998d99a6dc3e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 23 May 2025 08:23:11 -0400 Subject: [PATCH 122/216] refactor mean_aligned_angular_distance --- src/aspire/utils/coor_trans.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index 622f9f2c08..99ccb70b72 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -279,9 +279,11 @@ def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): :return: The mean angular distance between registered estimates and the ground truth (in degrees). """ - Q_mat, flag = register_rotations(rots_est, rots_gt) + rots_est = Rotation(rots_est) + rots_gt = Rotation(rots_gt) + Q_mat, flag = rots_est.find_registration(rots_gt) logger.debug(f"Registration Q_mat: {Q_mat}\nflag: {flag}") - regrot = get_aligned_rotations(rots_est, Q_mat, flag) + regrot = rots_est.apply_registration(Q_mat, flag) mean_ang_dist = Rotation.mean_angular_distance(regrot, rots_gt) * 180 / np.pi if degree_tol is not None: From 937610b0e604e7d2fe186642883874b893d5483b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 23 May 2025 09:12:07 -0400 Subject: [PATCH 123/216] remove duplicate rotation registration code --- src/aspire/utils/__init__.py | 3 - src/aspire/utils/coor_trans.py | 115 --------------------------------- tests/test_coor_trans.py | 13 ---- tests/test_orient_sdp.py | 13 +--- 4 files changed, 2 insertions(+), 142 deletions(-) diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index 90bfb67cdc..4a8fb78c67 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -4,12 +4,9 @@ mean_aligned_angular_distance, crop_pad_2d, crop_pad_3d, - get_aligned_rotations, - get_rots_mse, grid_1d, grid_2d, grid_3d, - register_rotations, rots_to_clmatrix, ) diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index 99ccb70b72..2bf317b81a 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -7,8 +7,6 @@ from functools import lru_cache import numpy as np -from numpy.linalg import norm -from scipy.linalg import svd from aspire import config from aspire.numeric import xp @@ -152,119 +150,6 @@ def grid_3d(n, shifted=False, normalized=True, indexing="zyx", dtype=np.float32) return {"x": x, "y": y, "z": z, "phi": phi, "theta": theta, "r": r} -def register_rotations(rots, rots_ref): - """ - Register estimated orientations to reference ones. - - Finds the orthogonal transformation that best aligns the estimated rotations - to the reference rotations. - - :param rots: The rotations to be aligned in the form of a n-by-3-by-3 array. - :param rots_ref: The reference rotations to which we would like to align in - the form of a n-by-3-by-3 array. - :return: o_mat, optimal orthogonal 3x3 matrix to align the two sets; - flag, flag==1 then J conjugacy is required and 0 is not. - """ - - assert ( - rots.shape == rots_ref.shape - ), "Two sets of rotations must have same dimensions." - K = rots.shape[0] - - # Reflection matrix - J = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]]) - - Q1 = np.zeros((3, 3), dtype=rots.dtype) - Q2 = np.zeros((3, 3), dtype=rots.dtype) - - for k in range(K): - R = rots[k, :, :] - Rref = rots_ref[k, :, :] - Q1 = Q1 + R @ Rref.T - Q2 = Q2 + (J @ R @ J) @ Rref.T - - # Compute the two possible orthogonal matrices which register the - # estimated rotations to the true ones. - Q1 = Q1 / K - Q2 = Q2 / K - - # We are registering one set of rotations (the estimated ones) to - # another set of rotations (the true ones). Thus, the transformation - # matrix between the two sets of rotations should be orthogonal. This - # matrix is either Q1 if we recover the non-reflected solution, or Q2, - # if we got the reflected one. In any case, one of them should be - # orthogonal. - - err1 = norm(Q1 @ Q1.T - np.eye(3, dtype=rots.dtype), ord="fro") - err2 = norm(Q2 @ Q2.T - np.eye(3, dtype=rots.dtype), ord="fro") - - # In any case, enforce the registering matrix O to be a rotation. - if err1 < err2: - # Use Q1 as the registering matrix - U, _, V = svd(Q1) - flag = 0 - else: - # Use Q2 as the registering matrix - U, _, V = svd(Q2) - flag = 1 - - Q_mat = U @ V - - return Q_mat, flag - - -def get_aligned_rotations(rots, Q_mat, flag): - """ - Get aligned rotation matrices to reference ones. - - Calculated aligned rotation matrices from the orthogonal transformation - that best aligns the estimated rotations to the reference rotations. - - :param rots: The reference rotations to which we would like to align in - the form of a n-by-3-by-3 array. - :param Q_mat: optimal orthogonal 3x3 transformation matrix - :param flag: flag==1 then J conjugacy is required and 0 is not - :return: regrot, aligned rotation matrices - """ - - K = rots.shape[0] - - # Reflection matrix - J = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]]) - - regrot = np.zeros_like(rots) - for k in range(K): - R = rots[k, :, :] - if flag == 1: - R = J @ R @ J - regrot[k, :, :] = Q_mat.T @ R - - return regrot - - -def get_rots_mse(rots_reg, rots_ref): - """ - Calculate MSE between the estimated orientations to reference ones. - - :param rots_reg: The estimated rotations after alignment in the form of - a n-by-3-by-3 array. - :param rots_ref: The reference rotations. - :return: The MSE value between two sets of rotations. - """ - assert ( - rots_reg.shape == rots_ref.shape - ), "Two sets of rotations must have same dimensions." - K = rots_reg.shape[0] - - diff = np.zeros(K) - mse = 0 - for k in range(K): - diff[k] = norm(rots_reg[k, :, :] - rots_ref[k, :, :], ord="fro") - mse += diff[k] ** 2 - mse = mse / K - return mse - - def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): """ Register estimates to ground truth rotations and compute the diff --git a/tests/test_coor_trans.py b/tests/test_coor_trans.py index 05f9bc3d18..74294fd9bc 100644 --- a/tests/test_coor_trans.py +++ b/tests/test_coor_trans.py @@ -7,11 +7,9 @@ Rotation, crop_pad_2d, crop_pad_3d, - get_aligned_rotations, grid_2d, grid_3d, mean_aligned_angular_distance, - register_rotations, ) DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data") @@ -68,17 +66,6 @@ def testGrid3d(self): ) ) - def testRegisterRots(self): - rots_ref = Rotation.generate_random_rotations(32, seed=0).matrices - - q_ang = [[np.pi / 4, np.pi / 4, np.pi / 4]] - q_mat = Rotation.from_euler(q_ang).matrices[0] - flag = 0 - regrots_ref = get_aligned_rotations(rots_ref, q_mat, flag) - q_mat_est, flag_est = register_rotations(rots_ref, regrots_ref) - - self.assertTrue(np.allclose(flag_est, flag) and np.allclose(q_mat_est, q_mat)) - def testSquareCrop2D(self): # Test even/odd cases based on the convention that the center of a sequence of length n # is (n+1)/2 if n is odd and n/2 + 1 if even. diff --git a/tests/test_orient_sdp.py b/tests/test_orient_sdp.py index 22658ee06a..867ec7b68a 100644 --- a/tests/test_orient_sdp.py +++ b/tests/test_orient_sdp.py @@ -4,13 +4,7 @@ from aspire.abinitio import CommonlineSDP from aspire.nufft import backend_available from aspire.source import Simulation -from aspire.utils import ( - Rotation, - get_aligned_rotations, - mean_aligned_angular_distance, - register_rotations, - rots_to_clmatrix, -) +from aspire.utils import Rotation, mean_aligned_angular_distance, rots_to_clmatrix from aspire.volume import AsymmetricVolume RESOLUTION = [ @@ -189,7 +183,4 @@ def test_deterministic_rounding(src_orient_est_fixture): est_rots = orient_est._deterministic_rounding(gt_gram) # Check that the estimated rotations are close to ground truth after global alignment. - Q_mat, flag = register_rotations(est_rots, gt_rots) - regrot = get_aligned_rotations(est_rots, Q_mat, flag) - - np.testing.assert_allclose(regrot, gt_rots) + mean_aligned_angular_distance(est_rots, gt_rots, degree_tol=1e-5) From b994c6b0555134bfc91cf7584e490e0d17a29e87 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 May 2025 09:35:32 -0400 Subject: [PATCH 124/216] only convert to Rotation if necessary --- src/aspire/utils/coor_trans.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index 2bf317b81a..a0528ea700 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -164,8 +164,11 @@ def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): :return: The mean angular distance between registered estimates and the ground truth (in degrees). """ - rots_est = Rotation(rots_est) - rots_gt = Rotation(rots_gt) + if not isinstance(rots_est, Rotation): + rots_est = Rotation(rots_est) + if not isinstance(rots_gt, Rotation): + rots_gt = Rotation(rots_gt) + Q_mat, flag = rots_est.find_registration(rots_gt) logger.debug(f"Registration Q_mat: {Q_mat}\nflag: {flag}") regrot = rots_est.apply_registration(Q_mat, flag) From 9bfd38873be0eb8794ce4d7e8fd037ec4c011afd Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 May 2025 11:00:51 -0400 Subject: [PATCH 125/216] Remove duplicate common_line_from_rots code --- src/aspire/abinitio/commonline_base.py | 10 ++++------ src/aspire/utils/__init__.py | 1 - src/aspire/utils/coor_trans.py | 25 ------------------------- src/aspire/utils/rotation.py | 2 +- tests/test_rotation.py | 2 +- 5 files changed, 6 insertions(+), 34 deletions(-) diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index fe456c645e..041c78d736 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/src/aspire/abinitio/commonline_base.py @@ -7,7 +7,7 @@ from aspire.image import Image from aspire.operators import PolarFT -from aspire.utils import common_line_from_rots, complex_type, fuzzy_mask, tqdm +from aspire.utils import Rotation, complex_type, fuzzy_mask, tqdm from aspire.utils.random import choice logger = logging.getLogger(__name__) @@ -536,7 +536,7 @@ def _get_shift_equations_approx(self, equations_factor=1, max_memory=4000): n_img = self.n_img # `estimate_shifts()` requires that rotations have already been estimated. - rotations = self.rotations + rotations = Rotation(self.rotations) pf = self.pf.copy() @@ -741,16 +741,14 @@ def _get_cl_indices(self, rotations, i, j, n_theta): """ Get common line indices based on the rotations from i and j images - :param rotations: Array of rotation matrices + :param rotations: Rotation object :param i: Index for i image :param j: Index for j image :param n_theta: Total number of common lines :return: Common line indices for i and j images """ # get the common line indices based on the rotations from i and j images - r_i = rotations[i] - r_j = rotations[j] - c_ij, c_ji = common_line_from_rots(r_i.T, r_j.T, 2 * n_theta) + c_ij, c_ji = rotations.invert().common_lines(i, j, 2 * n_theta) # To match clmatrix, c_ij is always less than PI # and c_ji may be be larger than PI. diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index 4a8fb78c67..b7924f0fb6 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -1,6 +1,5 @@ from .types import complex_type, real_type, utest_tolerance # isort:skip from .coor_trans import ( # isort:skip - common_line_from_rots, mean_aligned_angular_distance, crop_pad_2d, crop_pad_3d, diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index a0528ea700..841105c553 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -180,31 +180,6 @@ def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): return mean_ang_dist -def common_line_from_rots(r1, r2, ell): - """ - Compute the common line induced by rotation matrices r1 and r2. - - :param r1: The first rotation matrix of 3-by-3 array. - :param r2: The second rotation matrix of 3-by-3 array. - :param ell: The total number of common lines. - :return: The common line indices for both first and second rotations. - """ - - assert r1.dtype == r2.dtype, "Ambiguous dtypes" - - ut = np.dot(r2, r1.T) - alpha_ij = np.arctan2(ut[2, 0], -ut[2, 1]) + np.pi - alpha_ji = np.arctan2(-ut[0, 2], ut[1, 2]) + np.pi - - ell_ij = alpha_ij * ell / (2 * np.pi) - ell_ji = alpha_ji * ell / (2 * np.pi) - - ell_ij = int(np.mod(np.round(ell_ij), ell)) - ell_ji = int(np.mod(np.round(ell_ji), ell)) - - return ell_ij, ell_ji - - def rots_to_clmatrix(rots, n_theta): """ Compute the common lines matrix induced by all pairs of rotation diff --git a/src/aspire/utils/rotation.py b/src/aspire/utils/rotation.py index 2a2e0f98a0..de814f31f9 100644 --- a/src/aspire/utils/rotation.py +++ b/src/aspire/utils/rotation.py @@ -233,7 +233,7 @@ def common_lines(self, i, j, ell): r2 = self._matrices[j] ut = np.dot(r2, r1.T) alpha_ij = np.arctan2(ut[2, 0], -ut[2, 1]) + np.pi - alpha_ji = np.arctan2(ut[0, 2], -ut[1, 2]) + np.pi + alpha_ji = np.arctan2(-ut[0, 2], ut[1, 2]) + np.pi ell_ij = alpha_ij * ell / (2 * np.pi) ell_ji = alpha_ji * ell / (2 * np.pi) diff --git a/tests/test_rotation.py b/tests/test_rotation.py index 5ecf92841c..979f185ce1 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -116,7 +116,7 @@ def test_mse(rot_obj): def test_common_lines(rot_obj): ell_ij, ell_ji = rot_obj.common_lines(8, 11, 360) - np.testing.assert_equal([ell_ij, ell_ji], [235, 284]) + np.testing.assert_equal([ell_ij, ell_ji], [235, 104]) def test_string(rot_obj): From 254b7eddb593351cd54564b06d7dd9fdc38a2505 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 3 Jun 2025 08:57:37 -0400 Subject: [PATCH 126/216] update docstring --- src/aspire/utils/coor_trans.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py index 841105c553..600d8355b1 100644 --- a/src/aspire/utils/coor_trans.py +++ b/src/aspire/utils/coor_trans.py @@ -155,8 +155,10 @@ def mean_aligned_angular_distance(rots_est, rots_gt, degree_tol=None): Register estimates to ground truth rotations and compute the mean angular distance between them (in degrees). - :param rots_est: A set of estimated rotations of size nx3x3. - :param rots_gt: A set of ground truth rotations of size nx3x3. + :param rots_est: A set of estimated rotations. A Rotation object or + array of size nx3x3. + :param rots_gt: A set of ground truth rotations. A Rotation object or + array of size nx3x3. :param degree_tol: Option to assert if the mean angular distance is less than `degree_tol` degrees. If `None`, returns the mean aligned angular distance. From 8af3314eb2d5e485d5825e176335673e65e922e5 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 11 Jun 2025 10:06:44 -0400 Subject: [PATCH 127/216] permit corrupt MRC files lots of published datasets are full of junk metadata --- src/aspire/source/coordinates.py | 2 +- src/aspire/source/relion.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index 495f7a87f6..045ae5d8ae 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -77,7 +77,7 @@ def __init__(self, files, particle_size, max_rows, B, symmetry_group): dtype = "float32" # If file is MRC, we can get dtype from header. if self._ext == ".mrc": - with mrcfile.open(first_mrc) as mrc: + with mrcfile.open(first_mrc, mode="r", permissive=True) as mrc: # get dtype from first micrograph mode = int(mrc.header.mode) dtypes = {0: "int8", 1: "int16", 2: "float32", 6: "uint16"} diff --git a/src/aspire/source/relion.py b/src/aspire/source/relion.py index ffab840b17..a70367643b 100644 --- a/src/aspire/source/relion.py +++ b/src/aspire/source/relion.py @@ -72,7 +72,7 @@ def __init__( # Peek into the first image and populate some attributes first_mrc_filepath = metadata["__mrc_filepath"][0] - mrc = mrcfile.open(first_mrc_filepath) + mrc = mrcfile.open(first_mrc_filepath, mode="r", permissive=True) # Get the 'mode' (data type) - TODO: There's probably a more direct way to do this. mode = int(mrc.header.mode) @@ -247,7 +247,7 @@ def _images(self, indices): logger.debug(f"Indices: {indices}") def load_single_mrcs(filepath, indices): - arr = mrcfile.open(filepath).data + arr = mrcfile.open(filepath, mode="r", permissive=True).data # if the stack only contains one image, arr will have shape (resolution, resolution) # the code below reshapes it to (1, resolution, resolution) if len(arr.shape) == 2: From 9ab9d969d82959db2e071b6aff4d47b84cd624c3 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 27 Mar 2025 11:51:11 -0400 Subject: [PATCH 128/216] stub in BFT work from notebook --- src/aspire/classification/__init__.py | 1 + src/aspire/classification/averager2d.py | 191 +++++++++++++++++++++++- 2 files changed, 191 insertions(+), 1 deletion(-) diff --git a/src/aspire/classification/__init__.py b/src/aspire/classification/__init__.py index df92d36932..e62139349f 100644 --- a/src/aspire/classification/__init__.py +++ b/src/aspire/classification/__init__.py @@ -4,6 +4,7 @@ BFRAverager2D, BFSRAverager2D, BFSReddyChatterjiAverager2D, + BFTAverager2D, EMAverager2D, FTKAverager2D, ReddyChatterjiAverager2D, diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index e25187d5a2..8c0f98604d 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -6,7 +6,8 @@ from aspire.basis import Coef from aspire.classification.reddy_chatterji import reddy_chatterji_register from aspire.image import Image, ImageStacker, MeanImageStacker -from aspire.numeric import xp +from aspire.numeric import fft, xp +from aspire.operators import PolarFT from aspire.utils import tqdm, trange from aspire.utils.coor_trans import grid_2d @@ -725,6 +726,194 @@ def average( return AligningAverager2D.average(self, classes, reflections, coefs) +class BFTAverager2D(AligningAverager2D): + """ + This perfoms a Brute Force Translations and fast rotational alignment. + + For each pair of x_shifts and y_shifts, + Perform cross correlation based rotational alignment + + Return the rotation and shift yielding the best results. + """ + + def __init__( + self, + composite_basis, + src, + alignment_basis=None, + n_angles=360, + radius=None, + batch_size=512, + dtype=None, + ): + """ + See AligningAverager2D adds `n_angles` and `radius`. + + :params n_angles: Number of brute force rotations to attempt, defaults 360. + :param radius: Brute force translation search radius. + Defaults to src.L//16. + """ + super().__init__( + composite_basis, + src, + alignment_basis, + batch_size=batch_size, + dtype=dtype, + ) + + self.n_angles = n_angles + + # # XXX Will use polar for rotate + # if not hasattr(self.alignment_basis, "rotate"): + # raise RuntimeError( + # f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `rotate` method." + # ) + + self.radius = radius if radius is not None else src.L // 16 + + if self.radius != 0: + + if not hasattr(self.alignment_basis, "shift"): + raise RuntimeError( + f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method." + ) + + # XXX todo config + ntheta = 360 + nrad = self.src.L + + # Setup Polar Transform + self._pft = PolarFT(self.src.L, ntheta=ntheta, nrad=nrad, dtype=self.dtype) + + def _fast_rotational_alignment(self, A, B): + """ + Perform fast rotational alignment using Polar Fourier cross correlation. + """ + + if not isinstance(A, Image): + A = Image(A) + if not isinstance(B, Image): + B = Image(B) + + pftA = self._pft.half_to_full(self._pft.transform(A)) + pftB = self._pft.half_to_full(self._pft.transform(B)) + + # 2 hats one sum + pftA = fft.fft(pftA, axis=-2) + pftB = fft.fft(pftB, axis=-2) + x = (pftA * pftB.conj())[0] + # x = x /np.linalg.norm(x) # waste of compute, just for diagnostics + circ_corr = abs(fft.ifft2(x)) + angular = np.sum(circ_corr, axis=-1) # sum all radial contributions + + # Resolve the angle maximizing the correlation through the angular dimension + ind = np.argmax(angular) + max_theta_deg = 360 / self._pft.ntheta * ind + max_theta = np.deg2rad(max_theta_deg) + peak = angular[ind] + + return max_theta, peak + + def align(self, classes, reflections, basis_coefficients=None): + """ + See `AligningAverager2D.align` + """ + + # Admit simple case of single case alignment + classes = np.atleast_2d(classes) + reflections = np.atleast_2d(reflections) + + # Result arrays + # These arrays will incrementally store our best alignment. + n_classes, n_nbor = classes.shape + rotations = np.zeros((n_classes, n_nbor), dtype=self.dtype) + dot_products = np.ones((n_classes, n_nbor), dtype=self.dtype) * -np.inf + shifts = np.zeros((*classes.shape, 2), dtype=int) + + # Work arrays + _rotations = np.zeros((n_nbor), dtype=self.dtype) + _dot_products = np.ones((n_nbor), dtype=self.dtype) * -np.inf + + # Create a search grid and force initial pair to (0,0) + # This is done primarily in case of a tie later, we would take unshifted. + x_shifts, y_shifts = self._shift_search_grid( + self.src.L, self.radius, roll_zero=True + ) + + for k in trange(n_classes, desc="Rotationally aligning classes"): + # 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 + # The coefficient for the base images are also computed here. + if basis_coefficients is None: + original_images = Image(self._cls_images(classes[k], src=self.src)) + else: + original_coef = basis_coefficients[classes[k], :] + original_images = self.alignment_basis.evaluate(original_coef) + + template_image = original_images[0] + + # Loop over shift search space, updating best result + for x, y in tqdm( + zip(x_shifts, y_shifts), + total=len(x_shifts), + desc="\tmaximizing over shifts", + disable=len(x_shifts) == 1, + leave=False, + ): + shift = np.array([x, y], dtype=int) + logger.debug(f"Computing rotational alignment after shift ({x},{y}).") + + # For each shift, the set of neighbor images is shifted. + # This order is chosen because: + # i) allows concatenation of shifts and rotation + # operations after orientation estimation + # ii) because generally the number of neighbors << the + # number of test rotations. + + # Note the base original_image[0] should remain unprocessed + _images = original_images[1:] + # Skip zero shifting. + # XXXX we can try inverting this later, shifting base image for less compute + if np.any(shift != 0): + _images = _images.shift(shift) + _images = _images.asnumpy().copy() + + # XXXX I think we might need to do this before shifting?! + # Handle reflections + refl = reflections[k][1:] # skips original_image 0 + _images[refl] = np.flipud(_images[refl]) + + # Compute and assign the best rotation found with this translation + # TODO, vectorize FRA + for i in range(1, n_nbor): + # note offset of 1 for skipped original_image 0 + _rotations[i], _dot_products[i] = self._fast_rotational_alignment( + template_image, _images[i - 1] + ) + + # Test and update + # Each base-neighbor pair may have a best shift+rot from a different shift iteration. + improved_indices = _dot_products > dot_products[k] + rotations[k, improved_indices] = _rotations[improved_indices] + dot_products[k, improved_indices] = _dot_products[improved_indices] + shifts[k, improved_indices] = ( + shift # when conver to shifting original_image, commute_shift_rot(shift, rot) + ) + + if (x, y) == (0, 0): + logger.debug("Initial rotational alignment complete (shift (0,0))") + # skipped original_image 0, f"{np.sum(improved_indices)} =?= {np.size(classes)}" + expected = np.size(classes[0]) - 1 + assert np.sum(improved_indices) == expected + else: + logger.debug( + f"Shift ({x},{y}) complete. Improved {np.sum(improved_indices)} alignments." + ) + + return rotations, shifts, dot_products + + class EMAverager2D(Averager2D): """ Citation needed. From b6a93428f44339af685511413ca667de25eb7ed1 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 27 Mar 2025 19:03:19 -0400 Subject: [PATCH 129/216] fixup mixing with translations --- src/aspire/classification/averager2d.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 8c0f98604d..4645965be7 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -245,9 +245,11 @@ def _shift_search_grid(self, L, radius, roll_zero=False): """ # We'll brute force all shifts in a grid. - g = grid_2d(L, normalized=False) - disc = g["r"] <= radius + sub_pixel = 1 # XXX + g = grid_2d(sub_pixel * L, normalized=False) + disc = g["r"] <= sub_pixel * radius X, Y = g["x"][disc], g["y"][disc] + X, Y = X / sub_pixel, Y / sub_pixel # Optionally roll arrays so 0 is first. if roll_zero: @@ -778,7 +780,7 @@ def __init__( f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method." ) - # XXX todo config + # XXX todo, better config ntheta = 360 nrad = self.src.L @@ -828,7 +830,7 @@ def align(self, classes, reflections, basis_coefficients=None): n_classes, n_nbor = classes.shape rotations = np.zeros((n_classes, n_nbor), dtype=self.dtype) dot_products = np.ones((n_classes, n_nbor), dtype=self.dtype) * -np.inf - shifts = np.zeros((*classes.shape, 2), dtype=int) + shifts = np.zeros((*classes.shape, 2), dtype=self.dtype) # Work arrays _rotations = np.zeros((n_nbor), dtype=self.dtype) @@ -839,6 +841,9 @@ def align(self, classes, reflections, basis_coefficients=None): x_shifts, y_shifts = self._shift_search_grid( self.src.L, self.radius, roll_zero=True ) + print(x_shifts, y_shifts) + + mask = grid_2d(self.src.L, normalized=True)["r"] < 1 for k in trange(n_classes, desc="Rotationally aligning classes"): # We want to locally cache the original images, @@ -851,7 +856,7 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) - template_image = original_images[0] + template_image = original_images[0] * mask # Loop over shift search space, updating best result for x, y in tqdm( @@ -861,8 +866,8 @@ def align(self, classes, reflections, basis_coefficients=None): disable=len(x_shifts) == 1, leave=False, ): - shift = np.array([x, y], dtype=int) - logger.debug(f"Computing rotational alignment after shift ({x},{y}).") + shift = np.array([x, y], dtype=shifts.dtype) + logger.debug(f"Computing rotational alignment after shift {shift}.") # For each shift, the set of neighbor images is shifted. # This order is chosen because: @@ -876,9 +881,12 @@ def align(self, classes, reflections, basis_coefficients=None): # Skip zero shifting. # XXXX we can try inverting this later, shifting base image for less compute if np.any(shift != 0): - _images = _images.shift(shift) + _images = _images.shift(-shift) _images = _images.asnumpy().copy() + # mask + _images = _images * mask + # XXXX I think we might need to do this before shifting?! # Handle reflections refl = reflections[k][1:] # skips original_image 0 From 4083c0e2696b0879a6ac0ccd482037d2ad25eb72 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 27 Mar 2025 20:14:23 -0400 Subject: [PATCH 130/216] vector fast polar align --- src/aspire/classification/averager2d.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 4645965be7..fd872079f4 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -803,18 +803,17 @@ def _fast_rotational_alignment(self, A, B): # 2 hats one sum pftA = fft.fft(pftA, axis=-2) pftB = fft.fft(pftB, axis=-2) - x = (pftA * pftB.conj())[0] - # x = x /np.linalg.norm(x) # waste of compute, just for diagnostics + x = pftA * pftB.conj() circ_corr = abs(fft.ifft2(x)) angular = np.sum(circ_corr, axis=-1) # sum all radial contributions # Resolve the angle maximizing the correlation through the angular dimension - ind = np.argmax(angular) - max_theta_deg = 360 / self._pft.ntheta * ind - max_theta = np.deg2rad(max_theta_deg) - peak = angular[ind] + inds = np.argmax(angular, axis=-1) + max_thetas_deg = 360 / self._pft.ntheta * inds + max_thetas = np.deg2rad(max_thetas_deg) + peaks = np.take_along_axis(angular, inds.reshape(-1, 1), axis=1).flatten() - return max_theta, peak + return max_thetas, peaks def align(self, classes, reflections, basis_coefficients=None): """ @@ -893,12 +892,10 @@ def align(self, classes, reflections, basis_coefficients=None): _images[refl] = np.flipud(_images[refl]) # Compute and assign the best rotation found with this translation - # TODO, vectorize FRA - for i in range(1, n_nbor): - # note offset of 1 for skipped original_image 0 - _rotations[i], _dot_products[i] = self._fast_rotational_alignment( - template_image, _images[i - 1] - ) + # note offset of 1 for skipped original_image 0 + _rotations[1:], _dot_products[1:] = self._fast_rotational_alignment( + template_image, _images[:] + ) # Test and update # Each base-neighbor pair may have a best shift+rot from a different shift iteration. From 76173865a41313b81e182edb3184efb2c0afdae3 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 27 Mar 2025 21:26:27 -0400 Subject: [PATCH 131/216] shift base image and commute shift --- src/aspire/classification/averager2d.py | 53 ++++++++++++++++--------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index fd872079f4..d87b14cac8 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -14,6 +14,17 @@ logger = logging.getLogger(__name__) +def commute_shift_rot(shifts, rots): + """ + Rotate `shifts` points by `rots` ccw radians. + """ + sx = shifts[:, 0] + sy = shifts[:, 1] + x = sx * np.cos(rots) - sy * np.sin(rots) + y = sx * np.sin(rots) + sy * np.cos(rots) + return np.stack((x, y), axis=1) + + class Averager2D(ABC): """ Base class for 2D Image Averaging methods. @@ -235,21 +246,22 @@ def _innerloop(i): return Image(avgs) - def _shift_search_grid(self, L, radius, roll_zero=False): + def _shift_search_grid(self, L, radius, roll_zero=False, sub_pixel=1): """ Returns two 1-D arrays representing the X and Y grid points in the defined shift search space (disc <= self.radius). :param radius: Disc radius in pixels + :param roll_zero: Roll (0,0) to zero'th element. Defaults to False. + :param sub_pixel: Sub pixel decimation. 1 is integer, 0.1 is 1/10 pixel, etc. :returns: Grid points as 2-tuple of vectors X,Y. """ # We'll brute force all shifts in a grid. - sub_pixel = 1 # XXX - g = grid_2d(sub_pixel * L, normalized=False) - disc = g["r"] <= sub_pixel * radius + g = grid_2d(1 / sub_pixel * L, normalized=False) + disc = g["r"] <= 1 / sub_pixel * radius X, Y = g["x"][disc], g["y"][disc] - X, Y = X / sub_pixel, Y / sub_pixel + X, Y = X * sub_pixel, Y * sub_pixel # Optionally roll arrays so 0 is first. if roll_zero: @@ -745,6 +757,7 @@ def __init__( alignment_basis=None, n_angles=360, radius=None, + sub_pixel=0.1, batch_size=512, dtype=None, ): @@ -780,12 +793,15 @@ def __init__( f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method." ) + self.sub_pixel = sub_pixel + # XXX todo, better config ntheta = 360 nrad = self.src.L # Setup Polar Transform self._pft = PolarFT(self.src.L, ntheta=ntheta, nrad=nrad, dtype=self.dtype) + self._mask = grid_2d(self.src.L, normalized=True)["r"] < 1 def _fast_rotational_alignment(self, A, B): """ @@ -838,12 +854,13 @@ def align(self, classes, reflections, basis_coefficients=None): # Create a search grid and force initial pair to (0,0) # This is done primarily in case of a tie later, we would take unshifted. x_shifts, y_shifts = self._shift_search_grid( - self.src.L, self.radius, roll_zero=True + self.src.L, + self.radius, + roll_zero=True, + sub_pixel=self.sub_pixel, ) print(x_shifts, y_shifts) - mask = grid_2d(self.src.L, normalized=True)["r"] < 1 - for k in trange(n_classes, desc="Rotationally aligning classes"): # We want to locally cache the original images, # because we will mutate them with shifts in the next loop. @@ -855,8 +872,6 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) - template_image = original_images[0] * mask - # Loop over shift search space, updating best result for x, y in tqdm( zip(x_shifts, y_shifts), @@ -876,15 +891,15 @@ def align(self, classes, reflections, basis_coefficients=None): # number of test rotations. # Note the base original_image[0] should remain unprocessed - _images = original_images[1:] + _images = original_images[1:].asnumpy().copy() # Skip zero shifting. - # XXXX we can try inverting this later, shifting base image for less compute + template_image = original_images[0] if np.any(shift != 0): - _images = _images.shift(-shift) - _images = _images.asnumpy().copy() + template_image = template_image.shift(shift) # mask - _images = _images * mask + template_image = template_image * self._mask + _images = _images * self._mask # XXXX I think we might need to do this before shifting?! # Handle reflections @@ -902,9 +917,8 @@ def align(self, classes, reflections, basis_coefficients=None): improved_indices = _dot_products > dot_products[k] rotations[k, improved_indices] = _rotations[improved_indices] dot_products[k, improved_indices] = _dot_products[improved_indices] - shifts[k, improved_indices] = ( - shift # when conver to shifting original_image, commute_shift_rot(shift, rot) - ) + # base shifts assigned here, commutation resolved end of loop + shifts[k, improved_indices] = shift if (x, y) == (0, 0): logger.debug("Initial rotational alignment complete (shift (0,0))") @@ -916,6 +930,9 @@ def align(self, classes, reflections, basis_coefficients=None): f"Shift ({x},{y}) complete. Improved {np.sum(improved_indices)} alignments." ) + # Commute the rotation and shift (code shifted the base image instead of all class members) + shifts[k] = commute_shift_rot(shifts[k], -rotations[k]) + return rotations, shifts, dot_products From 4043179c779de3b015498e1fbdbc6acfa2e692b4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 27 Mar 2025 21:29:26 -0400 Subject: [PATCH 132/216] cleanup --- src/aspire/classification/averager2d.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index d87b14cac8..aecf6d1495 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -820,8 +820,7 @@ def _fast_rotational_alignment(self, A, B): pftA = fft.fft(pftA, axis=-2) pftB = fft.fft(pftB, axis=-2) x = pftA * pftB.conj() - circ_corr = abs(fft.ifft2(x)) - angular = np.sum(circ_corr, axis=-1) # sum all radial contributions + angular = np.sum(np.abs(fft.ifft2(x)), axis=-1) # sum all radial contributions # Resolve the angle maximizing the correlation through the angular dimension inds = np.argmax(angular, axis=-1) @@ -852,14 +851,13 @@ def align(self, classes, reflections, basis_coefficients=None): _dot_products = np.ones((n_nbor), dtype=self.dtype) * -np.inf # Create a search grid and force initial pair to (0,0) - # This is done primarily in case of a tie later, we would take unshifted. + # This is done primarily in case of a tie later, we would prefer unshifted. x_shifts, y_shifts = self._shift_search_grid( self.src.L, self.radius, roll_zero=True, sub_pixel=self.sub_pixel, ) - print(x_shifts, y_shifts) for k in trange(n_classes, desc="Rotationally aligning classes"): # We want to locally cache the original images, @@ -901,7 +899,7 @@ def align(self, classes, reflections, basis_coefficients=None): template_image = template_image * self._mask _images = _images * self._mask - # XXXX I think we might need to do this before shifting?! + # XXXX think if we need to do this before shifting?! # Handle reflections refl = reflections[k][1:] # skips original_image 0 _images[refl] = np.flipud(_images[refl]) From 54be4579832e72aec1945edf4ea1a7bd643d3efb Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 27 Mar 2025 22:27:36 -0400 Subject: [PATCH 133/216] hack in gpu code, dirty [skip ci] --- src/aspire/classification/averager2d.py | 26 ++++++++++++------------- src/aspire/operators/polar_ft.py | 24 ++++++++++++++--------- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index aecf6d1495..387acc370c 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -801,17 +801,17 @@ def __init__( # Setup Polar Transform self._pft = PolarFT(self.src.L, ntheta=ntheta, nrad=nrad, dtype=self.dtype) - self._mask = grid_2d(self.src.L, normalized=True)["r"] < 1 + self._mask = xp.asarray(grid_2d(self.src.L, normalized=True)["r"] < 1) def _fast_rotational_alignment(self, A, B): """ Perform fast rotational alignment using Polar Fourier cross correlation. """ - if not isinstance(A, Image): - A = Image(A) - if not isinstance(B, Image): - B = Image(B) + # if not isinstance(A, Image): + # A = Image(A) + # if not isinstance(B, Image): + # B = Image(B) pftA = self._pft.half_to_full(self._pft.transform(A)) pftB = self._pft.half_to_full(self._pft.transform(B)) @@ -820,15 +820,15 @@ def _fast_rotational_alignment(self, A, B): pftA = fft.fft(pftA, axis=-2) pftB = fft.fft(pftB, axis=-2) x = pftA * pftB.conj() - angular = np.sum(np.abs(fft.ifft2(x)), axis=-1) # sum all radial contributions + angular = xp.sum(xp.abs(fft.ifft2(x)), axis=-1) # sum all radial contributions # Resolve the angle maximizing the correlation through the angular dimension - inds = np.argmax(angular, axis=-1) + inds = xp.argmax(angular, axis=-1) max_thetas_deg = 360 / self._pft.ntheta * inds - max_thetas = np.deg2rad(max_thetas_deg) - peaks = np.take_along_axis(angular, inds.reshape(-1, 1), axis=1).flatten() + max_thetas = xp.deg2rad(max_thetas_deg) + peaks = xp.take_along_axis(angular, inds.reshape(-1, 1), axis=1).flatten() - return max_thetas, peaks + return xp.asnumpy(max_thetas), xp.asnumpy(peaks) def align(self, classes, reflections, basis_coefficients=None): """ @@ -889,20 +889,20 @@ def align(self, classes, reflections, basis_coefficients=None): # number of test rotations. # Note the base original_image[0] should remain unprocessed - _images = original_images[1:].asnumpy().copy() + _images = xp.array(original_images[1:].asnumpy()) # implicit .copy() # Skip zero shifting. template_image = original_images[0] if np.any(shift != 0): template_image = template_image.shift(shift) # mask - template_image = template_image * self._mask + template_image = xp.asarray(template_image) * self._mask _images = _images * self._mask # XXXX think if we need to do this before shifting?! # Handle reflections refl = reflections[k][1:] # skips original_image 0 - _images[refl] = np.flipud(_images[refl]) + _images[refl] = xp.flipud(_images[refl]) # Compute and assign the best rotation found with this translation # note offset of 1 for skipped original_image 0 diff --git a/src/aspire/operators/polar_ft.py b/src/aspire/operators/polar_ft.py index 2b90d5cad6..6f661f46cf 100644 --- a/src/aspire/operators/polar_ft.py +++ b/src/aspire/operators/polar_ft.py @@ -4,6 +4,7 @@ from aspire.image import Image from aspire.nufft import nufft +from aspire.numeric import xp from aspire.utils import complex_type logger = logging.getLogger(__name__) @@ -100,13 +101,13 @@ def transform(self, x): f" Inconsistent dtypes x: {x.dtype} self: {self.dtype}" ) - if not isinstance(x, Image): - raise TypeError( - f"{self.__class__.__name__}.transform" - f" passed numpy array instead of {Image}." - ) - else: - x = x.asnumpy() + # if not isinstance(x, Image): + # raise TypeError( + # f"{self.__class__.__name__}.transform" + # f" passed numpy array instead of {Image}." + # ) + if isinstance(x, Image): + x = xp.asarray(x.asnumpy()) # Flatten stack stack_shape = x.shape[: -self.ndim] @@ -114,7 +115,7 @@ def transform(self, x): # We expect the Image `x` to be real in order to take advantage of the conjugate # symmetry of the Fourier transform of a real valued image. - if not np.isreal(x).all(): + if not xp.isreal(x).all(): raise TypeError( f"The Image `x` must be real valued. Found dtype {x.dtype}." ) @@ -136,4 +137,9 @@ def half_to_full(pf): :return: The full polar Fourier transform with shape (*stack_shape, ntheta, nrad) """ - return np.concatenate((pf, np.conj(pf)), axis=-2) + # cheap way to interop for now + concatenate = xp.concatenate + if isinstance(pf, np.ndarray): + concatenate = np.concatenate + + return concatenate((pf, pf.conj()), axis=-2) From d5af43a462c2faac46dd6ef8ee37d9c6c5d49eec Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 28 Mar 2025 08:37:52 -0400 Subject: [PATCH 134/216] factor out the pft --- src/aspire/classification/averager2d.py | 50 +++++++++++-------------- 1 file changed, 22 insertions(+), 28 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 387acc370c..c3b98652a2 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -803,23 +803,15 @@ def __init__( self._pft = PolarFT(self.src.L, ntheta=ntheta, nrad=nrad, dtype=self.dtype) self._mask = xp.asarray(grid_2d(self.src.L, normalized=True)["r"] < 1) - def _fast_rotational_alignment(self, A, B): + def _fast_rotational_alignment(self, pfA, pfB): """ Perform fast rotational alignment using Polar Fourier cross correlation. """ - # if not isinstance(A, Image): - # A = Image(A) - # if not isinstance(B, Image): - # B = Image(B) - - pftA = self._pft.half_to_full(self._pft.transform(A)) - pftB = self._pft.half_to_full(self._pft.transform(B)) - # 2 hats one sum - pftA = fft.fft(pftA, axis=-2) - pftB = fft.fft(pftB, axis=-2) - x = pftA * pftB.conj() + pfA = fft.fft(pfA, axis=-2) + pfB = fft.fft(pfB, axis=-2) + x = pfA * pfB.conj() angular = xp.sum(xp.abs(fft.ifft2(x)), axis=-1) # sum all radial contributions # Resolve the angle maximizing the correlation through the angular dimension @@ -870,7 +862,19 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) - # Loop over shift search space, updating best result + _images = xp.array(original_images[1:].asnumpy()) # implicit .copy() + + # Handle reflections + refl = reflections[k][1:] # skips original_image 0 + _images[refl] = xp.flipud(_images[refl]) + + # Mask off + _images = _images * self._mask + + # Convert to polar Fourier + pf_images = self._pft.half_to_full(self._pft.transform(_images)) + + # Loop over shift search space, updating best results for x, y in tqdm( zip(x_shifts, y_shifts), total=len(x_shifts), @@ -881,33 +885,23 @@ def align(self, classes, reflections, basis_coefficients=None): shift = np.array([x, y], dtype=shifts.dtype) logger.debug(f"Computing rotational alignment after shift {shift}.") - # For each shift, the set of neighbor images is shifted. - # This order is chosen because: - # i) allows concatenation of shifts and rotation - # operations after orientation estimation - # ii) because generally the number of neighbors << the - # number of test rotations. - # Note the base original_image[0] should remain unprocessed - _images = xp.array(original_images[1:].asnumpy()) # implicit .copy() # Skip zero shifting. template_image = original_images[0] if np.any(shift != 0): template_image = template_image.shift(shift) - # mask + # Mask off template_image = xp.asarray(template_image) * self._mask - _images = _images * self._mask - # XXXX think if we need to do this before shifting?! - # Handle reflections - refl = reflections[k][1:] # skips original_image 0 - _images[refl] = xp.flipud(_images[refl]) + pf_template_image = self._pft.half_to_full( + self._pft.transform(template_image) + ) # Compute and assign the best rotation found with this translation # note offset of 1 for skipped original_image 0 _rotations[1:], _dot_products[1:] = self._fast_rotational_alignment( - template_image, _images[:] + pf_template_image, pf_images ) # Test and update From 32d56bf03a05970a5ff41d3c9881343645e2c878 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 28 Mar 2025 09:47:11 -0400 Subject: [PATCH 135/216] begin batching, two places to broadcast --- src/aspire/classification/averager2d.py | 80 ++++++++++++++----------- 1 file changed, 45 insertions(+), 35 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index c3b98652a2..3c799a80ca 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -851,6 +851,13 @@ def align(self, classes, reflections, basis_coefficients=None): sub_pixel=self.sub_pixel, ) + # XXX + # maybe just change _shift_search_grid to output this later, + # it was first written that way to use sequential pixelrolls in each dimension, + + # (num_shifts, 2) + test_shifts = np.stack((x_shifts, y_shifts), axis=1) + for k in trange(n_classes, desc="Rotationally aligning classes"): # We want to locally cache the original images, # because we will mutate them with shifts in the next loop. @@ -874,54 +881,57 @@ def align(self, classes, reflections, basis_coefficients=None): # Convert to polar Fourier pf_images = self._pft.half_to_full(self._pft.transform(_images)) - # Loop over shift search space, updating best results - for x, y in tqdm( - zip(x_shifts, y_shifts), - total=len(x_shifts), + # Batch over shift search space, updating best results + for start in trange( + 0, + len(test_shifts), + self.batch_size, desc="\tmaximizing over shifts", disable=len(x_shifts) == 1, leave=False, ): - shift = np.array([x, y], dtype=shifts.dtype) - logger.debug(f"Computing rotational alignment after shift {shift}.") - # Note the base original_image[0] should remain unprocessed - # Skip zero shifting. - template_image = original_images[0] - if np.any(shift != 0): - template_image = template_image.shift(shift) + end = min(start + self.batch_size, len(test_shifts)) + batch_shifts = test_shifts[start:end] + + template_images = xp.zeros( + (len(batch_shifts), self.src.L, self.src.L), dtype=self.dtype + ) + + # HACK, unwind later by broadcasting in `.shift` + for i, shift in enumerate(batch_shifts): + + # Note the base original_image[0] should remain unprocessed + # Skip zero shifting. + template_image = original_images[0] + if np.any(shift != 0): + template_image = template_image.shift(shift) + template_images[i] = xp.asarray(template_image) # Mask off - template_image = xp.asarray(template_image) * self._mask + template_images = template_images * self._mask - pf_template_image = self._pft.half_to_full( - self._pft.transform(template_image) + pf_template_images = self._pft.half_to_full( + self._pft.transform(template_images) ) - # Compute and assign the best rotation found with this translation - # note offset of 1 for skipped original_image 0 - _rotations[1:], _dot_products[1:] = self._fast_rotational_alignment( - pf_template_image, pf_images - ) + # unwind, table broadcast in _fast_rotational_alignment + for shift, pf_template_image in zip(batch_shifts, pf_template_images): - # Test and update - # Each base-neighbor pair may have a best shift+rot from a different shift iteration. - improved_indices = _dot_products > dot_products[k] - rotations[k, improved_indices] = _rotations[improved_indices] - dot_products[k, improved_indices] = _dot_products[improved_indices] - # base shifts assigned here, commutation resolved end of loop - shifts[k, improved_indices] = shift - - if (x, y) == (0, 0): - logger.debug("Initial rotational alignment complete (shift (0,0))") - # skipped original_image 0, f"{np.sum(improved_indices)} =?= {np.size(classes)}" - expected = np.size(classes[0]) - 1 - assert np.sum(improved_indices) == expected - else: - logger.debug( - f"Shift ({x},{y}) complete. Improved {np.sum(improved_indices)} alignments." + # Compute and assign the best rotation found with this translation + # note offset of 1 for skipped original_image 0 + _rotations[1:], _dot_products[1:] = self._fast_rotational_alignment( + pf_template_image, pf_images ) + # Test and update + # Each base-neighbor pair may have a best shift+rot from a different shift iteration. + improved_indices = _dot_products > dot_products[k] + rotations[k, improved_indices] = _rotations[improved_indices] + dot_products[k, improved_indices] = _dot_products[improved_indices] + # base shifts assigned here, commutation resolved end of loop + shifts[k, improved_indices] = shift + # Commute the rotation and shift (code shifted the base image instead of all class members) shifts[k] = commute_shift_rot(shifts[k], -rotations[k]) From e9d44bdda3433274579f10653fd7cd31c582386c Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 28 Mar 2025 10:35:35 -0400 Subject: [PATCH 136/216] table broadcast polar cross corr --- src/aspire/classification/averager2d.py | 75 ++++++++++++++++--------- 1 file changed, 48 insertions(+), 27 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 3c799a80ca..e4cc80d542 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -806,19 +806,35 @@ def __init__( def _fast_rotational_alignment(self, pfA, pfB): """ Perform fast rotational alignment using Polar Fourier cross correlation. + + Note broadcasting is specialized for this problem. + pfA.shape (m, nt, nr) + pfB.shape (n, nt, nr) + yields thetas (m,n), peaks (m,n) + """ + if pfA.ndim == 2: + pfA = pfA[None] + if pfB.ndim == 2: + pfB = pfB[None] + # 2 hats one sum pfA = fft.fft(pfA, axis=-2) pfB = fft.fft(pfB, axis=-2) - x = pfA * pfB.conj() + # x = pfA * pfB.conj() + x = xp.expand_dims(pfA, 1) * xp.expand_dims(pfB.conj(), 0) angular = xp.sum(xp.abs(fft.ifft2(x)), axis=-1) # sum all radial contributions # Resolve the angle maximizing the correlation through the angular dimension inds = xp.argmax(angular, axis=-1) max_thetas_deg = 360 / self._pft.ntheta * inds max_thetas = xp.deg2rad(max_thetas_deg) - peaks = xp.take_along_axis(angular, inds.reshape(-1, 1), axis=1).flatten() + peaks = xp.take_along_axis(angular, inds[..., None], axis=-1).squeeze(-1) + + # sanity check, can mv to unit test later + assert max_thetas.shape == peaks.shape + assert max_thetas.shape == (pfA.shape[0], pfB.shape[0]) return xp.asnumpy(max_thetas), xp.asnumpy(peaks) @@ -838,10 +854,6 @@ def align(self, classes, reflections, basis_coefficients=None): dot_products = np.ones((n_classes, n_nbor), dtype=self.dtype) * -np.inf shifts = np.zeros((*classes.shape, 2), dtype=self.dtype) - # Work arrays - _rotations = np.zeros((n_nbor), dtype=self.dtype) - _dot_products = np.ones((n_nbor), dtype=self.dtype) * -np.inf - # Create a search grid and force initial pair to (0,0) # This is done primarily in case of a tie later, we would prefer unshifted. x_shifts, y_shifts = self._shift_search_grid( @@ -858,6 +870,12 @@ def align(self, classes, reflections, basis_coefficients=None): # (num_shifts, 2) test_shifts = np.stack((x_shifts, y_shifts), axis=1) + # Work arrays + bs = min(self.batch_size, len(test_shifts)) + _rotations = np.zeros((bs, n_nbor), dtype=self.dtype) + _dot_products = np.ones((bs, n_nbor), dtype=self.dtype) * -np.inf + template_images = xp.empty((bs, self.src.L, self.src.L), dtype=self.dtype) + for k in trange(n_classes, desc="Rotationally aligning classes"): # We want to locally cache the original images, # because we will mutate them with shifts in the next loop. @@ -882,23 +900,19 @@ def align(self, classes, reflections, basis_coefficients=None): pf_images = self._pft.half_to_full(self._pft.transform(_images)) # Batch over shift search space, updating best results - for start in trange( - 0, - len(test_shifts), - self.batch_size, + pbar = tqdm( + total=len(test_shifts), desc="\tmaximizing over shifts", disable=len(x_shifts) == 1, leave=False, - ): - + ) + for start in range(0, len(test_shifts), self.batch_size): end = min(start + self.batch_size, len(test_shifts)) + bs = end - start # handle a small last batch batch_shifts = test_shifts[start:end] - template_images = xp.zeros( - (len(batch_shifts), self.src.L, self.src.L), dtype=self.dtype - ) - # HACK, unwind later by broadcasting in `.shift` + # template_images[:bs] = template_image.shift(batch_shifts) for i, shift in enumerate(batch_shifts): # Note the base original_image[0] should remain unprocessed @@ -915,22 +929,29 @@ def align(self, classes, reflections, basis_coefficients=None): self._pft.transform(template_images) ) - # unwind, table broadcast in _fast_rotational_alignment - for shift, pf_template_image in zip(batch_shifts, pf_template_images): + # # Compute and assign the best rotation found with this translation + # # note offset of 1 for skipped original_image 0 + _rotations[:bs, 1:], _dot_products[:bs, 1:] = ( + self._fast_rotational_alignment(pf_template_images[:bs], pf_images) + ) - # Compute and assign the best rotation found with this translation - # note offset of 1 for skipped original_image 0 - _rotations[1:], _dot_products[1:] = self._fast_rotational_alignment( - pf_template_image, pf_images - ) + # vectorize these + for i in range(bs): # Test and update # Each base-neighbor pair may have a best shift+rot from a different shift iteration. - improved_indices = _dot_products > dot_products[k] - rotations[k, improved_indices] = _rotations[improved_indices] - dot_products[k, improved_indices] = _dot_products[improved_indices] + improved_indices = _dot_products[i] > dot_products[k] + rotations[k, improved_indices] = _rotations[i, improved_indices] + dot_products[k, improved_indices] = _dot_products[ + i, improved_indices + ] # base shifts assigned here, commutation resolved end of loop - shifts[k, improved_indices] = shift + shifts[k, improved_indices] = batch_shifts[i] + + pbar.update(bs) + + # Completed batching over shifts + pbar.close() # Commute the rotation and shift (code shifted the base image instead of all class members) shifts[k] = commute_shift_rot(shifts[k], -rotations[k]) From b86ee437e824ee83b651d79859178242dab10dcc Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 28 Mar 2025 14:29:48 -0400 Subject: [PATCH 137/216] table broadcast shifts, resuse arrays, reduce mem cost some speed [skip ci] --- src/aspire/classification/averager2d.py | 24 ++++++++++-------------- src/aspire/image/image.py | 25 +++++++++++++++++-------- 2 files changed, 27 insertions(+), 22 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index e4cc80d542..d39db88a20 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -875,6 +875,7 @@ def align(self, classes, reflections, basis_coefficients=None): _rotations = np.zeros((bs, n_nbor), dtype=self.dtype) _dot_products = np.ones((bs, n_nbor), dtype=self.dtype) * -np.inf template_images = xp.empty((bs, self.src.L, self.src.L), dtype=self.dtype) + _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"): # We want to locally cache the original images, @@ -887,14 +888,14 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) - _images = xp.array(original_images[1:].asnumpy()) # implicit .copy() + _images[:] = xp.asarray(original_images[1:].asnumpy()) - # Handle reflections + # Handle reflections, XXX confirm location in sequence refl = reflections[k][1:] # skips original_image 0 _images[refl] = xp.flipud(_images[refl]) # Mask off - _images = _images * self._mask + _images[:] = _images[:] * self._mask # Convert to polar Fourier pf_images = self._pft.half_to_full(self._pft.transform(_images)) @@ -911,19 +912,14 @@ def align(self, classes, reflections, basis_coefficients=None): bs = end - start # handle a small last batch batch_shifts = test_shifts[start:end] - # HACK, unwind later by broadcasting in `.shift` - # template_images[:bs] = template_image.shift(batch_shifts) - for i, shift in enumerate(batch_shifts): - - # Note the base original_image[0] should remain unprocessed - # Skip zero shifting. - template_image = original_images[0] - if np.any(shift != 0): - template_image = template_image.shift(shift) - template_images[i] = xp.asarray(template_image) + # Note the base original_image[0] needs to remain unprocessed + # XXX This is shifting for zero, consider carving that out. + template_images[:bs] = xp.asarray( + original_images[0].shift(batch_shifts) + ) # Mask off - template_images = template_images * self._mask + template_images[:] = template_images[:] * self._mask pf_template_images = self._pft.half_to_full( self._pft.transform(template_images) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 8fcf60fa5f..bb5ec9d999 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -418,7 +418,8 @@ def shift(self, shifts): if not shifts.shape[1] == 2: raise ValueError("Input shifts must be of shape (n_images, 2) or (1, 2).") - if not n_shifts == 1 and not n_shifts == self.n_images: + + if not (n_shifts == 1 or self.n_images == 1 or n_shifts == self.n_images): raise ValueError( "The number of shifts must be 1 or equal to self.n_images." ) @@ -686,26 +687,34 @@ def load(filepath, dtype=None): def _im_translate(self, shifts): """ - Translate image by shifts + Translate image by `shifts`. + + Note broadcasting special case + Image shape (n,L,L) x shifts shape (n,2) -> (n,L,L) shifted images + Image shape (1,L,L) x shifts shape (n,2) -> (n,L,L) shifted images - :param im: An array of size n-by-L-by-L containing images to be translated. + :param im: An array of size m-by-L-by-L containing images to be translated. + m may be 1 or n. :param shifts: An array of size n-by-2 specifying the shifts in pixels. Alternatively, it can be a row vector of length 2, in which case the same shifts is applied to each image. :return: The images translated by the shifts, with periodic boundaries. """ - # Note original stack shape and flatten stack - stack_shape = self.stack_shape - im = self.stack_reshape(-1)._data - if shifts.ndim == 1: shifts = shifts[np.newaxis, :] n_shifts = shifts.shape[0] assert shifts.shape[-1] == 2, "shifts must be nx2" + # Note original stack shape and flatten stack + stack_shape = self.stack_shape + if self.n_images == 1 and n_shifts > 1: + # XXX special case, cleanup broadcasting later + stack_shape = n_shifts + im = self.stack_reshape(-1)._data + assert ( - n_shifts == 1 or n_shifts == self.n_images + n_shifts == 1 or self.n_images == 1 or n_shifts == self.n_images ), "number of shifts must be 1 or match the number of images" # Cast shifts to this instance's internal dtype shifts = xp.asarray(shifts, dtype=self.dtype) From 15e6a9b9468f3e36d9f7b72dcfffea2923fb211c Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 23 Apr 2025 09:17:28 -0400 Subject: [PATCH 138/216] Cleanup unit test for broadcast case --- src/aspire/image/image.py | 2 +- tests/test_image.py | 20 +++++++++++++++----- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index bb5ec9d999..ed33e335c6 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -408,7 +408,7 @@ def shift(self, shifts): :param shifts: An array of size n-by-2 specifying the shifts in pixels. Alternatively, it can be a column vector of length 2, in which case - the same shifts is applied to each image. + the same shift is applied to each image. :return: The Image translated by the shifts, with periodic boundaries. """ if shifts.ndim == 1: diff --git a/tests/test_image.py b/tests/test_image.py index 4e8513a2f5..70f5906a64 100644 --- a/tests/test_image.py +++ b/tests/test_image.py @@ -148,14 +148,24 @@ def testImShiftStack(get_stacks, dtype): np.testing.assert_allclose(im0.asnumpy(), im3, atol=atol) -def testImageShiftErrors(get_images): - _, im = get_images - # test bad shift shape +def testImageShiftShapeErrors(): + # Test images + im = Image(np.ones((1, 8, 8))) + im3 = Image(np.ones((3, 8, 8))) + + # Single image, broadcast multiple shifts is allowed + _ = im.shift(np.array([[100, 200], [100, 200]])) + + # Multiple image, broadcast single shifts is allowed + _ = im3.shift(np.array([[100, 200]])) + + # Bad shift shape, must be (..., 2) with pytest.raises(ValueError, match="Input shifts must be of shape"): _ = im.shift(np.array([100, 100, 100])) - # test bad number of shifts + + # Incoherent number of shifts (number of images != number of shifts when neither 1). with pytest.raises(ValueError, match="The number of shifts"): - _ = im.shift(np.array([[100, 200], [100, 200]])) + _ = im3.shift(np.array([[100, 200], [100, 200]])) def testImageSqrt(get_images, get_stacks): From 491e9513dbcad5c3490702bcd8454a4731799d7b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 23 Apr 2025 10:29:57 -0400 Subject: [PATCH 139/216] cleanup pft interop --- src/aspire/classification/averager2d.py | 4 ++-- src/aspire/operators/polar_ft.py | 32 ++++++++++++++++++------- 2 files changed, 25 insertions(+), 11 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index d39db88a20..d7fbcbb11e 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -898,7 +898,7 @@ def align(self, classes, reflections, basis_coefficients=None): _images[:] = _images[:] * self._mask # Convert to polar Fourier - pf_images = self._pft.half_to_full(self._pft.transform(_images)) + pf_images = self._pft.half_to_full(self._pft._transform(_images)) # Batch over shift search space, updating best results pbar = tqdm( @@ -922,7 +922,7 @@ def align(self, classes, reflections, basis_coefficients=None): template_images[:] = template_images[:] * self._mask pf_template_images = self._pft.half_to_full( - self._pft.transform(template_images) + self._pft._transform(template_images) ) # # Compute and assign the best rotation found with this translation diff --git a/src/aspire/operators/polar_ft.py b/src/aspire/operators/polar_ft.py index 6f661f46cf..1218a22010 100644 --- a/src/aspire/operators/polar_ft.py +++ b/src/aspire/operators/polar_ft.py @@ -89,26 +89,39 @@ def transform(self, x): """ Evaluate coefficient in polar Fourier grid from those in standard 2D coordinate basis - :param x: The Image instance representing coefficient array in the + :param x: The `Image` instance representing coefficient array in the standard 2D coordinate basis to be evaluated. + :return: Numpy array holding the evaluation of the coefficient + array `x` in the polar Fourier grid. This is an array of + vectors whose first dimension corresponds to `x.shape[0]`, + and last dimension equals `self.count`. + """ + if not isinstance(x, Image): + raise TypeError( + f"{self.__class__.__name__}.transform" + f" passed numpy array instead of {Image}." + ) + + return xp.asnumpy(self._transform(x.asnumpy())) + + def _transform(self, x): + """ + Evaluate coefficient in polar Fourier grid from those in standard 2D coordinate basis + + :param x: Coefficients array in the standard 2D coordinate basis to be evaluated. :return: The evaluation of the coefficient array `x` in the polar Fourier grid. This is an array of vectors whose first dimension corresponds to `x.shape[0]`, and last dimension equals `self.count`. """ + + x = xp.asarray(x) + if x.dtype != self.dtype: raise TypeError( f"{self.__class__.__name__}.transform" f" Inconsistent dtypes x: {x.dtype} self: {self.dtype}" ) - # if not isinstance(x, Image): - # raise TypeError( - # f"{self.__class__.__name__}.transform" - # f" passed numpy array instead of {Image}." - # ) - if isinstance(x, Image): - x = xp.asarray(x.asnumpy()) - # Flatten stack stack_shape = x.shape[: -self.ndim] x = x.reshape(-1, *x.shape[-self.ndim :]) @@ -122,6 +135,7 @@ def transform(self, x): resolution = x.shape[-1] + # nufft call should return `pf` as array type (np or cp) of `x` pf = nufft(x, self.freqs) / resolution**2 return pf.reshape(*stack_shape, self.ntheta // 2, self.nrad) From 1e471bcadbd43aaf792fc990f1cad99bf79d54e5 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 23 Apr 2025 10:43:31 -0400 Subject: [PATCH 140/216] A little more cleanup --- src/aspire/classification/averager2d.py | 10 ++++------ src/aspire/image/image.py | 2 +- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index d7fbcbb11e..2953785128 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -17,6 +17,10 @@ def commute_shift_rot(shifts, rots): """ Rotate `shifts` points by `rots` ccw radians. + + :param shifts: Array of shift points shaped (..., 2) + :param rots: Array of rotations (radians) + :returns: Array of rotated shift points shaped (..., 2) """ sx = shifts[:, 0] sy = shifts[:, 1] @@ -778,12 +782,6 @@ def __init__( self.n_angles = n_angles - # # XXX Will use polar for rotate - # if not hasattr(self.alignment_basis, "rotate"): - # raise RuntimeError( - # f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `rotate` method." - # ) - self.radius = radius if radius is not None else src.L // 16 if self.radius != 0: diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index ed33e335c6..8ee9aec96d 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -709,7 +709,7 @@ def _im_translate(self, shifts): # Note original stack shape and flatten stack stack_shape = self.stack_shape if self.n_images == 1 and n_shifts > 1: - # XXX special case, cleanup broadcasting later + # Handle the shift broadcast special case stack_shape = n_shifts im = self.stack_reshape(-1)._data From eb89388a19445c34ad34cbadee8faf0384e4eee0 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 23 Apr 2025 11:29:34 -0400 Subject: [PATCH 141/216] stash --- src/aspire/classification/averager2d.py | 37 ++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 4 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 2953785128..fa2c79abba 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -801,7 +801,7 @@ def __init__( self._pft = PolarFT(self.src.L, ntheta=ntheta, nrad=nrad, dtype=self.dtype) self._mask = xp.asarray(grid_2d(self.src.L, normalized=True)["r"] < 1) - def _fast_rotational_alignment(self, pfA, pfB): + def _fast_rotational_alignment(self, pfA, pfB, do_interp=False): """ Perform fast rotational alignment using Polar Fourier cross correlation. @@ -826,9 +826,38 @@ def _fast_rotational_alignment(self, pfA, pfB): # Resolve the angle maximizing the correlation through the angular dimension inds = xp.argmax(angular, axis=-1) - max_thetas_deg = 360 / self._pft.ntheta * inds - max_thetas = xp.deg2rad(max_thetas_deg) - peaks = xp.take_along_axis(angular, inds[..., None], axis=-1).squeeze(-1) + breakpoint() + + if do_interp: + half_width = 5 + fine_steps = 100 + thetas = np.linspace(0, 2*np.pi, self._pft.ntheta) + shp = (pfA.shape[0], pfB.shape[0]) + max_thetas = np.empty(shp, dtype=self.dtype) + peaks = np.empty(shp, dtype=self.dtype) + + for i,ind in enumerate(inds): + # Select windows around peak + x = thetas[ind-half_width:ind+half_width] + y = angular[ind-half_width:ind+half_width] + + # Setup an interpolator for the window + f_interp = interp1d(x,y,kind="cubic") + + # fine grid + xfine = np.linspace(thetas[0], thetas[-1], fine_steps) + yfine = f_interp(xfine) + + indfine = xp.argmax(yfine) + max_thetas[i] = xfine[indfine] + peaks[i] = y_fine[indfine] + + # # opt + # # Search for max in interpolation window + + else: + max_thetas = 2*np.pi / self._pft.ntheta * inds + peaks = xp.take_along_axis(angular, inds[..., None], axis=-1).squeeze(-1) # sanity check, can mv to unit test later assert max_thetas.shape == peaks.shape From a6f8af133c451815640b4a4578590593f66b0794 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 24 Apr 2025 15:20:55 -0400 Subject: [PATCH 142/216] add fine interp and optimize methods --- src/aspire/classification/averager2d.py | 71 +++++++++++++++++-------- 1 file changed, 48 insertions(+), 23 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index fa2c79abba..1870165e29 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -2,6 +2,8 @@ from abc import ABC, abstractmethod import numpy as np +from scipy.interpolate import interp1d +from scipy.optimize import minimize_scalar from aspire.basis import Coef from aspire.classification.reddy_chatterji import reddy_chatterji_register @@ -826,37 +828,60 @@ def _fast_rotational_alignment(self, pfA, pfB, do_interp=False): # Resolve the angle maximizing the correlation through the angular dimension inds = xp.argmax(angular, axis=-1) - breakpoint() if do_interp: half_width = 5 fine_steps = 100 - thetas = np.linspace(0, 2*np.pi, self._pft.ntheta) + thetas = np.linspace(0, 2 * np.pi, self._pft.ntheta, endpoint=False) shp = (pfA.shape[0], pfB.shape[0]) max_thetas = np.empty(shp, dtype=self.dtype) peaks = np.empty(shp, dtype=self.dtype) - - for i,ind in enumerate(inds): - # Select windows around peak - x = thetas[ind-half_width:ind+half_width] - y = angular[ind-half_width:ind+half_width] - - # Setup an interpolator for the window - f_interp = interp1d(x,y,kind="cubic") - - # fine grid - xfine = np.linspace(thetas[0], thetas[-1], fine_steps) - yfine = f_interp(xfine) - - indfine = xp.argmax(yfine) - max_thetas[i] = xfine[indfine] - peaks[i] = y_fine[indfine] - - # # opt - # # Search for max in interpolation window - + + for i in range(inds.shape[0]): + for j in range(inds.shape[1]): + ind = inds[i, j] + + # Select windows around peak + # Want slice, [ind-half_width:ind+half_width], with wrapping + # Note, could alternatively use halfwidth "pad" with wrap + window = range(ind - half_width, ind + half_width) + xw = thetas.take(window, mode="wrap") + mask = xw < xw[0] + xw[mask] = xw[mask] + 2 * np.pi + yw = angular[i, j].take(window, mode="wrap") + + # Setup an interpolator for the window + f_interp = interp1d(xw, yw, kind="cubic") + + if do_interp == "opt": + # Negate the function we want to maximize + def f(x, _f=f_interp): + return -_f(x) + + # Call the optimizer + res = minimize_scalar(f, bounds=(xw[0], xw[-1])) + + # Assign results + max_thetas[i, j] = res.x + peaks[i, j] = f_interp(res.x) + + else: + # Create fine grid window + xfine = np.linspace(xw[0], xw[-1], fine_steps) + yfine = f_interp(xfine) + + # Find the maximal value in the fine grid window + indfine = xp.argmax(yfine) + + # Assign results + max_thetas[i, j] = xfine[indfine] + peaks[i, j] = yfine[indfine] + + # Modulate the interpolants wraping around the circle. + max_thetas = max_thetas % (2 * np.pi) + else: - max_thetas = 2*np.pi / self._pft.ntheta * inds + max_thetas = 2 * np.pi / self._pft.ntheta * inds peaks = xp.take_along_axis(angular, inds[..., None], axis=-1).squeeze(-1) # sanity check, can mv to unit test later From 2d7d8cb0ce3a641cb8ecdfe17d3ba820293b48c2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 25 Apr 2025 09:54:07 -0400 Subject: [PATCH 143/216] add BFTAverager2D to test suite --- src/aspire/classification/averager2d.py | 6 +++--- tests/test_averager2d.py | 5 +++++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 1870165e29..c309ab93cf 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -989,12 +989,12 @@ def align(self, classes, reflections, basis_coefficients=None): # Test and update # Each base-neighbor pair may have a best shift+rot from a different shift iteration. improved_indices = _dot_products[i] > dot_products[k] - rotations[k, improved_indices] = _rotations[i, improved_indices] + rotations[k, improved_indices] = -_rotations[i, improved_indices] dot_products[k, improved_indices] = _dot_products[ i, improved_indices ] # base shifts assigned here, commutation resolved end of loop - shifts[k, improved_indices] = batch_shifts[i] + shifts[k, improved_indices] = -batch_shifts[i] pbar.update(bs) @@ -1002,7 +1002,7 @@ def align(self, classes, reflections, basis_coefficients=None): pbar.close() # Commute the rotation and shift (code shifted the base image instead of all class members) - shifts[k] = commute_shift_rot(shifts[k], -rotations[k]) + shifts[k] = commute_shift_rot(shifts[k], rotations[k]) return rotations, shifts, dot_products diff --git a/tests/test_averager2d.py b/tests/test_averager2d.py index becf0ad2d1..2eaa355755 100644 --- a/tests/test_averager2d.py +++ b/tests/test_averager2d.py @@ -12,6 +12,7 @@ BFRAverager2D, BFSRAverager2D, BFSReddyChatterjiAverager2D, + BFTAverager2D, ReddyChatterjiAverager2D, ) from aspire.operators import PolarFT @@ -299,3 +300,7 @@ def testAverager(self): class BFSReddyChatterjiAverager2DTestCase(ReddyChatterjiAverager2DTestCase): averager = BFSReddyChatterjiAverager2D + + +class BFTAverager2DTestCase(BFSRAverager2DTestCase): + averager = BFTAverager2D From ee5f5cb4aad3c4580b736baf5bf8f63eef0aa15a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 28 Apr 2025 15:40:20 -0400 Subject: [PATCH 144/216] intial add BFT to source wrappers, remove 110 --- src/aspire/classification/averager2d.py | 29 +++++++---- src/aspire/denoising/class_avg.py | 67 +++++++++++++------------ tests/test_class_src.py | 2 - 3 files changed, 56 insertions(+), 42 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index c309ab93cf..e8b13ef191 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -460,6 +460,12 @@ def align(self, classes, reflections, basis_coefficients=None): class BFRAverager2D(BFSRAverager2D): + """ + Brute Force Rotation only reference implementation. + + See BFT with `radius=0` for a more performant implementation using a fast rotational alignment. + """ + def __init__(self, *args, **kwargs): super().__init__(*args, radius=0, **kwargs) @@ -751,7 +757,7 @@ class BFTAverager2D(AligningAverager2D): This perfoms a Brute Force Translations and fast rotational alignment. For each pair of x_shifts and y_shifts, - Perform cross correlation based rotational alignment + Perform polar Fourier cross correlation based rotational alignment. Return the rotation and shift yielding the best results. """ @@ -762,17 +768,21 @@ def __init__( src, alignment_basis=None, n_angles=360, + n_radial=None, radius=None, sub_pixel=0.1, batch_size=512, dtype=None, ): """ - See AligningAverager2D adds `n_angles` and `radius`. + See AligningAverager2D. Adds `n_angles`, `n_radial`, `radius`, `sub_pixel`. - :params n_angles: Number of brute force rotations to attempt, defaults 360. + :params n_angles: Number of PFT angular components, defaults 360. + :param n_radial: Number of PFT radial components, defaults `self.src.L`. :param radius: Brute force translation search radius. - Defaults to src.L//16. + `0` disables translation search, rotations only. + Defaults to `src.L//16`. + :param sub_pixel: Subpixel shift size used in brute force shift search. """ super().__init__( composite_basis, @@ -784,7 +794,7 @@ def __init__( self.n_angles = n_angles - self.radius = radius if radius is not None else src.L // 16 + self.radius = radius if radius is not None else src.L // 32 if self.radius != 0: @@ -795,12 +805,13 @@ def __init__( self.sub_pixel = sub_pixel - # XXX todo, better config - ntheta = 360 - nrad = self.src.L + # Configure number of radial points + self.n_radial = n_radial or self.src.L # Setup Polar Transform - self._pft = PolarFT(self.src.L, ntheta=ntheta, nrad=nrad, dtype=self.dtype) + self._pft = PolarFT( + self.src.L, ntheta=n_angles, nrad=n_radial, dtype=self.dtype + ) self._mask = xp.asarray(grid_2d(self.src.L, normalized=True)["r"] < 1) def _fast_rotational_alignment(self, pfA, pfB, do_interp=False): diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py index 8aff4c23e4..e46177b08b 100644 --- a/src/aspire/denoising/class_avg.py +++ b/src/aspire/denoising/class_avg.py @@ -8,6 +8,7 @@ BandedSNRImageQualityFunction, BFRAverager2D, BFSRAverager2D, + BFTAverager2D, Class2D, ClassSelector, GlobalVarianceClassSelector, @@ -459,7 +460,7 @@ class LegacyClassAvgSource(ClassAvgSource): Defaults to using global variance based class selection, and a brute force image alignment (rotational only). - This is most similar to what was reported for papers using the + This is similar to what was reported for papers using the MATLAB code. """ @@ -484,7 +485,7 @@ def __init__( :param class_selector: `ClassSelector` instance. Default `None` creates `GlobalVarianceClassSelector`. :param averager: `Averager2D` instance. - Default `None` ceates `BFRAverager2D` instance. + Default `None` ceates `BFTAverager2D` instance. See code for parameter details. :param averager_src: Optionally explicitly assign source to `BFRAverager2D` during initialization. Allows users to @@ -514,9 +515,10 @@ def __init__( basis_2d = self._get_classifier_basis(classifier) - averager = BFRAverager2D( + averager = BFTAverager2D( composite_basis=basis_2d, src=averager_src, + radius=0, # disables translation search batch_size=batch_size, dtype=dtype, ) @@ -573,10 +575,10 @@ def DefaultClassAvgSource( """ _versions = { - None: ClassAvgSourcev132, - "latest": ClassAvgSourcev132, + None: ClassAvgSourcev140, + "latest": ClassAvgSourcev140, + "0.14.0": ClassAvgSourcev140, "0.13.2": ClassAvgSourcev132, - "0.11.0": ClassAvgSourcev110, } if version not in _versions: @@ -594,13 +596,15 @@ def DefaultClassAvgSource( ) -class ClassAvgSourcev132(ClassAvgSource): +class ClassAvgSourcev140(ClassAvgSource): """ Source for denoised 2D images using class average methods. - Defaults to using SNR based class selection, - avoiding neighbors of previous classes, + Defaults to using global variance based class selection, and a brute force image alignment (rotational only). + + This is most similar to what was reported for papers using the + MATLAB code, but takes significant time to compute. """ def __init__( @@ -614,7 +618,7 @@ def __init__( batch_size=512, ): """ - Instantiates ClassAvgSourcev132 with the following parameters. + Instantiates ClassAvgSourcev140 with the following parameters. :param src: Source used for image classification. :param n_nbor: Number of nearest neighbors. Default 50. @@ -622,11 +626,9 @@ def __init__( Default `None` creates `RIRClass2D`. See code for parameter details. :param class_selector: `ClassSelector` instance. - Default `None` creates `GlobalWithRepulsionClassSelector` with - `BandedSNRImageQualityFunction`. This will select the - images with the highest banded SNR. + Default `None` creates `GlobalVarianceClassSelector`. :param averager: `Averager2D` instance. - Default `None` ceates `BFRAverager2D` instance. + Default `None` ceates `BFTAverager2D` instance. See code for parameter details. :param averager_src: Optionally explicitly assign source to `averager` during initialization. Allows users to @@ -656,7 +658,7 @@ def __init__( basis_2d = self._get_classifier_basis(classifier) - averager = BFRAverager2D( + averager = BFTAverager2D( composite_basis=basis_2d, src=averager_src, batch_size=batch_size, @@ -668,10 +670,7 @@ def __init__( ) if class_selector is None: - quality_function = BandedSNRImageQualityFunction() - class_selector = GlobalWithRepulsionClassSelector( - averager, quality_function - ) + class_selector = GlobalVarianceClassSelector(averager=averager) super().__init__( src=src, @@ -682,13 +681,13 @@ def __init__( ) -class ClassAvgSourcev110(ClassAvgSource): +class ClassAvgSourcev132(ClassAvgSource): """ Source for denoised 2D images using class average methods. - Defaults to using Contrast based class selection (on the fly, compressed), + Defaults to using SNR based class selection, avoiding neighbors of previous classes, - and a brute force image alignment. + and a brute force image alignment (rotational only). """ def __init__( @@ -702,7 +701,7 @@ def __init__( batch_size=512, ): """ - Instantiates ClassAvgSourcev110 with the following parameters. + Instantiates ClassAvgSourcev132 with the following parameters. :param src: Source used for image classification. :param n_nbor: Number of nearest neighbors. Default 50. @@ -710,14 +709,17 @@ def __init__( Default `None` creates `RIRClass2D`. See code for parameter details. :param class_selector: `ClassSelector` instance. - Default `None` creates `NeighborVarianceWithRepulsionClassSelector`. + Default `None` creates `GlobalWithRepulsionClassSelector` with + `BandedSNRImageQualityFunction`. This will select the + images with the highest banded SNR. :param averager: `Averager2D` instance. - Default `None` ceates `BFSRAverager2D` instance. + Default `None` ceates `BFRAverager2D` instance. See code for parameter details. - :param averager_src: Optionally explicitly assign source - to BFSRAverager2D during initialization. - Raises error when combined with an explicit `averager` - argument. + :param averager_src: Optionally explicitly assign source to + `averager` during initialization. Allows users to + provide distinct sources for classification and + averaging. Raises error when combined with an explicit + `averager` argument. :param batch_size: Integer size for batched operations. :return: ClassAvgSource instance. @@ -741,7 +743,7 @@ def __init__( basis_2d = self._get_classifier_basis(classifier) - averager = BFSRAverager2D( + averager = BFRAverager2D( composite_basis=basis_2d, src=averager_src, batch_size=batch_size, @@ -753,7 +755,10 @@ def __init__( ) if class_selector is None: - class_selector = NeighborVarianceWithRepulsionClassSelector() + quality_function = BandedSNRImageQualityFunction() + class_selector = GlobalWithRepulsionClassSelector( + averager, quality_function + ) super().__init__( src=src, diff --git a/tests/test_class_src.py b/tests/test_class_src.py index c105bf0a6d..4697f0e92b 100644 --- a/tests/test_class_src.py +++ b/tests/test_class_src.py @@ -29,7 +29,6 @@ DefaultClassAvgSource, LegacyClassAvgSource, ) -from aspire.denoising.class_avg import ClassAvgSourcev110 from aspire.image import Image from aspire.source import RelionSource, Simulation from aspire.utils import Rotation @@ -55,7 +54,6 @@ DebugClassAvgSource, DefaultClassAvgSource, LegacyClassAvgSource, - ClassAvgSourcev110, ] From 4ecbc53a2c4491ab3a1c15e0e5d293e25f43d03d Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 29 Apr 2025 08:58:18 -0400 Subject: [PATCH 145/216] tox checks --- src/aspire/denoising/class_avg.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py index e46177b08b..5479bc75de 100644 --- a/src/aspire/denoising/class_avg.py +++ b/src/aspire/denoising/class_avg.py @@ -7,13 +7,11 @@ Averager2D, BandedSNRImageQualityFunction, BFRAverager2D, - BFSRAverager2D, BFTAverager2D, Class2D, ClassSelector, GlobalVarianceClassSelector, GlobalWithRepulsionClassSelector, - NeighborVarianceWithRepulsionClassSelector, RIRClass2D, TopClassSelector, ) From b37f3c75d2381d6c2ccf4d71164a60a899abf18d Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 29 Apr 2025 16:55:13 -0400 Subject: [PATCH 146/216] flip bug fix --- src/aspire/classification/averager2d.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index e8b13ef191..0235eeeb23 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -953,9 +953,11 @@ def align(self, classes, reflections, basis_coefficients=None): _images[:] = xp.asarray(original_images[1:].asnumpy()) - # Handle reflections, XXX confirm location in sequence - refl = reflections[k][1:] # skips original_image 0 - _images[refl] = xp.flipud(_images[refl]) + # Handle reflections + refl = reflections[k] + _images[refl[1:]] = xp.flip( + _images[refl[1:]], axis=-2 + ) # 1: skips original_image 0 # Mask off _images[:] = _images[:] * self._mask From b3cd1cf464873b8229d8ab8bc5958556a6c0f8fe Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 30 Apr 2025 09:36:23 -0400 Subject: [PATCH 147/216] update shift grid to return array of tuples --- src/aspire/classification/averager2d.py | 37 ++++++++++--------------- 1 file changed, 15 insertions(+), 22 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 0235eeeb23..4d8cc4084d 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -260,7 +260,7 @@ def _shift_search_grid(self, L, radius, roll_zero=False, sub_pixel=1): :param radius: Disc radius in pixels :param roll_zero: Roll (0,0) to zero'th element. Defaults to False. :param sub_pixel: Sub pixel decimation. 1 is integer, 0.1 is 1/10 pixel, etc. - :returns: Grid points as 2-tuple of vectors X,Y. + :returns: Grid points as array of 2-tuples [(x0,y0),... (xi,yi)]. """ # We'll brute force all shifts in a grid. @@ -275,7 +275,9 @@ def _shift_search_grid(self, L, radius, roll_zero=False, sub_pixel=1): X, Y = np.roll(X, -zero_ind), np.roll(Y, -zero_ind) assert (X[0], Y[0]) == (0, 0), (radius, zero_ind, X, Y) - return X, Y + shifts = np.stack((X, Y), axis=1) + + return shifts class BFSRAverager2D(AligningAverager2D): @@ -358,9 +360,7 @@ def align(self, classes, reflections, basis_coefficients=None): # Create a search grid and force initial pair to (0,0) # This is done primarily in case of a tie later, we would take unshifted. - x_shifts, y_shifts = self._shift_search_grid( - self.src.L, self.radius, roll_zero=True - ) + test_shifts = self._shift_search_grid(self.src.L, self.radius, roll_zero=True) for k in trange(n_classes, desc="Rotationally aligning classes"): # We want to locally cache the original images, @@ -391,10 +391,10 @@ def align(self, classes, reflections, basis_coefficients=None): # Loop over shift search space, updating best result for x, y in tqdm( - zip(x_shifts, y_shifts), - total=len(x_shifts), + test_shifts, + total=len(shifts), desc="\tmaximizing over shifts", - disable=len(x_shifts) == 1, + disable=len(shifts) == 1, leave=False, ): shift = np.array([x, y], dtype=int) @@ -687,7 +687,7 @@ def align(self, classes, reflections, basis_coefficients=None): dot_products = np.ones(classes.shape, dtype=self.dtype) * -np.inf shifts = np.zeros((*classes.shape, 2), dtype=int) - X, Y = self._shift_search_grid(self.alignment_src.L, self.radius) + test_shifts = self._shift_search_grid(self.alignment_src.L, self.radius) def _innerloop(k): unshifted_images = self._cls_images(classes[k]) @@ -697,10 +697,10 @@ def _innerloop(k): _shifts = np.zeros((*classes.shape[1:], 2), dtype=int) for xs, ys in tqdm( - zip(X, Y), - total=len(X), + test_shifts, + total=len(test_shifts), desc="\tmaximizing over shifts", - disable=len(X) == 1, + disable=len(test_shifts) == 1, leave=False, ): @@ -756,7 +756,7 @@ class BFTAverager2D(AligningAverager2D): """ This perfoms a Brute Force Translations and fast rotational alignment. - For each pair of x_shifts and y_shifts, + For each shift, Perform polar Fourier cross correlation based rotational alignment. Return the rotation and shift yielding the best results. @@ -919,20 +919,13 @@ def align(self, classes, reflections, basis_coefficients=None): # Create a search grid and force initial pair to (0,0) # This is done primarily in case of a tie later, we would prefer unshifted. - x_shifts, y_shifts = self._shift_search_grid( + test_shifts = self._shift_search_grid( self.src.L, self.radius, roll_zero=True, sub_pixel=self.sub_pixel, ) - # XXX - # maybe just change _shift_search_grid to output this later, - # it was first written that way to use sequential pixelrolls in each dimension, - - # (num_shifts, 2) - test_shifts = np.stack((x_shifts, y_shifts), axis=1) - # Work arrays bs = min(self.batch_size, len(test_shifts)) _rotations = np.zeros((bs, n_nbor), dtype=self.dtype) @@ -969,7 +962,7 @@ def align(self, classes, reflections, basis_coefficients=None): pbar = tqdm( total=len(test_shifts), desc="\tmaximizing over shifts", - disable=len(x_shifts) == 1, + disable=len(test_shifts) == 1, leave=False, ) for start in range(0, len(test_shifts), self.batch_size): From b3b077d73520a5c3fddfc9d56849ff4713d79e14 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 1 May 2025 11:09:39 -0400 Subject: [PATCH 148/216] cleanup --- src/aspire/classification/averager2d.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 4d8cc4084d..b8cf835ccf 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -944,13 +944,11 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) - _images[:] = xp.asarray(original_images[1:].asnumpy()) + _images[:] = xp.asarray(original_images[1:].asnumpy(), copy=True) # Handle reflections - refl = reflections[k] - _images[refl[1:]] = xp.flip( - _images[refl[1:]], axis=-2 - ) # 1: skips original_image 0 + refl = reflections[k][1:] # skips original_image 0 + _images[refl] = xp.flip(_images[refl], axis=-2) # Mask off _images[:] = _images[:] * self._mask @@ -970,8 +968,8 @@ def align(self, classes, reflections, basis_coefficients=None): bs = end - start # handle a small last batch batch_shifts = test_shifts[start:end] - # Note the base original_image[0] needs to remain unprocessed - # XXX This is shifting for zero, consider carving that out. + # Shift the base, original_image[0], for each shift in this batch + # Note this includes shifting for the zero case template_images[:bs] = xp.asarray( original_images[0].shift(batch_shifts) ) From 8d9e744fd4780767f408a448c096c3342bf2496a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 1 May 2025 14:51:34 -0400 Subject: [PATCH 149/216] reversed the index mapping, whoops --- src/aspire/denoising/class_avg.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py index 5479bc75de..7c362fd0a4 100644 --- a/src/aspire/denoising/class_avg.py +++ b/src/aspire/denoising/class_avg.py @@ -289,12 +289,16 @@ def _images(self, indices): # Check if this src cached images. if self._cached_im is not None: - logger.debug(f"Loading {len(indices)} images from image cache") - im = self._cached_im[indices, :, :] + logger.debug( + f"Loading {len(indices)} images from image cache, indices {_indices}" + ) + im = self._cached_im[_indices, :, :] # Check for heap cached image sets from class_selector. elif heap_inds: - logger.debug(f"Mapping {len(heap_inds)} images from heap cache.") + logger.debug( + f"Mapping {len(heap_inds)} images from heap cache, indices {indices}" + ) # Images in heap_inds can be fetched from class_selector. # For others, create an indexing map that preserves @@ -346,9 +350,11 @@ def _images(self, indices): else: # Perform image averaging for the requested images (classes) - logger.debug(f"Averaging {len(_indices)} images from source") + logger.debug( + f"Averaging {len(indices)} images from source, indices: {indices}" + ) im = self.averager.average( - self.class_indices[_indices], self.class_refl[_indices] + self.class_indices[indices], self.class_refl[indices] ) # Finally, apply transforms to resulting Images From 67dcf40ca3f51e89b27a62cc00f7442010157f54 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 1 May 2025 16:01:41 -0400 Subject: [PATCH 150/216] copy syntax --- src/aspire/classification/averager2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index b8cf835ccf..baaf0edf36 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -944,7 +944,7 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) - _images[:] = xp.asarray(original_images[1:].asnumpy(), copy=True) + _images[:] = xp.asarray(original_images[1:].asnumpy().copy()) # Handle reflections refl = reflections[k][1:] # skips original_image 0 From 1ce1b1a5aaca0f8e185ebdfde96bd60fffbfef0e Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 8 May 2025 12:41:16 -0400 Subject: [PATCH 151/216] remove interp option from polar cross cor align --- src/aspire/classification/averager2d.py | 66 ++----------------------- 1 file changed, 4 insertions(+), 62 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index baaf0edf36..49f77a6fb3 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -2,8 +2,6 @@ from abc import ABC, abstractmethod import numpy as np -from scipy.interpolate import interp1d -from scipy.optimize import minimize_scalar from aspire.basis import Coef from aspire.classification.reddy_chatterji import reddy_chatterji_register @@ -814,7 +812,7 @@ def __init__( ) self._mask = xp.asarray(grid_2d(self.src.L, normalized=True)["r"] < 1) - def _fast_rotational_alignment(self, pfA, pfB, do_interp=False): + def _fast_rotational_alignment(self, pfA, pfB): """ Perform fast rotational alignment using Polar Fourier cross correlation. @@ -833,71 +831,15 @@ def _fast_rotational_alignment(self, pfA, pfB, do_interp=False): # 2 hats one sum pfA = fft.fft(pfA, axis=-2) pfB = fft.fft(pfB, axis=-2) - # x = pfA * pfB.conj() + # Tabulate elements of pfA cross pfB.conj() using broadcast multiply x = xp.expand_dims(pfA, 1) * xp.expand_dims(pfB.conj(), 0) angular = xp.sum(xp.abs(fft.ifft2(x)), axis=-1) # sum all radial contributions # Resolve the angle maximizing the correlation through the angular dimension inds = xp.argmax(angular, axis=-1) - if do_interp: - half_width = 5 - fine_steps = 100 - thetas = np.linspace(0, 2 * np.pi, self._pft.ntheta, endpoint=False) - shp = (pfA.shape[0], pfB.shape[0]) - max_thetas = np.empty(shp, dtype=self.dtype) - peaks = np.empty(shp, dtype=self.dtype) - - for i in range(inds.shape[0]): - for j in range(inds.shape[1]): - ind = inds[i, j] - - # Select windows around peak - # Want slice, [ind-half_width:ind+half_width], with wrapping - # Note, could alternatively use halfwidth "pad" with wrap - window = range(ind - half_width, ind + half_width) - xw = thetas.take(window, mode="wrap") - mask = xw < xw[0] - xw[mask] = xw[mask] + 2 * np.pi - yw = angular[i, j].take(window, mode="wrap") - - # Setup an interpolator for the window - f_interp = interp1d(xw, yw, kind="cubic") - - if do_interp == "opt": - # Negate the function we want to maximize - def f(x, _f=f_interp): - return -_f(x) - - # Call the optimizer - res = minimize_scalar(f, bounds=(xw[0], xw[-1])) - - # Assign results - max_thetas[i, j] = res.x - peaks[i, j] = f_interp(res.x) - - else: - # Create fine grid window - xfine = np.linspace(xw[0], xw[-1], fine_steps) - yfine = f_interp(xfine) - - # Find the maximal value in the fine grid window - indfine = xp.argmax(yfine) - - # Assign results - max_thetas[i, j] = xfine[indfine] - peaks[i, j] = yfine[indfine] - - # Modulate the interpolants wraping around the circle. - max_thetas = max_thetas % (2 * np.pi) - - else: - max_thetas = 2 * np.pi / self._pft.ntheta * inds - peaks = xp.take_along_axis(angular, inds[..., None], axis=-1).squeeze(-1) - - # sanity check, can mv to unit test later - assert max_thetas.shape == peaks.shape - assert max_thetas.shape == (pfA.shape[0], pfB.shape[0]) + max_thetas = 2 * np.pi / self._pft.ntheta * inds + peaks = xp.take_along_axis(angular, inds[..., None], axis=-1).squeeze(-1) return xp.asnumpy(max_thetas), xp.asnumpy(peaks) From 889543f70f9673759513d945a4f4e298435d6416 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 13 May 2025 09:03:37 -0400 Subject: [PATCH 152/216] cleanup comment --- src/aspire/classification/averager2d.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 49f77a6fb3..ce11407486 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -929,7 +929,9 @@ def align(self, classes, reflections, basis_coefficients=None): self._fast_rotational_alignment(pf_template_images[:bs], pf_images) ) - # vectorize these + # Note, these could be vectorized, but the code block + # wasn't appreciably faster when I compared them for + # current problem sizes. for i in range(bs): # Test and update From d3d111fcf6d2a4020e524fa59998046ca92348ee Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 May 2025 15:52:14 -0400 Subject: [PATCH 153/216] update //16 to //32 in shift search --- src/aspire/classification/averager2d.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index ce11407486..a34829bba1 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -304,7 +304,7 @@ def __init__( :params n_angles: Number of brute force rotations to attempt, defaults 360. :param radius: Brute force translation search radius. - Defaults to src.L//16. + Defaults to src.L//32. """ super().__init__( composite_basis, @@ -321,7 +321,7 @@ def __init__( f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `rotate` method." ) - self.radius = radius if radius is not None else src.L // 16 + self.radius = radius if radius is not None else src.L // 32 if self.radius != 0: @@ -779,7 +779,7 @@ def __init__( :param n_radial: Number of PFT radial components, defaults `self.src.L`. :param radius: Brute force translation search radius. `0` disables translation search, rotations only. - Defaults to `src.L//16`. + Defaults to `src.L//32`. :param sub_pixel: Subpixel shift size used in brute force shift search. """ super().__init__( From 24c5c8b27977e4db19023de17a62f4a9dfa91252 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 May 2025 15:57:18 -0400 Subject: [PATCH 154/216] default to self.n_radial --- src/aspire/classification/averager2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index a34829bba1..1e16fef9ce 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -808,7 +808,7 @@ def __init__( # Setup Polar Transform self._pft = PolarFT( - self.src.L, ntheta=n_angles, nrad=n_radial, dtype=self.dtype + self.src.L, ntheta=n_angles, nrad=self.n_radial, dtype=self.dtype ) self._mask = xp.asarray(grid_2d(self.src.L, normalized=True)["r"] < 1) From 3bd578b4e048190eb4754bb28fd7d8123fc00567 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 May 2025 15:57:29 -0400 Subject: [PATCH 155/216] typo ceates -> creates --- src/aspire/denoising/class_avg.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py index 7c362fd0a4..ea8cb99342 100644 --- a/src/aspire/denoising/class_avg.py +++ b/src/aspire/denoising/class_avg.py @@ -418,7 +418,7 @@ def __init__( :param class_selector: `ClassSelector` instance. Default `None` creates `TopClassSelector`. :param averager: `Averager2D` instance. - Default `None` ceates `BFRAverager2D` instance. + Default `None` creates `BFRAverager2D` instance. See code for parameter details. :param batch_size: Integer size for batched operations. @@ -489,7 +489,7 @@ def __init__( :param class_selector: `ClassSelector` instance. Default `None` creates `GlobalVarianceClassSelector`. :param averager: `Averager2D` instance. - Default `None` ceates `BFTAverager2D` instance. + Default `None` creates `BFTAverager2D` instance. See code for parameter details. :param averager_src: Optionally explicitly assign source to `BFRAverager2D` during initialization. Allows users to @@ -632,7 +632,7 @@ def __init__( :param class_selector: `ClassSelector` instance. Default `None` creates `GlobalVarianceClassSelector`. :param averager: `Averager2D` instance. - Default `None` ceates `BFTAverager2D` instance. + Default `None` creates `BFTAverager2D` instance. See code for parameter details. :param averager_src: Optionally explicitly assign source to `averager` during initialization. Allows users to @@ -717,7 +717,7 @@ def __init__( `BandedSNRImageQualityFunction`. This will select the images with the highest banded SNR. :param averager: `Averager2D` instance. - Default `None` ceates `BFRAverager2D` instance. + Default `None` creates `BFRAverager2D` instance. See code for parameter details. :param averager_src: Optionally explicitly assign source to `averager` during initialization. Allows users to From 08a19cc7e60520dd65f2231a52fd69264a2160b2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 15 May 2025 16:05:03 -0400 Subject: [PATCH 156/216] docstring updates --- src/aspire/denoising/class_avg.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py index ea8cb99342..01489556f1 100644 --- a/src/aspire/denoising/class_avg.py +++ b/src/aspire/denoising/class_avg.py @@ -461,8 +461,11 @@ class LegacyClassAvgSource(ClassAvgSource): """ Source for denoised 2D images using class average methods. - Defaults to using global variance based class selection, - and a brute force image alignment (rotational only). + Defaults to using global variance based class selection, and a + rotational image alignment. Translational alignment is skipped by + default (images are assumed reasonably centered), but can be + configured by supplying a custom `averager=BFTAverager2D(...)` + argument. This is similar to what was reported for papers using the MATLAB code. @@ -492,7 +495,7 @@ def __init__( Default `None` creates `BFTAverager2D` instance. See code for parameter details. :param averager_src: Optionally explicitly assign source to - `BFRAverager2D` during initialization. Allows users to + `averager` during initialization. Allows users to provide distinct sources for classification and averaging. Raises error when combined with an explicit `averager` argument. From 9c5a9dc115108766134bbfa9f2ec276cbdb1900b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 16 May 2025 12:14:00 -0400 Subject: [PATCH 157/216] use L//2 for n_radial --- src/aspire/classification/averager2d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 1e16fef9ce..aaef96d013 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -776,7 +776,7 @@ def __init__( See AligningAverager2D. Adds `n_angles`, `n_radial`, `radius`, `sub_pixel`. :params n_angles: Number of PFT angular components, defaults 360. - :param n_radial: Number of PFT radial components, defaults `self.src.L`. + :param n_radial: Number of PFT radial components, defaults `self.src.L//2`. :param radius: Brute force translation search radius. `0` disables translation search, rotations only. Defaults to `src.L//32`. @@ -804,7 +804,7 @@ def __init__( self.sub_pixel = sub_pixel # Configure number of radial points - self.n_radial = n_radial or self.src.L + self.n_radial = n_radial or self.src.L // 2 # Setup Polar Transform self._pft = PolarFT( From 4c6fedd2d2bfc17c0444f1e3bbde255f1916b749 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 21 May 2025 09:01:51 -0400 Subject: [PATCH 158/216] len(shifts) ~> len(test_shifts) --- src/aspire/classification/averager2d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index aaef96d013..300ef0983c 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -390,9 +390,9 @@ def align(self, classes, reflections, basis_coefficients=None): # Loop over shift search space, updating best result for x, y in tqdm( test_shifts, - total=len(shifts), + total=len(test_shifts), desc="\tmaximizing over shifts", - disable=len(shifts) == 1, + disable=len(test_shifts) == 1, leave=False, ): shift = np.array([x, y], dtype=int) From a1f20431ea4b358c0035db86b5c110aecc4a40a3 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 21 May 2025 10:18:15 -0400 Subject: [PATCH 159/216] cleanup minor review remarks --- src/aspire/classification/averager2d.py | 31 +++++++++++++------------ 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 300ef0983c..9138009d96 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -257,13 +257,15 @@ def _shift_search_grid(self, L, radius, roll_zero=False, sub_pixel=1): :param radius: Disc radius in pixels :param roll_zero: Roll (0,0) to zero'th element. Defaults to False. - :param sub_pixel: Sub pixel decimation. 1 is integer, 0.1 is 1/10 pixel, etc. + :param sub_pixel: Sub-pixel decimation . 1 yields 1 pixel, 10 yields 1/10 pixel, etc. + Values will be cast to integers. :returns: Grid points as array of 2-tuples [(x0,y0),... (xi,yi)]. """ + sub_pixel = int(sub_pixel) # We'll brute force all shifts in a grid. - g = grid_2d(1 / sub_pixel * L, normalized=False) - disc = g["r"] <= 1 / sub_pixel * radius + g = grid_2d(sub_pixel * L, normalized=False) + disc = g["r"] <= (sub_pixel * radius) X, Y = g["x"][disc], g["y"][disc] X, Y = X * sub_pixel, Y * sub_pixel @@ -768,7 +770,7 @@ def __init__( n_angles=360, n_radial=None, radius=None, - sub_pixel=0.1, + sub_pixel=10, batch_size=512, dtype=None, ): @@ -780,7 +782,8 @@ def __init__( :param radius: Brute force translation search radius. `0` disables translation search, rotations only. Defaults to `src.L//32`. - :param sub_pixel: Subpixel shift size used in brute force shift search. + :param sub_pixel: Sub-pixel decimation used in brute force shift search. + Defaults to 10 sub-pixel to pixel, ie 0.1 spaced sub-pixel. """ super().__init__( composite_basis, @@ -794,12 +797,10 @@ def __init__( self.radius = radius if radius is not None else src.L // 32 - if self.radius != 0: - - if not hasattr(self.alignment_basis, "shift"): - raise RuntimeError( - f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method." - ) + if self.radius != 0 and not hasattr(self.alignment_basis, "shift"): + raise RuntimeError( + f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method." + ) self.sub_pixel = sub_pixel @@ -817,8 +818,8 @@ def _fast_rotational_alignment(self, pfA, pfB): Perform fast rotational alignment using Polar Fourier cross correlation. Note broadcasting is specialized for this problem. - pfA.shape (m, nt, nr) - pfB.shape (n, nt, nr) + pfA.shape (m, ntheta, nrad) + pfB.shape (n, ntheta, nrad) yields thetas (m,n), peaks (m,n) """ @@ -923,8 +924,8 @@ def align(self, classes, reflections, basis_coefficients=None): self._pft._transform(template_images) ) - # # Compute and assign the best rotation found with this translation - # # note offset of 1 for skipped original_image 0 + # Compute and assign the best rotation found with this translation + # note offset of 1 for skipped original_image 0 _rotations[:bs, 1:], _dot_products[:bs, 1:] = ( self._fast_rotational_alignment(pf_template_images[:bs], pf_images) ) From ac0c07318927dc955d851a2fa52afbb16781674b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Fri, 23 May 2025 12:11:53 -0400 Subject: [PATCH 160/216] sub pixel review change bug --- src/aspire/classification/averager2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 9138009d96..0b442224c9 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -267,7 +267,7 @@ def _shift_search_grid(self, L, radius, roll_zero=False, sub_pixel=1): g = grid_2d(sub_pixel * L, normalized=False) disc = g["r"] <= (sub_pixel * radius) X, Y = g["x"][disc], g["y"][disc] - X, Y = X * sub_pixel, Y * sub_pixel + X, Y = X / sub_pixel, Y / sub_pixel # Optionally roll arrays so 0 is first. if roll_zero: From f9ed410c2de3365c9cc63b6f63b523d813e34fa2 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 27 May 2025 15:20:29 -0400 Subject: [PATCH 161/216] stub in PolarFT shifting --- src/aspire/operators/polar_ft.py | 28 +++++++++++++++++++++ tests/test_polar_ft.py | 43 ++++++++++++++++++++++++++------ 2 files changed, 64 insertions(+), 7 deletions(-) diff --git a/src/aspire/operators/polar_ft.py b/src/aspire/operators/polar_ft.py index 1218a22010..80593f8366 100644 --- a/src/aspire/operators/polar_ft.py +++ b/src/aspire/operators/polar_ft.py @@ -157,3 +157,31 @@ def half_to_full(pf): concatenate = np.concatenate return concatenate((pf, pf.conj()), axis=-2) + + def shift(self, pfx, shifts): + """ + Shift `pfx` by `shifts` pixels using `PolarFT`. + + :param pfx: Array of `PolarFT` coefs shaped `(n_img, ntheta//2, nrad)`. + :param shifts: Array of (x,y) shifts shaped `(n_img, 2). + :return: Array of shifted coefs shaped `(n_img, ntheta//2, nrad)`. + """ + + n_img = pfx.shape[0] + + # Handle a single shift + shifts = np.atleast_2d(shifts) + + # Flip shift XY axis?! + shifts = shifts[..., ::-1] + + # Broadcast and accumulate phase shifts + freqs = np.tile(self.freqs, (n_img, 1, 1)) + phase_shifts = np.exp(-1j * np.sum(freqs * -shifts[:, :, None], axis=1)) + + # Reshape flat frequency grid back to (..., ntheta//2, self.nrad) + phase_shifts = phase_shifts.reshape(n_img, self.ntheta // 2, self.nrad) + # Apply the phase shifts elementwise + shifted_pfx = phase_shifts * pfx + + return shifted_pfx diff --git a/tests/test_polar_ft.py b/tests/test_polar_ft.py index 425d5e14bb..6f1b7d5bc9 100644 --- a/tests/test_polar_ft.py +++ b/tests/test_polar_ft.py @@ -62,7 +62,7 @@ def gaussian(img_size, dtype): gauss = Image( gaussian_2d(img_size, sigma=(img_size // 10, img_size // 10), dtype=dtype) ) - pf = pf_transform(gauss) + pf = pf_transform(gauss)[0] return pf @@ -74,7 +74,7 @@ def symmetric_image(img_size, dtype): img_size, C=1, order=4, K=25, seed=10, dtype=dtype ).generate() symmetric_image = symmetric_vol.project(np.eye(3, dtype=dtype)) - pf = pf_transform(symmetric_image) + pf = pf_transform(symmetric_image)[0] return pf @@ -84,16 +84,16 @@ def asymmetric_image(img_size, dtype): """Asymetric image.""" asymmetric_vol = AsymmetricVolume(img_size, C=1, dtype=dtype).generate() asymmetric_image = asymmetric_vol.project(np.eye(3, dtype=dtype)) - pf = pf_transform(asymmetric_image) + pf, pft = pf_transform(asymmetric_image) - return asymmetric_image, pf + return asymmetric_image, pf, pft @pytest.fixture def radial_mode_image(img_size, dtype, radial_mode): g = grid_2d(img_size, dtype=dtype) image = Image(np.sin(radial_mode * np.pi * g["r"])) - pf = pf_transform(image) + pf = pf_transform(image)[0] return pf, radial_mode @@ -107,7 +107,7 @@ def pf_transform(image): pft = PolarFT(img_size, nrad=nrad, ntheta=ntheta, dtype=image.dtype) pf = pft.transform(image)[0] - return pf + return pf, pft # ============= @@ -117,7 +117,7 @@ def pf_transform(image): def test_dc_component(asymmetric_image): """Test that the DC component equals the mean of the signal.""" - image, pf = asymmetric_image + image, pf, _ = asymmetric_image signal_mean = np.mean(image) dc_components = abs(pf[:, 0]) @@ -220,3 +220,32 @@ def test_half_to_full_transform(stack_shape): np.testing.assert_allclose( full_pf[..., ray, :], np.conj(full_pf[..., ray + pft.ntheta // 2, :]) ) + + +def test_shift(asymmetric_image): + """ + Compare shifting using PolarFT.shift against Image.shift. + """ + + # Test image, PolarFT coef, and PolarFT instance. + img, pf_coef, pft = asymmetric_image + # For some reason the utils in this file strip off the stack axis. + # put it back so it matches what the `transform` function actually returns. + pf_coef = pf_coef[None] + + # Test shift + shift = np.array([[3, 5]], dtype=img.dtype) + + # Shift using `PolarFT` class + pf_shifted_coef = pft.shift(pf_coef, shift) + + # Shift using `Image` class + img_shifted = img.shift(shift) + # then transform to PolarFT coef + img_shifted_coef = pft.transform(img_shifted) + + # Compare resulting coefs, look for <1% error (loose). + err = np.linalg.norm(pf_shifted_coef - img_shifted_coef) + norm = np.linalg.norm(img_shifted_coef) + percent_error = err / norm * 100 + np.testing.assert_array_less(percent_error, 1, err_msg="Shifting error too high.") From cbb6bf316ace1ad38b2ed1912357c0d87f62b853 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 28 May 2025 09:59:56 -0400 Subject: [PATCH 162/216] stub in PolarFT shift test 2d --- tests/test_polar_ft.py | 35 ++++++++++++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/tests/test_polar_ft.py b/tests/test_polar_ft.py index 6f1b7d5bc9..7107526da0 100644 --- a/tests/test_polar_ft.py +++ b/tests/test_polar_ft.py @@ -222,12 +222,12 @@ def test_half_to_full_transform(stack_shape): ) -def test_shift(asymmetric_image): +def test_shift_1d(asymmetric_image): """ Compare shifting using PolarFT.shift against Image.shift. """ - # Test image, PolarFT coef, and PolarFT instance. + # Test image, `PolarFT` coef, and `PolarFT` instance. img, pf_coef, pft = asymmetric_image # For some reason the utils in this file strip off the stack axis. # put it back so it matches what the `transform` function actually returns. @@ -241,7 +241,7 @@ def test_shift(asymmetric_image): # Shift using `Image` class img_shifted = img.shift(shift) - # then transform to PolarFT coef + # then transform to `PolarFT` coef img_shifted_coef = pft.transform(img_shifted) # Compare resulting coefs, look for <1% error (loose). @@ -249,3 +249,32 @@ def test_shift(asymmetric_image): norm = np.linalg.norm(img_shifted_coef) percent_error = err / norm * 100 np.testing.assert_array_less(percent_error, 1, err_msg="Shifting error too high.") + + +def test_shift_2d(asymmetric_image): + """ + Compare shifting using PolarFT.shift against Image.shift. + """ + + # Test image, `PolarFT` coef, and `PolarFT` instance. + img, pf_coef, pft = asymmetric_image + # For some reason the utils in this file strip off the stack axis. + # put it back so it matches what the `transform` function actually returns. + pf_coef = np.tile(pf_coef, (3, 1, 1)) + + # Test shift + shift = np.array([[3, 5], [-2, -1], [4, -3]], dtype=img.dtype) + + # Shift using `PolarFT` class + pf_shifted_coef = pft.shift(pf_coef, shift) + + # Shift using `Image` class + img_shifted = img.shift(shift) + # then transform to `PolarFT` coef + img_shifted_coef = pft.transform(img_shifted) + + # Compare resulting coefs, look for <1% error (loose). + err = np.linalg.norm(pf_shifted_coef - img_shifted_coef, axis=(1, 2)) + norm = np.linalg.norm(img_shifted_coef, axis=(1, 2)) + percent_error = err / norm * 100 + np.testing.assert_array_less(percent_error, 1, err_msg="Shifting error too high.") From 9a4b52aeba42ea18639888ad178c83f350f54bb4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 28 May 2025 15:15:52 -0400 Subject: [PATCH 163/216] add broadcast polar shift test --- tests/test_polar_ft.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/tests/test_polar_ft.py b/tests/test_polar_ft.py index 7107526da0..93bba641b6 100644 --- a/tests/test_polar_ft.py +++ b/tests/test_polar_ft.py @@ -256,11 +256,41 @@ def test_shift_2d(asymmetric_image): Compare shifting using PolarFT.shift against Image.shift. """ + # Test image, `PolarFT` coef, and `PolarFT` instance. + img, pf_coef, pft = asymmetric_image + # Tile to a stack of 3 images + pf_coef = np.tile(pf_coef, (3, 1, 1)) + + # Test shift + shift = np.array([[3, 5], [-2, -1], [4, -3]], dtype=img.dtype) + + # Shift using `PolarFT` class + pf_shifted_coef = pft.shift(pf_coef, shift) + + # Shift using `Image` class + img_shifted = img.shift(shift) + # then transform to `PolarFT` coef + img_shifted_coef = pft.transform(img_shifted) + + # Compare resulting coefs, look for <1% error (loose). + err = np.linalg.norm(pf_shifted_coef - img_shifted_coef, axis=(1, 2)) + norm = np.linalg.norm(img_shifted_coef, axis=(1, 2)) + percent_error = err / norm * 100 + np.testing.assert_array_less(percent_error, 1, err_msg="Shifting error too high.") + + +def test_shift_broadcast(asymmetric_image): + """ + Compare shifting using PolarFT.shift against Image.shift. + + Shifts single image with multiple shifts. + """ + # Test image, `PolarFT` coef, and `PolarFT` instance. img, pf_coef, pft = asymmetric_image # For some reason the utils in this file strip off the stack axis. # put it back so it matches what the `transform` function actually returns. - pf_coef = np.tile(pf_coef, (3, 1, 1)) + pf_coef = pf_coef[None] # Test shift shift = np.array([[3, 5], [-2, -1], [4, -3]], dtype=img.dtype) @@ -277,4 +307,5 @@ def test_shift_2d(asymmetric_image): err = np.linalg.norm(pf_shifted_coef - img_shifted_coef, axis=(1, 2)) norm = np.linalg.norm(img_shifted_coef, axis=(1, 2)) percent_error = err / norm * 100 + np.testing.assert_array_less(percent_error, 1, err_msg="Shifting error too high.") From f147b251fb38fe61a8da5f2470969778189b4b4a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 28 May 2025 15:20:40 -0400 Subject: [PATCH 164/216] add multiple shift broadcast polar code --- src/aspire/operators/polar_ft.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/aspire/operators/polar_ft.py b/src/aspire/operators/polar_ft.py index 80593f8366..a0f577ee10 100644 --- a/src/aspire/operators/polar_ft.py +++ b/src/aspire/operators/polar_ft.py @@ -167,20 +167,31 @@ def shift(self, pfx, shifts): :return: Array of shifted coefs shaped `(n_img, ntheta//2, nrad)`. """ + # Number of input images n_img = pfx.shape[0] # Handle a single shift shifts = np.atleast_2d(shifts) + n_shifts = shifts.shape[0] + + # Handle broadcast case, calculate number of output images `n` + n = n_img + if n_img == 1: + n = n_shifts + elif n_shifts != n_img: + raise ValueError( + f"Incompatible number of images {n_img} and shifts {n_shifts}" + ) # Flip shift XY axis?! shifts = shifts[..., ::-1] # Broadcast and accumulate phase shifts - freqs = np.tile(self.freqs, (n_img, 1, 1)) + freqs = np.tile(self.freqs, (n, 1, 1)) phase_shifts = np.exp(-1j * np.sum(freqs * -shifts[:, :, None], axis=1)) # Reshape flat frequency grid back to (..., ntheta//2, self.nrad) - phase_shifts = phase_shifts.reshape(n_img, self.ntheta // 2, self.nrad) + phase_shifts = phase_shifts.reshape(n, self.ntheta // 2, self.nrad) # Apply the phase shifts elementwise shifted_pfx = phase_shifts * pfx From 87a9f485e15dc01f79cff7f968e06b1732e267d1 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 28 May 2025 15:21:11 -0400 Subject: [PATCH 165/216] Use PolarFT.shift in BFT class source --- src/aspire/classification/averager2d.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py index 0b442224c9..93845573ac 100644 --- a/src/aspire/classification/averager2d.py +++ b/src/aspire/classification/averager2d.py @@ -8,7 +8,7 @@ from aspire.image import Image, ImageStacker, MeanImageStacker from aspire.numeric import fft, xp from aspire.operators import PolarFT -from aspire.utils import tqdm, trange +from aspire.utils import complex_type, tqdm, trange from aspire.utils.coor_trans import grid_2d logger = logging.getLogger(__name__) @@ -873,7 +873,9 @@ def align(self, classes, reflections, basis_coefficients=None): bs = min(self.batch_size, len(test_shifts)) _rotations = np.zeros((bs, n_nbor), dtype=self.dtype) _dot_products = np.ones((bs, n_nbor), dtype=self.dtype) * -np.inf - template_images = xp.empty((bs, self.src.L, self.src.L), dtype=self.dtype) + template_images = xp.empty( + (bs, self._pft.ntheta // 2, self._pft.nrad), dtype=complex_type(self.dtype) + ) _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"): @@ -887,6 +889,7 @@ def align(self, classes, reflections, basis_coefficients=None): original_coef = basis_coefficients[classes[k], :] original_images = self.alignment_basis.evaluate(original_coef) + _img0 = original_images[0].asnumpy().copy() _images[:] = xp.asarray(original_images[1:].asnumpy().copy()) # Handle reflections @@ -897,6 +900,7 @@ def align(self, classes, reflections, basis_coefficients=None): _images[:] = _images[:] * self._mask # Convert to polar Fourier + pf_img0 = self._pft._transform(_img0) pf_images = self._pft.half_to_full(self._pft._transform(_images)) # Batch over shift search space, updating best results @@ -911,18 +915,13 @@ def align(self, classes, reflections, basis_coefficients=None): bs = end - start # handle a small last batch batch_shifts = test_shifts[start:end] - # Shift the base, original_image[0], for each shift in this batch - # Note this includes shifting for the zero case + # Shift the base, pf_img0, for each shift in this batch + # Note this includes shifting for the zero shift case template_images[:bs] = xp.asarray( - original_images[0].shift(batch_shifts) + self._pft.shift(pf_img0, batch_shifts) ) - # Mask off - template_images[:] = template_images[:] * self._mask - - pf_template_images = self._pft.half_to_full( - self._pft._transform(template_images) - ) + pf_template_images = self._pft.half_to_full(template_images) # Compute and assign the best rotation found with this translation # note offset of 1 for skipped original_image 0 From 8ad361b6514b0e2612d92e36324f12c13425c2b8 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 28 May 2025 15:32:12 -0400 Subject: [PATCH 166/216] exted PolarFT.shift to xp --- src/aspire/operators/polar_ft.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/src/aspire/operators/polar_ft.py b/src/aspire/operators/polar_ft.py index a0f577ee10..de57297700 100644 --- a/src/aspire/operators/polar_ft.py +++ b/src/aspire/operators/polar_ft.py @@ -167,11 +167,16 @@ def shift(self, pfx, shifts): :return: Array of shifted coefs shaped `(n_img, ntheta//2, nrad)`. """ + # Convert to xp array as needed + input_on_host = isinstance(pfx, np.ndarray) + pfx = xp.asarray(pfx) + shifts = xp.asarray(shifts) + # Number of input images n_img = pfx.shape[0] # Handle a single shift - shifts = np.atleast_2d(shifts) + shifts = xp.atleast_2d(shifts) n_shifts = shifts.shape[0] # Handle broadcast case, calculate number of output images `n` @@ -187,12 +192,16 @@ def shift(self, pfx, shifts): shifts = shifts[..., ::-1] # Broadcast and accumulate phase shifts - freqs = np.tile(self.freqs, (n, 1, 1)) - phase_shifts = np.exp(-1j * np.sum(freqs * -shifts[:, :, None], axis=1)) + freqs = xp.tile(xp.asarray(self.freqs), (n, 1, 1)) + phase_shifts = xp.exp(-1j * xp.sum(freqs * -shifts[:, :, None], axis=1)) # Reshape flat frequency grid back to (..., ntheta//2, self.nrad) phase_shifts = phase_shifts.reshape(n, self.ntheta // 2, self.nrad) # Apply the phase shifts elementwise shifted_pfx = phase_shifts * pfx + # If we started on host, return as host array. + if input_on_host: + shifted_pfx = xp.asnumpy(shifted_pfx) + return shifted_pfx From 83ac3699f52c6109b887b5384bdd85f1c4b924a3 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 25 Jun 2025 11:01:21 -0400 Subject: [PATCH 167/216] update pinned cu/finufft to 2.4.0 --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 66950c1770..a0e59d8dfa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ "click", "confuse >= 2.0.0", "cvxpy", - "finufft==2.3.0", + "finufft==2.4.0", "gemmi >= 0.6.5", "joblib", "matplotlib >= 3.2.0", @@ -56,7 +56,7 @@ dependencies = [ "Source" = "https://github.com/ComputationalCryoEM/ASPIRE-Python" [project.optional-dependencies] -gpu-12x = ["cupy-cuda12x", "cufinufft==2.3.0"] +gpu-12x = ["cupy-cuda12x", "cufinufft==2.4.0"] dev = [ "black", "bumpversion", From e04c0ead98576c40bc4dd1023fd3a4e0447c1009 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 26 Jun 2025 11:06:29 -0400 Subject: [PATCH 168/216] let's try updating the conda envs along with finufft --- environment-accelerate.yml | 4 ++-- environment-default.yml | 4 ++-- environment-intel.yml | 4 ++-- environment-openblas.yml | 4 ++-- environment-win64.yml | 4 ++-- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/environment-accelerate.yml b/environment-accelerate.yml index 5248c3a953..aaed932b12 100644 --- a/environment-accelerate.yml +++ b/environment-accelerate.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=1.24.1 - - scipy=1.10.1 + - numpy=2.0.2 + - scipy=1.13.1 - scikit-learn - scikit-image - libblas=*=*accelerate diff --git a/environment-default.yml b/environment-default.yml index f827ec55bc..0b79d4333b 100644 --- a/environment-default.yml +++ b/environment-default.yml @@ -7,7 +7,7 @@ channels: dependencies: - pip - python=3.9 - - numpy=1.23.5 - - scipy=1.9.3 + - numpy=2.0.2 + - scipy=1.13.1 - scikit-learn - scikit-image diff --git a/environment-intel.yml b/environment-intel.yml index 840dd92f76..68141d0a11 100644 --- a/environment-intel.yml +++ b/environment-intel.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=1.23.5 - - scipy=1.9.3 + - numpy=2.0.2 + - scipy=1.13.1 - scikit-learn - scikit-image - mkl_fft diff --git a/environment-openblas.yml b/environment-openblas.yml index 088035f88d..813264aa5c 100644 --- a/environment-openblas.yml +++ b/environment-openblas.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=1.23.5 - - scipy=1.9.3 + - numpy=2.0.2 + - scipy=1.13.1 - scikit-learn - scikit-image - libblas=*=*openblas diff --git a/environment-win64.yml b/environment-win64.yml index 34ca5d9fa6..ad74d1fd99 100644 --- a/environment-win64.yml +++ b/environment-win64.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=1.23.5 - - scipy=1.9.3 + - numpy=2.0.2 + - scipy=1.13.1 - scikit-learn - scikit-image - mkl=2024.1.* # possible regression impacts eig solver in later versions up to 2025.0 From 4c39b515a5fae53260d3b3d4c0548d71a1ace40a Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 30 Jun 2025 08:31:47 -0400 Subject: [PATCH 169/216] revert finufft for darwin --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a0e59d8dfa..25e400dfaa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,8 @@ dependencies = [ "click", "confuse >= 2.0.0", "cvxpy", - "finufft==2.4.0", + "finufft==2.4.0 ; sys_platform!='darwin'", + "finufft==2.3.0 ; sys_platform=='darwin'", "gemmi >= 0.6.5", "joblib", "matplotlib >= 3.2.0", From b16ae0df15e79cc3daf966a1b473b1fb607738e5 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Mon, 30 Jun 2025 15:15:46 -0400 Subject: [PATCH 170/216] revert finufft upsampfac --- src/aspire/nufft/finufft.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/aspire/nufft/finufft.py b/src/aspire/nufft/finufft.py index a955ab20d0..53645b5449 100644 --- a/src/aspire/nufft/finufft.py +++ b/src/aspire/nufft/finufft.py @@ -51,6 +51,7 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs): eps=self.epsilon, n_trans=self.ntransforms, dtype=self.complex_dtype, + upsampfac=2, # revert <2.4.0 default ) self._adjoint_plan = finufft.Plan( @@ -59,6 +60,7 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs): eps=self.epsilon, n_trans=self.ntransforms, dtype=self.complex_dtype, + upsampfac=2, # revert <2.4.0 default ) self._transform_plan.setpts(*self.fourier_pts) From ceb6cbd6b206a054ef9bf02a5d2708c9e62a2e80 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 1 Jul 2025 07:55:03 -0400 Subject: [PATCH 171/216] Revert env yamls --- environment-accelerate.yml | 4 ++-- environment-default.yml | 4 ++-- environment-intel.yml | 4 ++-- environment-openblas.yml | 4 ++-- environment-win64.yml | 4 ++-- 5 files changed, 10 insertions(+), 10 deletions(-) diff --git a/environment-accelerate.yml b/environment-accelerate.yml index aaed932b12..5248c3a953 100644 --- a/environment-accelerate.yml +++ b/environment-accelerate.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=2.0.2 - - scipy=1.13.1 + - numpy=1.24.1 + - scipy=1.10.1 - scikit-learn - scikit-image - libblas=*=*accelerate diff --git a/environment-default.yml b/environment-default.yml index 0b79d4333b..f827ec55bc 100644 --- a/environment-default.yml +++ b/environment-default.yml @@ -7,7 +7,7 @@ channels: dependencies: - pip - python=3.9 - - numpy=2.0.2 - - scipy=1.13.1 + - numpy=1.23.5 + - scipy=1.9.3 - scikit-learn - scikit-image diff --git a/environment-intel.yml b/environment-intel.yml index 68141d0a11..840dd92f76 100644 --- a/environment-intel.yml +++ b/environment-intel.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=2.0.2 - - scipy=1.13.1 + - numpy=1.23.5 + - scipy=1.9.3 - scikit-learn - scikit-image - mkl_fft diff --git a/environment-openblas.yml b/environment-openblas.yml index 813264aa5c..088035f88d 100644 --- a/environment-openblas.yml +++ b/environment-openblas.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=2.0.2 - - scipy=1.13.1 + - numpy=1.23.5 + - scipy=1.9.3 - scikit-learn - scikit-image - libblas=*=*openblas diff --git a/environment-win64.yml b/environment-win64.yml index ad74d1fd99..34ca5d9fa6 100644 --- a/environment-win64.yml +++ b/environment-win64.yml @@ -7,8 +7,8 @@ channels: dependencies: - pip - python=3.9 - - numpy=2.0.2 - - scipy=1.13.1 + - numpy=1.23.5 + - scipy=1.9.3 - scikit-learn - scikit-image - mkl=2024.1.* # possible regression impacts eig solver in later versions up to 2025.0 From 38783144c9754fec1b240bb70c8a6d7421e9e50b Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 24 Jun 2025 12:13:48 -0400 Subject: [PATCH 172/216] add more pixel_size attribute handling to CoordinateSources --- src/aspire/source/coordinates.py | 79 ++++++++++++++++++++++++++++++-- tests/test_coordinate_source.py | 59 +++++++++++++++++++++--- 2 files changed, 126 insertions(+), 12 deletions(-) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index 045ae5d8ae..28b8ef7ad1 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -1,5 +1,6 @@ import logging import os +import warnings from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Iterable @@ -46,7 +47,17 @@ class CoordinateSource(ImageSource, ABC): This also allows the CoordinateSource to be saved to an `.mrcs` stack. """ - def __init__(self, files, particle_size, max_rows, B, symmetry_group): + def __init__( + self, files, particle_size, max_rows, B, symmetry_group, pixel_size=None + ): + """ + :param files: A list of tuples of the form (path_to_mrc, path_to_coord) + :particle_size: Desired size of cropped particles (will override the size specified in coordinate file) + :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) + :param B: CTF envelope decay factor + :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. + :param pixel_size: Pixel size of the images in angstroms, default `None`. + """ mrc_paths, coord_paths = [f[0] for f in files], [f[1] for f in files] # the particle_size parameter is the *user-specified* argument # and is used in self._populate_particles @@ -134,7 +145,14 @@ def __init__(self, files, particle_size, max_rows, B, symmetry_group): # total particles loaded (specific to this instance) logger.info(f"CoordinateSource object contains {n} particles.") - ImageSource.__init__(self, L=L, n=n, dtype=dtype, symmetry_group=symmetry_group) + ImageSource.__init__( + self, + L=L, + n=n, + dtype=dtype, + symmetry_group=symmetry_group, + pixel_size=pixel_size, + ) # map mrc indices to particle indices # i'th element contains a list of particle indices corresponding to i'th mrc @@ -385,6 +403,32 @@ def _extract_ctf(self, data_block): # convert defocus_ang from degrees to radians filter_params[:, 3] *= np.pi / 180.0 + # Check pixel_size + # Get pixel_sizes from CTFFilters + ctf_pixel_sizes = list(set(filter_params[:, 6])) + # Compare with source.pixel_size if assigned + if (self.pixel_size is not None) and ( + not np.allclose(ctf_pixel_sizes, self.pixel_size) + ): + warnings.warn( + "Pixel size mismatch." + f"\n\tSource: {self.pixel_size}" + f"\n\tCTFs: {ctf_pixel_sizes}.", + stacklevel=2, + ) + # When source is not assigned we can try to assign it from CTF, + # but only do this if all the CTFFilter pixel_sizes are consistent + elif self.pixel_size is None: + if len(ctf_pixel_sizes) == 1: + self.pixel_size = ctf_pixel_sizes[0] # take the unique single element + logger.info( + f"Assigning source pixel_size={self.pixel_size} from CTFFilters." + ) + elif len(ctf_pixel_sizes) > 1: + logger.warning( + "Unable to assign source pixel_size from CTFFilters, multiple pixel_sizes found." + ) + # construct filters self.unique_filters = [ CTFFilter( @@ -516,16 +560,25 @@ def __init__( max_rows=None, B=0, symmetry_group=None, + pixel_size=None, ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) :particle_size: Desired size of cropped particles (will override the size specified in coordinate file) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) + :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. + :param pixel_size: Pixel size of the images in angstroms, default `None`. """ # instantiate super CoordinateSource.__init__( - self, files, particle_size, max_rows, B, symmetry_group + self, + files, + particle_size, + max_rows, + B, + symmetry_group, + pixel_size=pixel_size, ) def _extract_box_size(self, box_file): @@ -629,17 +682,33 @@ class CentersCoordinateSource(CoordinateSource): Represents a data source consisting of micrographs and coordinate files specifying particle centers only. Files can be text (.coord) or STAR files. """ - def __init__(self, files, particle_size, max_rows=None, B=0, symmetry_group=None): + def __init__( + self, + files, + particle_size, + max_rows=None, + B=0, + symmetry_group=None, + pixel_size=None, + ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) :particle_size: Desired size of cropped particles (mandatory) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) + :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. + :param pixel_size: Pixel size of the images in angstroms, default `None`. """ # instantiate super CoordinateSource.__init__( - self, files, particle_size, max_rows, B, symmetry_group + self, + files, + particle_size, + max_rows, + B, + symmetry_group, + pixel_size=pixel_size, ) def _validate_centers_file(self, coord_file): diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index 5304ef0b74..b864e592d0 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -9,6 +9,7 @@ import mrcfile import numpy as np +import pytest from click.testing import CliRunner import tests.saved_test_data @@ -31,6 +32,7 @@ def setUp(self): self.original_mrc_path = str(test_path) # save test data root dir self.test_dir_root = os.path.dirname(self.original_mrc_path) + self.pixel_size = 1.23 # Used for generating and comparing metadata # We will construct a source with two micrographs and two coordinate # files by using the same micrograph, but dividing the coordinates @@ -261,7 +263,7 @@ def createTestCtfFiles(self, index): "_rlnSphericalAberration": 700 + index, "_rlnAmplitudeContrast": 600 + index, "_rlnVoltage": 500 + index, - "_rlnMicrographPixelSize": 400 + index, + "_rlnMicrographPixelSize": self.pixel_size, } blocks = OrderedDict({"root": params_dict}) starfile = StarFile(blocks=blocks) @@ -284,8 +286,8 @@ def createTestRelionCtfFile(self, reverse_optics_block_rows=False): ] # using same unique values as in createTestCtfFiles optics_block = [ - ["opticsGroup1", 1, 500.0, 700.0, 600.0, 400.0], - ["opticsGroup2", 2, 501.0, 701.0, 601.0, 401.0], + ["opticsGroup1", 1, 500.0, 700.0, 600.0, self.pixel_size], + ["opticsGroup2", 2, 501.0, 701.0, 601.0, self.pixel_size], ] # Since optics block rows are self-contained, # reversing their order should have no affect anywhere. @@ -565,7 +567,15 @@ def _testCtfFilters(self, src): self.assertTrue( np.allclose( np.array( - [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, 400.0], + [ + 1000.0, + 900.0, + 800.0 * np.pi / 180.0, + 700.0, + 600.0, + 500.0, + self.pixel_size, + ], dtype=src.dtype, ), np.array( @@ -585,7 +595,15 @@ def _testCtfFilters(self, src): self.assertTrue( np.allclose( np.array( - [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, 401.0], + [ + 1001.0, + 901.0, + 801.0 * np.pi / 180.0, + 701.0, + 601.0, + 501.0, + self.pixel_size, + ], dtype=src.dtype, ), np.array( @@ -644,10 +662,10 @@ def _testCtfMetadata(self, src): ] ctf_metadata = np.zeros((src.n, len(ctf_cols)), dtype=src.dtype) ctf_metadata[:200] = np.array( - [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, 400.0] + [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, self.pixel_size] ) ctf_metadata[200:400] = np.array( - [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, 401.0] + [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, self.pixel_size] ) self.assertTrue(np.array_equal(ctf_metadata, src.get_metadata(ctf_cols))) @@ -700,6 +718,33 @@ def testCommand(self): self.assertTrue(result_star.exit_code == 0) self.assertTrue(result_preprocess.exit_code == 0) + def testPixelSizeWarning(self): + """ + Test source haveing a pixel size that conflicts with the CTFFilter instances. + """ + manual_pixel_size = 0.789 + src = BoxesCoordinateSource(self.files_box, pixel_size=manual_pixel_size) + # Capture and compare warning message + with pytest.warns(UserWarning, match=r".*Pixel size mismatch.*"): + src.import_relion_ctf(self.relion_ctf_file) + assert src.pixel_size == manual_pixel_size + + def testPixelSize(self): + """ + Test explicitly providing correct pixel_size. + """ + src = BoxesCoordinateSource(self.files_box, pixel_size=self.pixel_size) + src.import_relion_ctf(self.relion_ctf_file) + assert src.pixel_size == self.pixel_size + + def testPixelSizeNone(self): + """ + Test not providing pixel_size. + """ + src = BoxesCoordinateSource(self.files_box) + src.import_relion_ctf(self.relion_ctf_file) + assert src.pixel_size == self.pixel_size + def create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=(2.0, 2.0, 1.0)): # Create a rectangular micrograph (e.g., 128x256) From ba325f41e0c0aea298c1edca800042cfcc450335 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 25 Jun 2025 09:48:15 -0400 Subject: [PATCH 173/216] update utest to approx_equal --- tests/test_coordinate_source.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index b864e592d0..a39c404d24 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -727,7 +727,7 @@ def testPixelSizeWarning(self): # Capture and compare warning message with pytest.warns(UserWarning, match=r".*Pixel size mismatch.*"): src.import_relion_ctf(self.relion_ctf_file) - assert src.pixel_size == manual_pixel_size + np.testing.assert_approx_equal(src.pixel_size, manual_pixel_size) def testPixelSize(self): """ @@ -735,7 +735,7 @@ def testPixelSize(self): """ src = BoxesCoordinateSource(self.files_box, pixel_size=self.pixel_size) src.import_relion_ctf(self.relion_ctf_file) - assert src.pixel_size == self.pixel_size + np.testing.assert_approx_equal(src.pixel_size, self.pixel_size) def testPixelSizeNone(self): """ @@ -743,7 +743,7 @@ def testPixelSizeNone(self): """ src = BoxesCoordinateSource(self.files_box) src.import_relion_ctf(self.relion_ctf_file) - assert src.pixel_size == self.pixel_size + np.testing.assert_approx_equal(src.pixel_size, self.pixel_size) def create_test_rectangular_micrograph_and_star(tmp_path, voxel_size=(2.0, 2.0, 1.0)): From e61c015a8ec74735832d2fa77c07673813e2484e Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 25 Jun 2025 09:48:25 -0400 Subject: [PATCH 174/216] cleanup strings --- src/aspire/source/coordinates.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index 28b8ef7ad1..8dee8d5d45 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -52,7 +52,7 @@ def __init__( ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) - :particle_size: Desired size of cropped particles (will override the size specified in coordinate file) + :param particle_size: Desired size of cropped particles (will override the size specified in coordinate file) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. @@ -417,13 +417,14 @@ def _extract_ctf(self, data_block): stacklevel=2, ) # When source is not assigned we can try to assign it from CTF, - # but only do this if all the CTFFilter pixel_sizes are consistent elif self.pixel_size is None: + # but only do this if all the CTFFilter pixel_sizes are consistent if len(ctf_pixel_sizes) == 1: self.pixel_size = ctf_pixel_sizes[0] # take the unique single element logger.info( f"Assigning source pixel_size={self.pixel_size} from CTFFilters." ) + # otherwise let the user know elif len(ctf_pixel_sizes) > 1: logger.warning( "Unable to assign source pixel_size from CTFFilters, multiple pixel_sizes found." @@ -564,7 +565,7 @@ def __init__( ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) - :particle_size: Desired size of cropped particles (will override the size specified in coordinate file) + :param particle_size: Desired size of cropped particles (will override the size specified in coordinate file) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. @@ -693,7 +694,7 @@ def __init__( ): """ :param files: A list of tuples of the form (path_to_mrc, path_to_coord) - :particle_size: Desired size of cropped particles (mandatory) + :param particle_size: Desired size of cropped particles (mandatory) :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) :param B: CTF envelope decay factor From 567a38b2567a64a54affc39e6466d2e0856582db Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Wed, 25 Jun 2025 10:13:20 -0400 Subject: [PATCH 175/216] add multiple pixel size test --- tests/test_coordinate_source.py | 41 +++++++++++++++++++++++++++------ 1 file changed, 34 insertions(+), 7 deletions(-) diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index a39c404d24..eb45140ad2 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -138,6 +138,12 @@ def setUp(self): def tearDown(self): self.tmpdir.cleanup() + # This is a workaround to use a `pytest` fixture with `unittest` style cases. + # We use it below to capture and inspect the log + @pytest.fixture(autouse=True) + def inject_fixtures(self, caplog): + self._caplog = caplog + def createTestBoxFiles(self, centers, index): """ Create a .box file storing particle coordinates as @@ -252,6 +258,8 @@ def createFloatStarFile(self, centers): def createTestCtfFiles(self, index): """ Creates example ASPIRE-generated CTF files. + + Note two distinct pixel sizes. """ star_fp = os.path.join(self.data_folder, f"ctf{index+1}.star") # note that values are arbitrary and not representative of actual CTF data @@ -263,7 +271,7 @@ def createTestCtfFiles(self, index): "_rlnSphericalAberration": 700 + index, "_rlnAmplitudeContrast": 600 + index, "_rlnVoltage": 500 + index, - "_rlnMicrographPixelSize": self.pixel_size, + "_rlnMicrographPixelSize": self.pixel_size + index * 0.01, } blocks = OrderedDict({"root": params_dict}) starfile = StarFile(blocks=blocks) @@ -272,6 +280,8 @@ def createTestCtfFiles(self, index): def createTestRelionCtfFile(self, reverse_optics_block_rows=False): """ Creates example RELION-generated CTF file for a set of micrographs. + + Note uniform pixel size. """ star_fp = os.path.join(self.data_folder, "micrographs_ctf.star") blocks = OrderedDict() @@ -532,8 +542,8 @@ def testWrongNumberCtfFiles(self): def testImportCtfFromList(self): src = BoxesCoordinateSource(self.files_box) src.import_aspire_ctf(self.ctf_files) - self._testCtfFilters(src) - self._testCtfMetadata(src) + self._testCtfFilters(src, uniform_pixel_sizes=False) + self._testCtfMetadata(src, uniform_pixel_sizes=False) def testImportCtfFromRelion(self): src = BoxesCoordinateSource(self.files_box) @@ -556,7 +566,7 @@ def testImportCtfFromRelionLegacy(self): self._testCtfFilters(src) self._testCtfMetadata(src) - def _testCtfFilters(self, src): + def _testCtfFilters(self, src, uniform_pixel_sizes=True): # there are two micrographs and two CTF files, so there should be two # unique CTF filters self.assertEqual(len(src.unique_filters), 2) @@ -592,6 +602,9 @@ def _testCtfFilters(self, src): ) ) filter1 = src.unique_filters[1] + pixel_size1 = self.pixel_size + if not uniform_pixel_sizes: + pixel_size1 += 0.01 self.assertTrue( np.allclose( np.array( @@ -602,7 +615,7 @@ def _testCtfFilters(self, src): 701.0, 601.0, 501.0, - self.pixel_size, + pixel_size1, ], dtype=src.dtype, ), @@ -629,7 +642,7 @@ def _testCtfFilters(self, src): np.array_equal(np.where(src.filter_indices == 1)[0], np.arange(200, 400)) ) - def _testCtfMetadata(self, src): + def _testCtfMetadata(self, src, uniform_pixel_sizes=True): # ensure metadata is populated correctly when adding CTF info # __mrc_filepath mrc_fp_metadata = np.array( @@ -664,8 +677,11 @@ def _testCtfMetadata(self, src): ctf_metadata[:200] = np.array( [1000.0, 900.0, 800.0 * np.pi / 180.0, 700.0, 600.0, 500.0, self.pixel_size] ) + pixel_size1 = self.pixel_size + if not uniform_pixel_sizes: + pixel_size1 += 0.01 ctf_metadata[200:400] = np.array( - [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, self.pixel_size] + [1001.0, 901.0, 801.0 * np.pi / 180.0, 701.0, 601.0, 501.0, pixel_size1] ) self.assertTrue(np.array_equal(ctf_metadata, src.get_metadata(ctf_cols))) @@ -729,6 +745,17 @@ def testPixelSizeWarning(self): src.import_relion_ctf(self.relion_ctf_file) np.testing.assert_approx_equal(src.pixel_size, manual_pixel_size) + def testMultiplePixelSizeWarning(self): + """ + Test source haveing a multiple pixel sizes in CTFFilter instances. + """ + src = BoxesCoordinateSource(self.files_box) # pixel_size=None + # Capture and compare warning message + with self._caplog.at_level(logging.WARNING): + src.import_aspire_ctf(self.ctf_files) # not uniform_pixel_sizes + assert src.pixel_size is None + assert "multiple pixel_sizes found" in self._caplog.text + def testPixelSize(self): """ Test explicitly providing correct pixel_size. From 6eea91adc403d1017fa94523814e89ad2c0c5bda Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 26 Jun 2025 11:31:19 -0400 Subject: [PATCH 176/216] fix strings spelling typos --- tests/test_coordinate_source.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_coordinate_source.py b/tests/test_coordinate_source.py index eb45140ad2..c8e8704e69 100644 --- a/tests/test_coordinate_source.py +++ b/tests/test_coordinate_source.py @@ -736,7 +736,7 @@ def testCommand(self): def testPixelSizeWarning(self): """ - Test source haveing a pixel size that conflicts with the CTFFilter instances. + Test source having a pixel size that conflicts with the CTFFilter instances. """ manual_pixel_size = 0.789 src = BoxesCoordinateSource(self.files_box, pixel_size=manual_pixel_size) @@ -747,7 +747,7 @@ def testPixelSizeWarning(self): def testMultiplePixelSizeWarning(self): """ - Test source haveing a multiple pixel sizes in CTFFilter instances. + Test source having multiple pixel sizes in CTFFilter instances. """ src = BoxesCoordinateSource(self.files_box) # pixel_size=None # Capture and compare warning message From 6b4573800e911cb34bf17b6240b3ad098736b2b6 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 26 Jun 2025 11:31:38 -0400 Subject: [PATCH 177/216] replace list(set(...)) with np.unqiue(...) --- src/aspire/source/coordinates.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index 8dee8d5d45..ac6d515b81 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -405,7 +405,7 @@ def _extract_ctf(self, data_block): # Check pixel_size # Get pixel_sizes from CTFFilters - ctf_pixel_sizes = list(set(filter_params[:, 6])) + ctf_pixel_sizes = np.unique(filter_params[:, 6]) # Compare with source.pixel_size if assigned if (self.pixel_size is not None) and ( not np.allclose(ctf_pixel_sizes, self.pixel_size) From 04c9c04a40f5666679727a408c1365e4726009f8 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Tue, 1 Jul 2025 14:49:49 -0400 Subject: [PATCH 178/216] Update comment --- src/aspire/source/coordinates.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/aspire/source/coordinates.py b/src/aspire/source/coordinates.py index ac6d515b81..7c8973792c 100644 --- a/src/aspire/source/coordinates.py +++ b/src/aspire/source/coordinates.py @@ -56,7 +56,9 @@ def __init__( :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. - :param pixel_size: Pixel size of the images in angstroms, default `None`. + :param pixel_size: Pixel size of the images in angstroms. + Default `None` will attempt to infer `pixel_size` from + `CTFFilter` objects when available. """ mrc_paths, coord_paths = [f[0] for f in files], [f[1] for f in files] # the particle_size parameter is the *user-specified* argument @@ -569,7 +571,9 @@ def __init__( :param max_rows: Maximum number of particles to read. (If `None`, will attempt to load all particles) :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. - :param pixel_size: Pixel size of the images in angstroms, default `None`. + :param pixel_size: Pixel size of the images in angstroms. + Default `None` will attempt to infer `pixel_size` from + `CTFFilter` objects when available. """ # instantiate super CoordinateSource.__init__( @@ -699,7 +703,9 @@ def __init__( attempt to load all particles) :param B: CTF envelope decay factor :param symmetry_group: A `SymmetryGroup` object or string corresponding to the symmetry of the molecule. - :param pixel_size: Pixel size of the images in angstroms, default `None`. + :param pixel_size: Pixel size of the images in angstroms. + Default `None` will attempt to infer `pixel_size` from + `CTFFilter` objects when available. """ # instantiate super CoordinateSource.__init__( From d41cce97e3b0a8f3d9ba6f4675c3ae75ad79b051 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 28 Apr 2025 15:29:35 -0400 Subject: [PATCH 179/216] Initial port of IRLS w/ spectral norm constraint. --- src/aspire/abinitio/__init__.py | 1 + src/aspire/abinitio/commonline_irls.py | 388 +++++++++++++++++++++++++ 2 files changed, 389 insertions(+) create mode 100644 src/aspire/abinitio/commonline_irls.py diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index 964c54bc37..fe91b322ad 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -3,6 +3,7 @@ # isort: off from .commonline_sdp import CommonlineSDP from .commonline_lud import CommonlineLUD +from .commonline_irls import CommonlineIRLS from .sync_voting import SyncVotingMixin from .commonline_sync import CLSyncVoting from .commonline_sync3n import CLSync3N diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py new file mode 100644 index 0000000000..9c175e6051 --- /dev/null +++ b/src/aspire/abinitio/commonline_irls.py @@ -0,0 +1,388 @@ +import logging + +import numpy as np +from scipy.sparse.linalg import eigsh + +from aspire.abinitio import CLOrient3D, CommonlineLUD, CommonlineSDP + +logger = logging.getLogger(__name__) + + +class CommonlineIRLS(CommonlineLUD): + """ + Define a derived class to estimate 3D orientations using Iteratively Reweighted + Least Squares (IRLS) as described in the following publication: + L. Wang, A. Singer, and Z. Wen, Orientation Determination of Cryo-EM Images Using + Least Unsquared Deviations, SIAM J. Imaging Sciences, 6, 2450-2483 (2013). + """ + + def __init__( + self, + src, + *, + num_itrs=10, + ctype=False, + eps_weighting=1e-3, + alpha=2 / 3, + max_rankZ=None, + max_rankW=None, + **kwargs, + ): + """ + Initialize a class for estimating 3D orientations using an IRLS-based optimization. + + :param ctype: Constraint type for the optimization: + - 1: ||G||_2 as a regularization term in the objective. + - 0: ||G||_2 as a constraint in the optimization. + :param alpha: Spectral norm constraint for ADMM algorithm. Default is None, which + does not apply a spectral norm constraint. To apply a spectral norm constraint provide + a value in the range [2/3, 1), 2/3 is recommended. + :param tol: Tolerance for convergence. The algorithm stops when conditions reach this threshold. + Default is 1e-3. + :param mu: The penalty parameter (or dual variable scaling factor) in the optimization problem. + Default is 1. + :param gam: Relaxation factor for updating variables in the algorithm (typically between 1 and 2). + Default is 1.618. + :param eps: Small positive value used to filter out negligible eigenvalues. + Default is 1e-12. + :param maxit: Maximum number of iterations allowed for the algorithm. + Default is 1000. + :param adp_proj: Flag for using adaptive projection during eigenvalue computation: + - True: Adaptive rank selection (Default). + - False: Full eigenvalue decomposition. + :param max_rankZ: Maximum rank used for projecting the Z matrix (for adaptive projection). + If None, defaults to max(6, n_img // 4). + :param max_rankW: Maximum rank used for projecting the W matrix (for adaptive projection). + If None, defaults to max(6, n_img // 4). + :param adp_mu: Adaptive adjustment of the penalty parameter `mu`: + - True: Enabled (Default). + - False: Disabled. + :param dec_mu: Scaling factor for decreasing `mu`. + Default is 0.5. + :param inc_mu: Scaling factor for increasing `mu`. + Default is 2. + :param mu_min: Minimum allowable value for `mu`. + Default is 1e-4. + :param mu_max: Maximum allowable value for `mu`. + Default is 1e4. + :param min_mu_itr: Minimum number of iterations before `mu` is adjusted. + Default is 5. + :param max_mu_itr: Maximum number of iterations allowed for `mu` adjustment. + Default is 20. + :param delta_mu_l: Lower bound for relative drop ratio to trigger a decrease in `mu`. + Default is 0.1. + :param delta_mu_u: Upper bound for relative drop ratio to trigger an increase in `mu`. + Default is 10. + """ + + self.num_itrs = num_itrs + self.ctype = ctype + self.eps_weighting = eps_weighting + + # Adjust rank limits + max_rankZ = max_rankZ or max(6, src.n // 4) + max_rankW = max_rankW or max(6, src.n // 4) + + super().__init__( + src, + max_rankZ=max_rankZ, + max_rankW=max_rankW, + alpha=alpha, + **kwargs, + ) + + self.lambda_ = self.alpha * self.n_img # Spectral norm bound + + def estimate_rotations(self): + """ + Estimate rotation matrices using the common lines method with LUD optimization. + """ + logger.info("Computing the common lines matrix.") + self.build_clmatrix() + + self.S = CommonlineSDP._construct_S(self, self.clmatrix) + weights = np.ones(2 * self.n_img, dtype=self.dtype) + gram = np.eye(2 * self.n_img, dtype=self.dtype) + for itr in range(self.num_itrs): + S = weights * self.S + gram = self._compute_Gram(gram, S) + weights = self._update_weights(gram) + self.rotations = CommonlineSDP._deterministic_rounding(gram) + + return self.rotations + + def _compute_Gram(self, G, S): + """ + Perform the alternating direction method of multipliers (ADMM) for the SDP + problem: + + min sum_{i self.eps + V = V[:, eigs_mask] + W = V @ np.diag(D[eigs_mask]) @ V.T + else: + # Determine number of eigenvalues to compute for adaptive projection + if itr == 0: + num_eigs = self.max_rankW + else: + num_eigs = self._compute_num_eigs(num_eigs, eigs_H, num_eigs, 50, 5) + + # Compute Eigenvectors and sort by largest algebraic eigenvalue + eigs_H, V = eigsh( + -H.astype(np.float64), + k=num_eigs, + which="LA", + tol=1e-15, + ) + eigs_H = eigs_H[::-1].astype(self.dtype, copy=False) + V = V[:, ::-1].astype(self.dtype, copy=False) + + nD = eigs_H > self.eps + eigs_H = eigs_H[nD] + num_eigs = np.count_nonzero(nD) + W = V[:, nD] @ np.diag(eigs_H) @ V[:, nD].T + H if nD.any() else H + + ############ + # Update G # + ############ + G = (1 - self.gam) * G + self.gam * self._mu * (W - H) + + # Check optimality + if self.ctype: + spG = eigsh(G.astype(np.float64, copy=False), k=1, which="LM") + pobj = -np.sum(S * G) + self.lambda_ + dobj = b.T @ y + else: + pobj = -np.sum(S * G) + dobj = (b.T @ y) - self.lambda_ * np.sum(abs(eigs_Z)) + + gap = abs(dobj - pobj) / (1 + abs(dobj) + abs(pobj)) + + resi = self._compute_AX(G) - b + pinf = np.linalg.norm(resi) / max(np.linalg.norm(b), 1) + + dinf_term = S + W + ATy - Z + dinf = np.linalg.norm(dinf_term, "fro") / max(np.linalg.norm(S, np.inf), 1) + + logger.info( + f"Iteration: {itr}, residual: {max(pinf, dinf)}, target: {self.tol}" + ) + if max(pinf, dinf, gap) <= self.tol: + return G + + # Update mu adaptively + if self.adp_mu: + if pinf / dinf <= self.delta_mu_l: + itmu_pinf = itmu_pinf + 1 + itmu_dinf = 0 + if itmu_pinf > self.max_mu_itr: + self._mu = max(self._mu * self.inc_mu, self.mu_min) + itmu_pinf = 0 + elif pinf / dinf > self.delta_mu_u: + itmu_dinf = itmu_dinf + 1 + itmu_pinf = 0 + if itmu_dinf > self.max_mu_itr: + self._mu = min(self._mu * self.dec_mu, self.mu_max) + itmu_dinf = 0 + return G + + def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): + """ + Update ADMM subproblem for enforcing the spectral norm constraint. + + :param S: A 2*n_img x 2*n_img symmetric matrix representing the fidelity term. + :param W: A 2*n_img x 2*n_img array, primary ADMM subproblem matrix. + :param ATy: A 2*n_img x 2*n_img array. + :param G: Current value of the 2*n_img x 2*n_img optimization solution matrix. + :param eigs_Z: eigenvalues from previous iteration. + :param num_eigs_Z: Number of eigenvalues of Z to use to enforce spectral norm constraint. + :param num_eigs: Number of eigenvalues of W used in previous iteration of ADMM. + + :returns: + - Z, Updated 2*n_img x 2*n_img matrix for spectral norm constraint ADMM subproblem. + - num_eigs_Z, Number of eigenvalues of Z to use to enforce spectral norm constraint in next iteration. + - num_eigs, Number of eigenvalues of W to use in this iteration of ADMM. + """ + B = S + W + ATy + G / self._mu + B = (B + B.T) / 2 + + if not self.adp_proj: + U, pi = np.linalg.eigh(B) + else: + # Determine number of eigenvalues to compute for adaptive projection + if num_eigs_Z is None: + num_eigs_Z = self.max_rankZ + else: + num_eigs_Z = self._compute_num_eigs(num_eigs_Z, eigs_Z, num_eigs, 10, 3) + + pi, U = eigsh( + B.astype(np.float64, copy=False), + k=num_eigs_Z, + which="LM", + ) # Compute top `num_eigs_Z` eigenvalues and eigenvectors + + # Sort by eigenvalue magnitude. Note, eigsh does not return + # ordered eigenvalues/vectors for which="LM". + idx = np.argsort(np.abs(pi))[::-1] + pi = pi[idx] + U = U[:, idx] + + # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. + # Compute eigenvalues based on constraint type. + if self.ctype: + pass + # Need to make this branch work. Compute projection onto simplex. + # eigs_Z = projsplx(np.abs(pi) / self.lambda_) + # eigs_Z = (self.lambda_ * np.sign(pi)) * eigs_Z + else: + eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - self.lambda_ / self._mu, 0) + + nD = abs(eigs_Z) > 0 + num_eigs_Z = np.count_nonzero(nD) + if num_eigs_Z > 0: + eigs_Z = eigs_Z[nD] + Z = U[:, nD] @ np.diag(eigs_Z) @ U[:, nD].T + else: + Z = np.zeros_like(B) + + return Z, num_eigs_Z, eigs_Z + + def _update_weights(self, gram): + K = self.n_img + W = self.S * gram + weights = W[:K, :K] + W[:K, K:] + W[K:, :K] + W[K:, K:] + weights = np.sqrt(abs(2 - 2 * weights), dtype=self.dtype) + weights = 1 / np.sqrt(weights**2 + self.eps_weighting**2) + updated_W = np.block([[weights, weights], [weights, weights]]) + + return updated_W + + @staticmethod + def _compute_AX(X): + """ + Compute the application of the linear operator A to the symmetric input + matrix X, where A(X) is defined as: + + A(X) = [ + diag(X_11), + diag(X_22), + sqrt(2)/2 * diag(X_12 + sqrt(2)/2 * diag(X_21) + ] + + where X_{ij} is the (i, j)'th K x K sub-block of X. + + :param X: 2D square array of shape (2K, 2K).. + :return: Flattened array representing A(X) + """ + K = X.shape[0] // 2 + + # Extract the diagonal elements of (X_ii^(11) and X_ii^(22)) + diags = np.diag(X) + + # Extract the diagonal elements from upper right KxK sub-block. + upper_right_diag_vals = np.diag(X[K:, :K]) + + # Compute the second part of AX, which is sqrt(2)/2 times the sum of + # the diagonal entries of the KxK off-diagonal sub-blocks of X. + sqrt_2_X_col = np.sqrt(2, dtype=X.dtype) * upper_right_diag_vals + + # Form AX by concatenating the results. + AX = np.concatenate((diags, sqrt_2_X_col)) + + return AX + + @staticmethod + def _compute_ATy(y): + """ + Compute the application of the adjoint operator A^T to the input vector y, + where + + y = [ + y_i^1, + y_i^2, + y_i^3 + ] for i = 1, 2, ..., K, + + and the adjoint of the operator A is defined as: + + AT(y) = Y = [ + [Y_ii^(11), Y_ii^(12)], + [Y_ii^(21), Y_ii^(22)] + ], + + where for i = 1, 2, ..., K: + + Y_ii^(11) = y_i^1, + Y_ii^(22) = y_i^2, + Y_ii^(12) = Y_ii^(21) = y_i^3 / sqrt(2). + + :param y: 1D NumPy array of length 3K. + :return: 2D NumPy array of shape (2K, 2K). + """ + K = len(y) // 3 + n = 2 * K # Size of the output matrix + ATy = np.zeros((n, n), dtype=y.dtype) + + # Assign diagonal elements + ATy[::1, ::1] = np.diag(y[:n]) + + # Assign symmetric off-diagonal elements + off_diag_vals = np.sqrt(2, dtype=y.dtype) / 2 * y[n:] + ATy[:K, K:] = np.diag(off_diag_vals) + ATy[K:, :K] = np.diag(off_diag_vals) + + return ATy From 15424d5591f8ea78b765b122f206958d2391db36 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 28 Apr 2025 15:53:05 -0400 Subject: [PATCH 180/216] Add basic test --- src/aspire/abinitio/commonline_irls.py | 4 +- tests/test_orient_irls.py | 105 +++++++++++++++++++++++++ 2 files changed, 107 insertions(+), 2 deletions(-) create mode 100644 tests/test_orient_irls.py diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 9c175e6051..7e1993e3c8 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -278,8 +278,8 @@ def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): # Sort by eigenvalue magnitude. Note, eigsh does not return # ordered eigenvalues/vectors for which="LM". idx = np.argsort(np.abs(pi))[::-1] - pi = pi[idx] - U = U[:, idx] + pi = pi[idx].astype(self.dtype, copy=False) + U = U[:, idx].astype(self.dtype, copy=False) # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. # Compute eigenvalues based on constraint type. diff --git a/tests/test_orient_irls.py b/tests/test_orient_irls.py new file mode 100644 index 0000000000..fd31c6365e --- /dev/null +++ b/tests/test_orient_irls.py @@ -0,0 +1,105 @@ +import numpy as np +import pytest + +from aspire.abinitio import CommonlineIRLS +from aspire.source import Simulation +from aspire.utils import mean_aligned_angular_distance +from aspire.volume import AsymmetricVolume + +RESOLUTION = [ + 32, + 33, +] + +OFFSETS = [ + None, # Defaults to random offsets. +] + +DTYPES = [ + np.float64, +] + +SPECTRAL_NORM_CONSTRAINT = [ + 2 / 3, +] + +ADAPTIVE_PROJECTION = [ + True, + False, +] + + +@pytest.fixture(params=RESOLUTION, ids=lambda x: f"resolution={x}") +def resolution(request): + return request.param + + +@pytest.fixture(params=OFFSETS, ids=lambda x: f"offsets={x}") +def offsets(request): + return request.param + + +@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}") +def dtype(request): + return request.param + + +@pytest.fixture(params=SPECTRAL_NORM_CONSTRAINT, ids=lambda x: f"alpha={x}") +def alpha(request): + return request.param + + +@pytest.fixture(params=ADAPTIVE_PROJECTION, ids=lambda x: f"adp_proj={x}") +def adp_proj(request): + return request.param + + +@pytest.fixture +def source(resolution, offsets, dtype): + """Fixture for simulation source object.""" + src = Simulation( + n=60, + L=resolution, + vols=AsymmetricVolume( + L=resolution, C=1, K=100, seed=10, dtype=dtype + ).generate(), + offsets=offsets, + amplitudes=1, + seed=0, + dtype=dtype, + ) + + # Cache source to prevent regenerating images. + src = src.cache() + + return src + + +@pytest.fixture +def orient_est(source, alpha, adp_proj): + """Fixture for LUD orientation estimation object.""" + # Generate LUD orientation estimation object. + orient_est = CommonlineIRLS( + source, + alpha=alpha, + adp_proj=adp_proj, + delta_mu_l=0.4, # Ensures branch is tested + mask=False, + tol=0.005, # Improves test speed + ) + + return orient_est + + +def test_estimate_rotations(source, orient_est): + # Estimate rotations + est_rots = orient_est.estimate_rotations() + + # Register estimates to ground truth rotations and compute the + # angular distance between them (in degrees). + # Assert that mean aligned angular distance is less than 5 degrees. + tol = 3 + mean_aligned_angular_distance(est_rots, source.rotations, degree_tol=tol) + + # Check dtype pass-through + np.testing.assert_equal(source.dtype, est_rots.dtype) From 77266175e6adf9c4478bfbdb61fb7dc8dc122d81 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 28 Apr 2025 15:57:33 -0400 Subject: [PATCH 181/216] tox --- src/aspire/abinitio/commonline_irls.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 7e1993e3c8..befba311bc 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -3,7 +3,7 @@ import numpy as np from scipy.sparse.linalg import eigsh -from aspire.abinitio import CLOrient3D, CommonlineLUD, CommonlineSDP +from aspire.abinitio import CommonlineLUD, CommonlineSDP logger = logging.getLogger(__name__) @@ -103,7 +103,7 @@ def estimate_rotations(self): self.S = CommonlineSDP._construct_S(self, self.clmatrix) weights = np.ones(2 * self.n_img, dtype=self.dtype) gram = np.eye(2 * self.n_img, dtype=self.dtype) - for itr in range(self.num_itrs): + for _ in range(self.num_itrs): S = weights * self.S gram = self._compute_Gram(gram, S) weights = self._update_weights(gram) @@ -203,9 +203,10 @@ def _compute_Gram(self, G, S): # Check optimality if self.ctype: - spG = eigsh(G.astype(np.float64, copy=False), k=1, which="LM") - pobj = -np.sum(S * G) + self.lambda_ - dobj = b.T @ y + pass + # spG = eigsh(G.astype(np.float64, copy=False), k=1, which="LM") + # pobj = -np.sum(S * G) + self.lambda_ + # dobj = b.T @ y else: pobj = -np.sum(S * G) dobj = (b.T @ y) - self.lambda_ * np.sum(abs(eigs_Z)) From cbef6c0237af106c6d1d2c9bc92b0eb6c8874f7e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Apr 2025 08:15:45 -0400 Subject: [PATCH 182/216] use 5 degree tol --- tests/test_orient_irls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_orient_irls.py b/tests/test_orient_irls.py index fd31c6365e..013ca6a200 100644 --- a/tests/test_orient_irls.py +++ b/tests/test_orient_irls.py @@ -98,7 +98,7 @@ def test_estimate_rotations(source, orient_est): # Register estimates to ground truth rotations and compute the # angular distance between them (in degrees). # Assert that mean aligned angular distance is less than 5 degrees. - tol = 3 + tol = 5 mean_aligned_angular_distance(est_rots, source.rotations, degree_tol=tol) # Check dtype pass-through From 732fe613fab5994afffeaf5891967e31e6c381f9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 2 May 2025 11:03:52 -0400 Subject: [PATCH 183/216] fix adp_proj branch --- src/aspire/abinitio/commonline_irls.py | 13 +++++++++---- src/aspire/abinitio/commonline_lud.py | 2 +- tests/test_orient_irls.py | 6 ++++-- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index befba311bc..2c5e8dd772 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -172,8 +172,13 @@ def _compute_Gram(self, G, S): if not self.adp_proj: D, V = np.linalg.eigh(H) eigs_mask = D > self.eps - V = V[:, eigs_mask] - W = V @ np.diag(D[eigs_mask]) @ V.T + num_eigs = np.count_nonzero(eigs_mask) + if num_eigs < n / 2: # few positive eigenvalues + V = V[:, eigs_mask] + W = V @ np.diag(D[eigs_mask]) @ V.T + else: # few negative eigenvalues + V = V[:, ~eigs_mask] + W = V @ np.diag(-D[~eigs_mask]) @ V.T + H else: # Determine number of eigenvalues to compute for adaptive projection if itr == 0: @@ -220,7 +225,7 @@ def _compute_Gram(self, G, S): dinf = np.linalg.norm(dinf_term, "fro") / max(np.linalg.norm(S, np.inf), 1) logger.info( - f"Iteration: {itr}, residual: {max(pinf, dinf)}, target: {self.tol}" + f"Iteration: {itr}, residual: {max(pinf, dinf, gap)}, target: {self.tol}" ) if max(pinf, dinf, gap) <= self.tol: return G @@ -262,7 +267,7 @@ def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): B = (B + B.T) / 2 if not self.adp_proj: - U, pi = np.linalg.eigh(B) + pi, U = np.linalg.eigh(B) else: # Determine number of eigenvalues to compute for adaptive projection if num_eigs_Z is None: diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 1898df4792..110289003a 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -300,7 +300,7 @@ def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): B = (B + B.T) / 2 if not self.adp_proj: - U, pi = np.linalg.eigh(B) + pi, U = np.linalg.eigh(B) else: # Determine number of eigenvalues to compute for adaptive projection if num_eigs_Z is None: diff --git a/tests/test_orient_irls.py b/tests/test_orient_irls.py index 013ca6a200..3e58cb7b48 100644 --- a/tests/test_orient_irls.py +++ b/tests/test_orient_irls.py @@ -8,15 +8,17 @@ RESOLUTION = [ 32, - 33, + pytest.param(33, marks=pytest.mark.expensive), ] OFFSETS = [ + 0, None, # Defaults to random offsets. ] DTYPES = [ - np.float64, + np.float32, + pytest.param(np.float64, marks=pytest.mark.expensive), ] SPECTRAL_NORM_CONSTRAINT = [ From 68e10f30c331c15f8bc465e93013c4505cbf76aa Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 2 May 2025 15:18:52 -0400 Subject: [PATCH 184/216] Add branch for no-spectral-norm-constraint. Add testing. --- src/aspire/abinitio/commonline_irls.py | 22 ++++++++++++++-------- src/aspire/abinitio/commonline_sdp.py | 2 +- tests/test_orient_irls.py | 2 ++ 3 files changed, 17 insertions(+), 9 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 2c5e8dd772..ab6d733977 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -91,11 +91,9 @@ def __init__( **kwargs, ) - self.lambda_ = self.alpha * self.n_img # Spectral norm bound - def estimate_rotations(self): """ - Estimate rotation matrices using the common lines method with LUD optimization. + Estimate rotation matrices using the common lines method with IRLS optimization. """ logger.info("Computing the common lines matrix.") self.build_clmatrix() @@ -103,10 +101,18 @@ def estimate_rotations(self): self.S = CommonlineSDP._construct_S(self, self.clmatrix) weights = np.ones(2 * self.n_img, dtype=self.dtype) gram = np.eye(2 * self.n_img, dtype=self.dtype) - for _ in range(self.num_itrs): - S = weights * self.S - gram = self._compute_Gram(gram, S) - weights = self._update_weights(gram) + if self.alpha is None: + A, b = CommonlineSDP._sdp_prep(self) + for _ in range(self.num_itrs): + S = weights * self.S + gram = CommonlineSDP._compute_gram_matrix(self, S, A, b) + weights = self._update_weights(gram) + else: + self.lambda_ = self.alpha * self.n_img # Spectral norm bound + for _ in range(self.num_itrs): + S = weights * self.S + gram = self._compute_Gram(gram, S) + weights = self._update_weights(gram) self.rotations = CommonlineSDP._deterministic_rounding(gram) return self.rotations @@ -120,7 +126,7 @@ def _compute_Gram(self, G, S): s.t. A(G) = b, G psd ||G||_2 <= lambda - Equivalent to matlab functions cryoEMSDPL12N/cryoEMSDPL12N_vsimple. + Equivalent to matlab function cryoEMADM. :param G: Gram matrix from previous iteration. :param S: Reweighted S matrix. :return: The updated gram matrix G. diff --git a/src/aspire/abinitio/commonline_sdp.py b/src/aspire/abinitio/commonline_sdp.py index b2604f1e17..6858e0f3fc 100644 --- a/src/aspire/abinitio/commonline_sdp.py +++ b/src/aspire/abinitio/commonline_sdp.py @@ -149,7 +149,7 @@ def _compute_gram_matrix(self, S, A, b): prob = cp.Problem(cp.Minimize(cp.trace(-S @ G)), constraints) prob.solve() - return G.value + return G.value.astype(self.dtype, copy=False) @staticmethod def _deterministic_rounding(gram): diff --git a/tests/test_orient_irls.py b/tests/test_orient_irls.py index 3e58cb7b48..d529c4a4bf 100644 --- a/tests/test_orient_irls.py +++ b/tests/test_orient_irls.py @@ -22,6 +22,7 @@ ] SPECTRAL_NORM_CONSTRAINT = [ + None, 2 / 3, ] @@ -88,6 +89,7 @@ def orient_est(source, alpha, adp_proj): delta_mu_l=0.4, # Ensures branch is tested mask=False, tol=0.005, # Improves test speed + num_itrs=3, # Improves test speed ) return orient_est From f678f8d56c7e043152c70ed47f1fc47b829183b9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 7 May 2025 11:58:52 -0400 Subject: [PATCH 185/216] remove ctype (dead branch in matlab). --- src/aspire/abinitio/commonline_irls.py | 30 ++++++-------------------- 1 file changed, 7 insertions(+), 23 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index ab6d733977..6effba5376 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -21,9 +21,8 @@ def __init__( src, *, num_itrs=10, - ctype=False, eps_weighting=1e-3, - alpha=2 / 3, + alpha=None, max_rankZ=None, max_rankW=None, **kwargs, @@ -31,9 +30,8 @@ def __init__( """ Initialize a class for estimating 3D orientations using an IRLS-based optimization. - :param ctype: Constraint type for the optimization: - - 1: ||G||_2 as a regularization term in the objective. - - 0: ||G||_2 as a constraint in the optimization. + :param num_itrs: Number of iterations for iterative reweighting. Default is 10. + :param eps_weighting: Regularization value for reweighting factor. Default is 1e-3. :param alpha: Spectral norm constraint for ADMM algorithm. Default is None, which does not apply a spectral norm constraint. To apply a spectral norm constraint provide a value in the range [2/3, 1), 2/3 is recommended. @@ -76,7 +74,6 @@ def __init__( """ self.num_itrs = num_itrs - self.ctype = ctype self.eps_weighting = eps_weighting # Adjust rank limits @@ -159,8 +156,7 @@ def _compute_Gram(self, G, S): # Compute y # ############# y = -(AS + self._compute_AX(W) - self._compute_AX(Z)) - resi / self._mu - # if self.ITR == 2: - # breakpoint() + ############# # Compute Z # ############# @@ -213,14 +209,8 @@ def _compute_Gram(self, G, S): G = (1 - self.gam) * G + self.gam * self._mu * (W - H) # Check optimality - if self.ctype: - pass - # spG = eigsh(G.astype(np.float64, copy=False), k=1, which="LM") - # pobj = -np.sum(S * G) + self.lambda_ - # dobj = b.T @ y - else: - pobj = -np.sum(S * G) - dobj = (b.T @ y) - self.lambda_ * np.sum(abs(eigs_Z)) + pobj = -np.sum(S * G) + dobj = (b.T @ y) - self.lambda_ * np.sum(abs(eigs_Z)) gap = abs(dobj - pobj) / (1 + abs(dobj) + abs(pobj)) @@ -295,13 +285,7 @@ def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. # Compute eigenvalues based on constraint type. - if self.ctype: - pass - # Need to make this branch work. Compute projection onto simplex. - # eigs_Z = projsplx(np.abs(pi) / self.lambda_) - # eigs_Z = (self.lambda_ * np.sign(pi)) * eigs_Z - else: - eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - self.lambda_ / self._mu, 0) + eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - self.lambda_ / self._mu, 0) nD = abs(eigs_Z) > 0 num_eigs_Z = np.count_nonzero(nD) From 8744b141ae3c1c764c42725f98e4832f4757af17 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 8 May 2025 08:54:28 -0400 Subject: [PATCH 186/216] Adapt lud compute_Z and use for IRLS --- src/aspire/abinitio/commonline_irls.py | 56 -------------------------- src/aspire/abinitio/commonline_lud.py | 41 ++++++++++--------- 2 files changed, 23 insertions(+), 74 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 6effba5376..7caf65f4b8 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -105,7 +105,6 @@ def estimate_rotations(self): gram = CommonlineSDP._compute_gram_matrix(self, S, A, b) weights = self._update_weights(gram) else: - self.lambda_ = self.alpha * self.n_img # Spectral norm bound for _ in range(self.num_itrs): S = weights * self.S gram = self._compute_Gram(gram, S) @@ -242,61 +241,6 @@ def _compute_Gram(self, G, S): itmu_dinf = 0 return G - def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): - """ - Update ADMM subproblem for enforcing the spectral norm constraint. - - :param S: A 2*n_img x 2*n_img symmetric matrix representing the fidelity term. - :param W: A 2*n_img x 2*n_img array, primary ADMM subproblem matrix. - :param ATy: A 2*n_img x 2*n_img array. - :param G: Current value of the 2*n_img x 2*n_img optimization solution matrix. - :param eigs_Z: eigenvalues from previous iteration. - :param num_eigs_Z: Number of eigenvalues of Z to use to enforce spectral norm constraint. - :param num_eigs: Number of eigenvalues of W used in previous iteration of ADMM. - - :returns: - - Z, Updated 2*n_img x 2*n_img matrix for spectral norm constraint ADMM subproblem. - - num_eigs_Z, Number of eigenvalues of Z to use to enforce spectral norm constraint in next iteration. - - num_eigs, Number of eigenvalues of W to use in this iteration of ADMM. - """ - B = S + W + ATy + G / self._mu - B = (B + B.T) / 2 - - if not self.adp_proj: - pi, U = np.linalg.eigh(B) - else: - # Determine number of eigenvalues to compute for adaptive projection - if num_eigs_Z is None: - num_eigs_Z = self.max_rankZ - else: - num_eigs_Z = self._compute_num_eigs(num_eigs_Z, eigs_Z, num_eigs, 10, 3) - - pi, U = eigsh( - B.astype(np.float64, copy=False), - k=num_eigs_Z, - which="LM", - ) # Compute top `num_eigs_Z` eigenvalues and eigenvectors - - # Sort by eigenvalue magnitude. Note, eigsh does not return - # ordered eigenvalues/vectors for which="LM". - idx = np.argsort(np.abs(pi))[::-1] - pi = pi[idx].astype(self.dtype, copy=False) - U = U[:, idx].astype(self.dtype, copy=False) - - # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. - # Compute eigenvalues based on constraint type. - eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - self.lambda_ / self._mu, 0) - - nD = abs(eigs_Z) > 0 - num_eigs_Z = np.count_nonzero(nD) - if num_eigs_Z > 0: - eigs_Z = eigs_Z[nD] - Z = U[:, nD] @ np.diag(eigs_Z) @ U[:, nD].T - else: - Z = np.zeros_like(B) - - return Z, num_eigs_Z, eigs_Z - def _update_weights(self, gram): K = self.n_img W = self.S * gram diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 110289003a..ab78a9d8bc 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -88,9 +88,14 @@ def __init__( # Handle parameters specific to CommonlineLUD self.alpha = alpha # Spectral norm constraint bound - if self.alpha is not None and not (2 / 3 <= self.alpha < 1): - raise ValueError("Spectral norm constraint, alpha, must be in [2/3, 1).") - + if self.alpha is not None: + if not (2 / 3 <= self.alpha < 1): + raise ValueError( + "Spectral norm constraint, alpha, must be in [2/3, 1)." + ) + else: + # Set spectral norm bound + self.lambda_ = self.alpha * src.n logger.info( f"Initializing LUD algorithm using ADMM with spectral norm constraint: {self.alpha}" ) @@ -158,6 +163,7 @@ def _compute_Gram(self): logger.info("Performing ADMM to compute Gram matrix.") # Initialize problem parameters + self._mu = self.mu n = 2 * self.n_img b = np.concatenate( [np.ones(n, dtype=self.dtype), np.zeros(self.n_img, dtype=self.dtype)] @@ -167,7 +173,7 @@ def _compute_Gram(self): G = np.eye(n, dtype=self.dtype) W = np.eye(n, dtype=self.dtype) Z = np.eye(n, dtype=self.dtype) - Phi = G / self.mu + Phi = G / self._mu if self.alpha is None: Phi += W @@ -186,7 +192,7 @@ def _compute_Gram(self): ############# # Compute y # ############# - y = -(AS + self._compute_AX(W)) - resi / self.mu + y = -(AS + self._compute_AX(W)) - resi / self._mu if self.alpha is not None: y += self._compute_AX(Z) @@ -194,7 +200,7 @@ def _compute_Gram(self): # Compute theta # ################# ATy = self._compute_ATy(y) - Phi = W + ATy + G / self.mu + Phi = W + ATy + G / self._mu if self.alpha is not None: Phi -= Z S, theta = self._Q_theta(Phi) @@ -210,7 +216,7 @@ def _compute_Gram(self): ############# # Compute W # ############# - H = -S - ATy - G / self.mu + H = -S - ATy - G / self._mu if self.alpha is not None: H += Z H = (H + H.T) / 2 @@ -244,7 +250,7 @@ def _compute_Gram(self): ############ # Update G # ############ - G = (1 - self.gam) * G + self.gam * self.mu * (W - H) + G = (1 - self.gam) * G + self.gam * self._mu * (W - H) # Check optimality resi = self._compute_AX(G) - b @@ -267,13 +273,13 @@ def _compute_Gram(self): itmu_pinf = itmu_pinf + 1 itmu_dinf = 0 if itmu_pinf > self.max_mu_itr: - self.mu = max(self.mu * self.inc_mu, self.mu_min) + self._mu = max(self._mu * self.inc_mu, self.mu_min) itmu_pinf = 0 elif pinf / dinf > self.delta_mu_u: itmu_dinf = itmu_dinf + 1 itmu_pinf = 0 if itmu_dinf > self.max_mu_itr: - self.mu = min(self.mu * self.dec_mu, self.mu_max) + self._mu = min(self._mu * self.dec_mu, self.mu_max) itmu_dinf = 0 return G @@ -295,8 +301,7 @@ def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): - num_eigs_Z, Number of eigenvalues of Z to use to enforce spectral norm constraint in next iteration. - num_eigs, Number of eigenvalues of W to use in this iteration of ADMM. """ - lambda_ = self.alpha * self.n_img # Spectral norm bound - B = S + W + ATy + G / self.mu + B = S + W + ATy + G / self._mu B = (B + B.T) / 2 if not self.adp_proj: @@ -316,12 +321,12 @@ def _compute_Z(self, S, W, ATy, G, eigs_Z, num_eigs_Z, num_eigs): # Sort by eigenvalue magnitude. Note, eigsh does not return # ordered eigenvalues/vectors for which="LM". idx = np.argsort(np.abs(pi))[::-1] - pi = pi[idx] - U = U[:, idx] + pi = pi[idx].astype(self.dtype, copy=False) + U = U[:, idx].astype(self.dtype, copy=False) # Apply soft-threshold to eigenvalues to enforce spectral norm constraint. - eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - lambda_ / self.mu, 0) - nD = eigs_Z > 0 + eigs_Z = np.sign(pi) * np.maximum(np.abs(pi) - self.lambda_ / self._mu, 0) + nD = abs(eigs_Z) > 0 num_eigs_Z = np.count_nonzero(nD) if num_eigs_Z > 0: eigs_Z = eigs_Z[nD] @@ -356,14 +361,14 @@ def _Q_theta(self, phi): # Compute theta phi = phi.reshape(self.n_img, 2, self.n_img, 2).transpose(0, 2, 1, 3) sum_prod = (phi[self.ut_mask] * self.C_t[self.ut_mask, None]).sum(axis=2) - theta[self.ut_mask] = self.C[self.ut_mask] - self.mu * sum_prod + theta[self.ut_mask] = self.C[self.ut_mask] - self._mu * sum_prod # Normalize theta theta_norm = np.linalg.norm(theta[self.ut_mask], axis=-1)[..., None] if self.alpha is not None: theta[self.ut_mask] /= theta_norm else: - theta[self.ut_mask] /= np.maximum(theta_norm, self.mu) + theta[self.ut_mask] /= np.maximum(theta_norm, self._mu) # Construct S S = theta[..., None] * self.C_t[:, :, None] From 9dee720a29bdbdd9acf1a21304aa83e3f2f2faef Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 8 May 2025 15:26:05 -0400 Subject: [PATCH 187/216] Adjoint property test --- tests/test_orient_irls.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_orient_irls.py b/tests/test_orient_irls.py index d529c4a4bf..6b1b584cd9 100644 --- a/tests/test_orient_irls.py +++ b/tests/test_orient_irls.py @@ -107,3 +107,22 @@ def test_estimate_rotations(source, orient_est): # Check dtype pass-through np.testing.assert_equal(source.dtype, est_rots.dtype) + + +def test_adjoint_property_A(dtype): + """ + Test = for random symmetric matrix `u` and + random vector `v`. + """ + n = 10 + u = np.random.rand(2 * n, 2 * n).astype(dtype, copy=False) + u = (u + u.T) / 2 + v = np.random.rand(3 * n).astype(dtype, copy=False) + + Au = CommonlineIRLS._compute_AX(u) + ATv = CommonlineIRLS._compute_ATy(v) + + lhs = np.dot(Au, v) + rhs = np.dot(u.flatten(), ATv.flatten()) + + np.testing.assert_allclose(lhs, rhs, rtol=1e-05, atol=1e-08) From 8a1172bfcc087b6b836701048430fb91e3f57e60 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 12 May 2025 13:44:28 -0400 Subject: [PATCH 188/216] add IRLS to LUD gallery --- .../commonline_lud_simulated_data.py | 55 ++++++++++--------- 1 file changed, 30 insertions(+), 25 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index 4883c575ae..0d95d0bdf1 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -16,10 +16,11 @@ import logging from fractions import Fraction +from itertools import product import numpy as np -from aspire.abinitio import CommonlineLUD +from aspire.abinitio import CommonlineIRLS, CommonlineLUD from aspire.noise import WhiteNoiseAdder from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance @@ -31,14 +32,15 @@ # %% # Parameters # ---------- -# Set up some initializing parameters. We will run the LUD algorithm -# for various levels of noise and output a table of results. +# Set up some initializing parameters. We will run the LUD algorithm using ADMM +# and IRLS methods under various spectral norm constraints and levels of noise. SNR = ["1/8", "1/16", "1/32"] # Signal-to-noise ratio +METHOD = ["ADMM", "IRLS"] +ALPHA = [None, 0.90, 0.75, 0.67] # Spectral norm constraint n_imgs = 500 # Number of images in our source dtype = np.float64 pad_size = 129 -results = {} # Dictionary to store results # %% # Load Volume Map @@ -60,10 +62,20 @@ # Generate Noisy Images and Estimate Rotations # -------------------------------------------- # A ``Simulation`` object is used to generate simulated data at various -# noise levels. Then rotations are estimated using the ``CommonlineLUD`` algorithm. -# Results are measured by computing the mean aligned angular distance between -# the ground truth rotations and the globally aligned estimated rotations. -for snr in SNR: +# noise levels. Then rotations are estimated using the ``CommonlineLUD`` and +# ``CommonlineIRLS`` algorithms. Results are measured by computing the mean +# aligned angular distance between the ground truth rotations and the globally +# aligned estimated rotations. + +# Build table to dislay results. +col_width = 15 +table = [] +table.append( + f"{'METHOD':<{col_width}} {'SNR':<{col_width}} {'ALPHA':<{col_width}} {'Mean Angular Distance':<{col_width}}" +) +table.append("-" * (col_width * 4)) + +for method, snr, alpha in product(METHOD, SNR, ALPHA): # Generate a white noise adder with specified SNR. noise_adder = WhiteNoiseAdder.from_snr(snr=Fraction(snr)) @@ -78,31 +90,24 @@ ).cache() # Estimate rotations using the LUD algorithm. - orient_est = CommonlineLUD(src) + if method == "ADMM": + orient_est = CommonlineLUD(src, alpha=alpha) + else: + orient_est = CommonlineIRLS(src, alpha=alpha) est_rotations = orient_est.estimate_rotations() # Find the mean aligned angular distance between estimates and ground truth rotations. mean_ang_dist = mean_aligned_angular_distance(est_rotations, src.rotations) - # Store results. - results[snr] = mean_ang_dist + # Append results to table. + table.append( + f"{method:<{col_width}} {snr:<{col_width}} {str(alpha):<{col_width}} {mean_ang_dist:<{col_width}}" + ) # %% # Display Results # --------------- -# Display table of results for various noise levels. - -# Column widths -col1_width = 10 -col2_width = 22 - -# Create table as a string -table = [] -table.append(f"{'SNR':<{col1_width}} {'Mean Angular Distance':<{col2_width}}") -table.append("-" * (col1_width + col2_width)) - -for snr, angle in results.items(): - table.append(f"{snr:<{col1_width}} {angle:<{col2_width}}") +# Display table of results for both methods using various spectral norm +# constraints and noise levels. -# Log the table logger.info("\n" + "\n".join(table)) From 40c4cbc1dfb61ca72b7e6d7767bc549215d1ed10 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 May 2025 13:13:16 -0400 Subject: [PATCH 189/216] edit gallery --- gallery/experiments/commonline_lud_simulated_data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gallery/experiments/commonline_lud_simulated_data.py b/gallery/experiments/commonline_lud_simulated_data.py index 0d95d0bdf1..da6f85e0fb 100644 --- a/gallery/experiments/commonline_lud_simulated_data.py +++ b/gallery/experiments/commonline_lud_simulated_data.py @@ -37,7 +37,7 @@ SNR = ["1/8", "1/16", "1/32"] # Signal-to-noise ratio METHOD = ["ADMM", "IRLS"] -ALPHA = [None, 0.90, 0.75, 0.67] # Spectral norm constraint +ALPHA = [0.90, 0.75, 0.67] # Spectral norm constraint n_imgs = 500 # Number of images in our source dtype = np.float64 pad_size = 129 @@ -68,7 +68,7 @@ # aligned estimated rotations. # Build table to dislay results. -col_width = 15 +col_width = 21 table = [] table.append( f"{'METHOD':<{col_width}} {'SNR':<{col_width}} {'ALPHA':<{col_width}} {'Mean Angular Distance':<{col_width}}" From 56bff84a931afb365e76918fb3d10adfb313e7d9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 May 2025 15:19:53 -0400 Subject: [PATCH 190/216] docstrings --- src/aspire/abinitio/commonline_irls.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 7caf65f4b8..6e784ccce5 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -32,7 +32,7 @@ def __init__( :param num_itrs: Number of iterations for iterative reweighting. Default is 10. :param eps_weighting: Regularization value for reweighting factor. Default is 1e-3. - :param alpha: Spectral norm constraint for ADMM algorithm. Default is None, which + :param alpha: Spectral norm constraint for IRLS algorithm. Default is None, which does not apply a spectral norm constraint. To apply a spectral norm constraint provide a value in the range [2/3, 1), 2/3 is recommended. :param tol: Tolerance for convergence. The algorithm stops when conditions reach this threshold. @@ -242,6 +242,12 @@ def _compute_Gram(self, G, S): return G def _update_weights(self, gram): + """ + Update the weight matrix for the IRLS algorithm. + + :param gram: 2K x 2K Gram matrix. + :return: 2K x 2K updated weight matrix. + """ K = self.n_img W = self.S * gram weights = W[:K, :K] + W[:K, K:] + W[K:, :K] + W[K:, K:] From c8e35463094e177b2c01ad8679448c2591a397d1 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 21 May 2025 14:06:15 -0400 Subject: [PATCH 191/216] Rearrange class structure --- src/aspire/abinitio/commonline_irls.py | 10 +++++----- src/aspire/abinitio/commonline_lud.py | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 6e784ccce5..26f5824263 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -3,7 +3,7 @@ import numpy as np from scipy.sparse.linalg import eigsh -from aspire.abinitio import CommonlineLUD, CommonlineSDP +from aspire.abinitio import CommonlineLUD logger = logging.getLogger(__name__) @@ -95,21 +95,21 @@ def estimate_rotations(self): logger.info("Computing the common lines matrix.") self.build_clmatrix() - self.S = CommonlineSDP._construct_S(self, self.clmatrix) + self.S = self._construct_S(self.clmatrix) weights = np.ones(2 * self.n_img, dtype=self.dtype) gram = np.eye(2 * self.n_img, dtype=self.dtype) if self.alpha is None: - A, b = CommonlineSDP._sdp_prep(self) + A, b = self._sdp_prep() for _ in range(self.num_itrs): S = weights * self.S - gram = CommonlineSDP._compute_gram_matrix(self, S, A, b) + gram = self._compute_gram_matrix(S, A, b) weights = self._update_weights(gram) else: for _ in range(self.num_itrs): S = weights * self.S gram = self._compute_Gram(gram, S) weights = self._update_weights(gram) - self.rotations = CommonlineSDP._deterministic_rounding(gram) + self.rotations = self._deterministic_rounding(gram) return self.rotations diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index ab78a9d8bc..79569cd80a 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -3,12 +3,12 @@ import numpy as np from scipy.sparse.linalg import eigsh -from aspire.abinitio import CLOrient3D, CommonlineSDP +from aspire.abinitio import CommonlineSDP logger = logging.getLogger(__name__) -class CommonlineLUD(CLOrient3D): +class CommonlineLUD(CommonlineSDP): """ Define a derived class to estimate 3D orientations using Least Unsquared Deviations as described in the following publication: @@ -140,7 +140,7 @@ def estimate_rotations(self): self._cl_to_C(self.clmatrix) gram = self._compute_Gram() gram = self._restructure_Gram(gram) - self.rotations = CommonlineSDP._deterministic_rounding(gram) + self.rotations = self._deterministic_rounding(gram) return self.rotations From 64609dae70a4562c5793e8f28262a23aaa3c785c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 12 Jun 2025 10:12:37 -0400 Subject: [PATCH 192/216] add comment about self._mu --- src/aspire/abinitio/commonline_lud.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 79569cd80a..2f5eb4efcc 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -163,6 +163,8 @@ def _compute_Gram(self): logger.info("Performing ADMM to compute Gram matrix.") # Initialize problem parameters + # Note, local self._mu must be reset each iteration when + # this method is used for IRLS. self._mu = self.mu n = 2 * self.n_img b = np.concatenate( From 635e4023b7a2ed39a1f3664fc31721cab42f9210 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 12 Jun 2025 11:10:47 -0400 Subject: [PATCH 193/216] update class docstring --- src/aspire/abinitio/commonline_irls.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 26f5824263..c4b075b2b8 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -10,8 +10,8 @@ class CommonlineIRLS(CommonlineLUD): """ - Define a derived class to estimate 3D orientations using Iteratively Reweighted - Least Squares (IRLS) as described in the following publication: + Estimate 3D orientations using Iteratively Reweighted Least Squares + (IRLS) as described in the following publication: L. Wang, A. Singer, and Z. Wen, Orientation Determination of Cryo-EM Images Using Least Unsquared Deviations, SIAM J. Imaging Sciences, 6, 2450-2483 (2013). """ From 76d8949fe9eb06c42ec5b6313e94cf71b4c1b179 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 12 Jun 2025 14:40:37 -0400 Subject: [PATCH 194/216] refer parent class for argument docstring --- src/aspire/abinitio/commonline_irls.py | 33 +------------------------- 1 file changed, 1 insertion(+), 32 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index c4b075b2b8..544229f289 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -29,48 +29,17 @@ def __init__( ): """ Initialize a class for estimating 3D orientations using an IRLS-based optimization. + See CommonlineLUD for additional arguments. :param num_itrs: Number of iterations for iterative reweighting. Default is 10. :param eps_weighting: Regularization value for reweighting factor. Default is 1e-3. :param alpha: Spectral norm constraint for IRLS algorithm. Default is None, which does not apply a spectral norm constraint. To apply a spectral norm constraint provide a value in the range [2/3, 1), 2/3 is recommended. - :param tol: Tolerance for convergence. The algorithm stops when conditions reach this threshold. - Default is 1e-3. - :param mu: The penalty parameter (or dual variable scaling factor) in the optimization problem. - Default is 1. - :param gam: Relaxation factor for updating variables in the algorithm (typically between 1 and 2). - Default is 1.618. - :param eps: Small positive value used to filter out negligible eigenvalues. - Default is 1e-12. - :param maxit: Maximum number of iterations allowed for the algorithm. - Default is 1000. - :param adp_proj: Flag for using adaptive projection during eigenvalue computation: - - True: Adaptive rank selection (Default). - - False: Full eigenvalue decomposition. :param max_rankZ: Maximum rank used for projecting the Z matrix (for adaptive projection). If None, defaults to max(6, n_img // 4). :param max_rankW: Maximum rank used for projecting the W matrix (for adaptive projection). If None, defaults to max(6, n_img // 4). - :param adp_mu: Adaptive adjustment of the penalty parameter `mu`: - - True: Enabled (Default). - - False: Disabled. - :param dec_mu: Scaling factor for decreasing `mu`. - Default is 0.5. - :param inc_mu: Scaling factor for increasing `mu`. - Default is 2. - :param mu_min: Minimum allowable value for `mu`. - Default is 1e-4. - :param mu_max: Maximum allowable value for `mu`. - Default is 1e4. - :param min_mu_itr: Minimum number of iterations before `mu` is adjusted. - Default is 5. - :param max_mu_itr: Maximum number of iterations allowed for `mu` adjustment. - Default is 20. - :param delta_mu_l: Lower bound for relative drop ratio to trigger a decrease in `mu`. - Default is 0.1. - :param delta_mu_u: Upper bound for relative drop ratio to trigger an increase in `mu`. - Default is 10. """ self.num_itrs = num_itrs From 6916d6d4b007f426b56a9432e81553a9fefcde8d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 13 Jun 2025 10:05:17 -0400 Subject: [PATCH 195/216] S -> S_weighted --- src/aspire/abinitio/commonline_irls.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 544229f289..fc3a2b19be 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -70,13 +70,13 @@ def estimate_rotations(self): if self.alpha is None: A, b = self._sdp_prep() for _ in range(self.num_itrs): - S = weights * self.S - gram = self._compute_gram_matrix(S, A, b) + S_weighted = weights * self.S + gram = self._compute_gram_matrix(S_weighted, A, b) weights = self._update_weights(gram) else: for _ in range(self.num_itrs): - S = weights * self.S - gram = self._compute_Gram(gram, S) + S_weighted = weights * self.S + gram = self._compute_Gram(gram, S_weighted) weights = self._update_weights(gram) self.rotations = self._deterministic_rounding(gram) From 493ef77019a4027548f2a6e7aff320d675428757 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 13 Jun 2025 11:57:33 -0400 Subject: [PATCH 196/216] use broadcasting --- src/aspire/abinitio/commonline_irls.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index fc3a2b19be..876a376869 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -144,11 +144,10 @@ def _compute_Gram(self, G, S): eigs_mask = D > self.eps num_eigs = np.count_nonzero(eigs_mask) if num_eigs < n / 2: # few positive eigenvalues - V = V[:, eigs_mask] - W = V @ np.diag(D[eigs_mask]) @ V.T + # Equivalent to V D V', computed with broadcasting for efficiency + W = (V[:, eigs_mask] * D[eigs_mask][None, :]) @ V[:, eigs_mask].T else: # few negative eigenvalues - V = V[:, ~eigs_mask] - W = V @ np.diag(-D[~eigs_mask]) @ V.T + H + W = (V[:, ~eigs_mask] * (-D[~eigs_mask])[None, :]) @ V[:, ~eigs_mask].T + H else: # Determine number of eigenvalues to compute for adaptive projection if itr == 0: @@ -169,7 +168,11 @@ def _compute_Gram(self, G, S): nD = eigs_H > self.eps eigs_H = eigs_H[nD] num_eigs = np.count_nonzero(nD) - W = V[:, nD] @ np.diag(eigs_H) @ V[:, nD].T + H if nD.any() else H + if nD.any(): + # Low-rank update: V diag(eigs_H) V^T + H, done via broadcasting + W1 = (V[:, nD] * eigs_H[None, :]) @ V[:, nD].T + H + else: + W = H ############ # Update G # From b71570252f2322d6d8fe641609986e90eb3a592d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 13 Jun 2025 14:02:15 -0400 Subject: [PATCH 197/216] eigenvalue eps for singles/doubles --- src/aspire/abinitio/commonline_irls.py | 4 +++- src/aspire/abinitio/commonline_lud.py | 10 +++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 876a376869..63d2b72a36 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -147,7 +147,9 @@ def _compute_Gram(self, G, S): # Equivalent to V D V', computed with broadcasting for efficiency W = (V[:, eigs_mask] * D[eigs_mask][None, :]) @ V[:, eigs_mask].T else: # few negative eigenvalues - W = (V[:, ~eigs_mask] * (-D[~eigs_mask])[None, :]) @ V[:, ~eigs_mask].T + H + W = (V[:, ~eigs_mask] * (-D[~eigs_mask])[None, :]) @ V[ + :, ~eigs_mask + ].T + H else: # Determine number of eigenvalues to compute for adaptive projection if itr == 0: diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index 2f5eb4efcc..e53af01142 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -23,7 +23,7 @@ def __init__( tol=1e-3, mu=1, gam=1.618, - eps=1e-12, + eps=None, maxit=1000, adp_proj=True, max_rankZ=None, @@ -55,7 +55,7 @@ def __init__( :param gam: Relaxation factor for updating variables in the algorithm (typically between 1 and 2). Default is 1.618. :param eps: Small positive value used to filter out negligible eigenvalues. - Default is 1e-12. + Default is 1e-5 if src.dtype is singles, otherwise 1e-12. :param maxit: Maximum number of iterations allowed for the algorithm. Default is 1000. :param adp_proj: Flag for using adaptive projection during eigenvalue computation: @@ -103,7 +103,6 @@ def __init__( self.tol = tol self.mu = mu self.gam = gam - self.eps = eps self.maxit = maxit self.adp_proj = adp_proj @@ -121,6 +120,11 @@ def __init__( # Initialize commonline base class super().__init__(src, **kwargs) + # Set eps for eigenvalue filter + if eps is None: + eps = 1e-5 if self.dtype == np.float32 else 1e-12 + self.eps = eps + # Adjust rank limits self.max_rankZ = max_rankW or max(6, self.n_img // 2) self.max_rankW = max_rankW or max(6, self.n_img // 2) From 9ad9f030a7d11b1fc92c848a4dffdd9a37e33eb3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 13 Jun 2025 14:07:49 -0400 Subject: [PATCH 198/216] variable name --- src/aspire/abinitio/commonline_irls.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 63d2b72a36..3e4df2fe73 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -172,7 +172,7 @@ def _compute_Gram(self, G, S): num_eigs = np.count_nonzero(nD) if nD.any(): # Low-rank update: V diag(eigs_H) V^T + H, done via broadcasting - W1 = (V[:, nD] * eigs_H[None, :]) @ V[:, nD].T + H + W = (V[:, nD] * eigs_H[None, :]) @ V[:, nD].T + H else: W = H From 14096cdc0d1aa50425633bfce200d77c44a8b631 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 16 Jun 2025 10:11:25 -0400 Subject: [PATCH 199/216] update docstring --- src/aspire/abinitio/commonline_irls.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 3e4df2fe73..0f253015c3 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -84,10 +84,10 @@ def estimate_rotations(self): def _compute_Gram(self, G, S): """ - Perform the alternating direction method of multipliers (ADMM) for the SDP - problem: + Given G^(k), solve for G^(k+1) using the alternating direction method of multipliers (ADMM) + for the IRLS problem: - min sum_{i s.t. A(G) = b, G psd ||G||_2 <= lambda From a914fd17b219a24f81f81fd1bcf677a1b2ef21b7 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 16 Jun 2025 10:43:46 -0400 Subject: [PATCH 200/216] Fix compute_AX/compute_ATy docstrings to distinguish between LUD/IRLS. --- src/aspire/abinitio/commonline_irls.py | 45 ++++++++++++++----------- src/aspire/abinitio/commonline_lud.py | 46 ++++++++++++++------------ 2 files changed, 50 insertions(+), 41 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 0f253015c3..5ebae48496 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -234,19 +234,27 @@ def _update_weights(self, gram): @staticmethod def _compute_AX(X): """ - Compute the application of the linear operator A to the symmetric input - matrix X, where A(X) is defined as: + Compute the application of the linear operator A to a symmetric input matrix X, + where X is a 2K x 2K matrix composed of four K x K blocks: - A(X) = [ - diag(X_11), - diag(X_22), - sqrt(2)/2 * diag(X_12 + sqrt(2)/2 * diag(X_21) - ] + X = [X_11 X_12] + [X_21 X_22] - where X_{ij} is the (i, j)'th K x K sub-block of X. + The operator A maps X to a 3K-dimensional vector as: - :param X: 2D square array of shape (2K, 2K).. - :return: Flattened array representing A(X) + A(X) = [ + diag(X_11), + diag(X_22), + sqrt(2) * diag(X_21) + ] + + where: + - diag(X_11) and diag(X_22) extract the diagonals of the top-left and bottom-right K x K blocks, + - diag(X_21) (i.e., the lower-left block) extracts the diagonal of the off-diagonal block, + scaled by sqrt(2) to account for symmetry. + + :param X: 2D symmetric NumPy array of shape (2K, 2K). + :return: 1D NumPy array of length 3K representing A(X). """ K = X.shape[0] // 2 @@ -277,18 +285,17 @@ def _compute_ATy(y): y_i^3 ] for i = 1, 2, ..., K, - and the adjoint of the operator A is defined as: + and the adjoint operator produces a 2K×2K symmetric matrix Y, where: - AT(y) = Y = [ - [Y_ii^(11), Y_ii^(12)], - [Y_ii^(21), Y_ii^(22)] - ], + - The first K elements y_i^1 are placed on the diagonal entries Y_ii for i = 0 to K−1. + - The next K elements y_i^2 are placed on the diagonal entries Y_ii for i = K to 2K−1. + - The final K elements y_i^3 are scaled by 1/√2 and placed as diagonal entries in the + upper-right block ([:K, K:]) and symmetrically in the lower-left block ([K:, :K]). - where for i = 1, 2, ..., K: + This results in a matrix of the form: - Y_ii^(11) = y_i^1, - Y_ii^(22) = y_i^2, - Y_ii^(12) = Y_ii^(21) = y_i^3 / sqrt(2). + Y = [ diag(y^1) diag(y^3 / sqrt(2)) ] + [ diag(y^3 / sqrt(2)) diag(y^2) ] :param y: 1D NumPy array of length 3K. :return: 2D NumPy array of shape (2K, 2K). diff --git a/src/aspire/abinitio/commonline_lud.py b/src/aspire/abinitio/commonline_lud.py index e53af01142..759818d8df 100644 --- a/src/aspire/abinitio/commonline_lud.py +++ b/src/aspire/abinitio/commonline_lud.py @@ -454,21 +454,26 @@ def _compute_num_eigs(num_eigs_prev, eig_vec, num_eigs_W, rel_drp_thresh, eigs_i @staticmethod def _compute_AX(X): """ - Compute the application of the linear operator A to the symmetric input - matrix X, where A(X) is defined as: + Compute the application of the linear operator A to a symmetric input matrix X, + where X is a 2K x 2K matrix consisting of K diagonal 2 x 2 blocks. - A(X) = [ - X_ii^(11), - X_ii^(22), - sqrt(2)/2 * X_ii^(12) + sqrt(2)/2 * X_ii^(21) - ] + The operator A maps X to a 3K-dimensional vector as follows: + + For each 2 x 2 diagonal block X_i: - i = 1, 2, ..., K + A(X_i) = [ + X_i[0, 0], + X_i[1, 1], + sqrt(2) * X_i[0, 1] + ] - where X_{ii}^{pq} denotes the (p,q)-th element in the 2x2 sub-block X_{ii}. + That is: + - The first K entries are the (0,0) elements from each 2 x 2 block, + - The next K are the (1,1) elements, + - The final K are sqrt(2) times the off-diagonal (0,1) elements. - :param X: 2D square array of shape (2K, 2K).. - :return: Flattened array representing A(X) + :param X: 2D symmetric NumPy array of shape (2K, 2K), with 2 x 2 blocks along the diagonal. + :return: 1D NumPy array of length 3K representing A(X). """ # Extract the diagonal elements of (X_ii^(11) and X_ii^(22)) diags = np.diag(X) @@ -499,21 +504,18 @@ def _compute_ATy(y): y_i^3 ] for i = 1, 2, ..., K, - and the adjoint of the operator A is defined as: + and the adjoint operator produces a 2K×2K matrix Y, where each 2×2 block + along the diagonal has the form: - AT(y) = Y = [ - [Y_ii^(11), Y_ii^(12)], - [Y_ii^(21), Y_ii^(22)] - ], - - where for i = 1, 2, ..., K: + Y_i = [ + [y_i^1, y_i^3 / sqrt(2)], + [y_i^3 / sqrt(2), y_i^2] + ] - Y_ii^(11) = y_i^1, - Y_ii^(22) = y_i^2, - Y_ii^(12) = Y_ii^(21) = y_i^3 / sqrt(2). + All off-diagonal blocks in Y are zero. :param y: 1D NumPy array of length 3K. - :return: 2D NumPy array of shape (2K, 2K). + :return: 2D NumPy array of shape (2K, 2K) with 2×2 symmetric blocks on the diagonal.. """ K = len(y) // 3 n = 2 * K # Size of the output matrix From bd35533bb7938d36d5154960c5ba998833f09df4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 16 Jun 2025 12:59:55 -0400 Subject: [PATCH 201/216] add a pbar --- src/aspire/abinitio/commonline_irls.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 5ebae48496..b083a7fe39 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -4,6 +4,7 @@ from scipy.sparse.linalg import eigsh from aspire.abinitio import CommonlineLUD +from aspire.utils import trange logger = logging.getLogger(__name__) @@ -69,12 +70,12 @@ def estimate_rotations(self): gram = np.eye(2 * self.n_img, dtype=self.dtype) if self.alpha is None: A, b = self._sdp_prep() - for _ in range(self.num_itrs): + for _ in trange(self.num_itrs, desc="Performing iterative re-weighting."): S_weighted = weights * self.S gram = self._compute_gram_matrix(S_weighted, A, b) weights = self._update_weights(gram) else: - for _ in range(self.num_itrs): + for _ in trange(self.num_itrs, desc="Performing iterative re-weighting."): S_weighted = weights * self.S gram = self._compute_Gram(gram, S_weighted) weights = self._update_weights(gram) From 8c7b84548e8ff34ba47a5c0f393309f767fa1b01 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 16 Jun 2025 13:04:09 -0400 Subject: [PATCH 202/216] use *args --- src/aspire/abinitio/commonline_irls.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index b083a7fe39..15be2269f8 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -20,7 +20,7 @@ class CommonlineIRLS(CommonlineLUD): def __init__( self, src, - *, + *args, num_itrs=10, eps_weighting=1e-3, alpha=None, @@ -52,6 +52,7 @@ def __init__( super().__init__( src, + *args, max_rankZ=max_rankZ, max_rankW=max_rankW, alpha=alpha, From cd93f26407bfc7769176ac0972a187a8c4b3d6f8 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Jul 2025 15:01:45 -0400 Subject: [PATCH 203/216] _compute_gram_matrix ->_compute_gram_SDP --- src/aspire/abinitio/commonline_irls.py | 2 +- src/aspire/abinitio/commonline_sdp.py | 6 +++--- tests/test_orient_sdp.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 15be2269f8..4977535ddf 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -73,7 +73,7 @@ def estimate_rotations(self): A, b = self._sdp_prep() for _ in trange(self.num_itrs, desc="Performing iterative re-weighting."): S_weighted = weights * self.S - gram = self._compute_gram_matrix(S_weighted, A, b) + gram = self._compute_gram_SDP(S_weighted, A, b) weights = self._update_weights(gram) else: for _ in trange(self.num_itrs, desc="Performing iterative re-weighting."): diff --git a/src/aspire/abinitio/commonline_sdp.py b/src/aspire/abinitio/commonline_sdp.py index 6858e0f3fc..140fe27733 100644 --- a/src/aspire/abinitio/commonline_sdp.py +++ b/src/aspire/abinitio/commonline_sdp.py @@ -32,7 +32,7 @@ def estimate_rotations(self): S = self._construct_S(self.clmatrix) A, b = self._sdp_prep() - gram = self._compute_gram_matrix(S, A, b) + gram = self._compute_gram_SDP(S, A, b) self.rotations = self._deterministic_rounding(gram) return self.rotations @@ -83,7 +83,7 @@ def _sdp_prep(self): """ Prepare optimization problem constraints. - The constraints for the SDP optimization, max tr(SG), performed in `_compute_gram_matrix()` + The constraints for the SDP optimization, max tr(SG), performed in `_compute_gram_SDP()` as min tr(-SG), are that the Gram matrix, G, is semidefinite positive and G11_ii = G22_ii = 1, G12_ii = G21_ii = 0, i=1,2,...,N, for the block representation of G = [[G11, G12], [G21, G22]]. @@ -117,7 +117,7 @@ def _sdp_prep(self): return A, b - def _compute_gram_matrix(self, S, A, b): + def _compute_gram_SDP(self, S, A, b): """ Compute the Gram matrix by solving an SDP optimization. diff --git a/tests/test_orient_sdp.py b/tests/test_orient_sdp.py index 867ec7b68a..eb6d36e104 100644 --- a/tests/test_orient_sdp.py +++ b/tests/test_orient_sdp.py @@ -122,7 +122,7 @@ def test_gram_matrix(src_orient_est_fixture): # Estimate the Gram matrix A, b = orient_est._sdp_prep() - gram = orient_est._compute_gram_matrix(S, A, b) + gram = orient_est._compute_gram_SDP(S, A, b) # Construct the ground truth Gram matrix, G = R @ R.T, where R = [R1, R2] # with R1 and R2 being the concatenation of the first and second columns From 52d3ce61415bf0cd06f03dc559443f3127148d9f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Jul 2025 15:04:06 -0400 Subject: [PATCH 204/216] _compute_Gram ->_compute_gram_IRLS --- src/aspire/abinitio/commonline_irls.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_irls.py b/src/aspire/abinitio/commonline_irls.py index 4977535ddf..f608adc679 100644 --- a/src/aspire/abinitio/commonline_irls.py +++ b/src/aspire/abinitio/commonline_irls.py @@ -78,13 +78,13 @@ def estimate_rotations(self): else: for _ in trange(self.num_itrs, desc="Performing iterative re-weighting."): S_weighted = weights * self.S - gram = self._compute_Gram(gram, S_weighted) + gram = self._compute_gram_IRLS(gram, S_weighted) weights = self._update_weights(gram) self.rotations = self._deterministic_rounding(gram) return self.rotations - def _compute_Gram(self, G, S): + def _compute_gram_IRLS(self, G, S): """ Given G^(k), solve for G^(k+1) using the alternating direction method of multipliers (ADMM) for the IRLS problem: From 16858783224dfcb67d9bcb98f41fab1e2d9d4ee4 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 3 Jul 2025 09:09:59 -0400 Subject: [PATCH 205/216] resolve ambiguous dtype in Downsample [skip ci] --- src/aspire/image/image.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 8ee9aec96d..3989699e43 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -559,7 +559,12 @@ def downsample(self, ds_res, zero_nyquist=True, legacy=False): out = fft.ifft2(fft.ifftshift(crop_fx)) else: out = fft.centered_ifft2(crop_fx) - out = xp.asnumpy(out.real * ds_res**2 / self.resolution**2) + + # The parenths are required because dtype casting semantics + # differs between Numpy 1, 2, and CuPy. + # At time of writing CuPy is consistent with Numpy1. + # The additional parenths yield consistent out.dtype. + out = xp.asnumpy(out.real * (ds_res**2 / self.resolution**2)) # Optionally scale pixel size ds_pixel_size = self.pixel_size From 5cd8bc7399404b1f538a468f172814e014773072 Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 3 Jul 2025 09:26:42 -0400 Subject: [PATCH 206/216] Update comment with GH issue --- src/aspire/image/image.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py index 3989699e43..ed061ecec7 100644 --- a/src/aspire/image/image.py +++ b/src/aspire/image/image.py @@ -564,6 +564,7 @@ def downsample(self, ds_res, zero_nyquist=True, legacy=False): # differs between Numpy 1, 2, and CuPy. # At time of writing CuPy is consistent with Numpy1. # The additional parenths yield consistent out.dtype. + # See #1298 for relevant debugger output. out = xp.asnumpy(out.real * (ds_res**2 / self.resolution**2)) # Optionally scale pixel size From 0afa283b953298eb3b8b19bc16ab37e4393bfe6e Mon Sep 17 00:00:00 2001 From: Garrett Wright Date: Thu, 3 Jul 2025 11:26:11 -0400 Subject: [PATCH 207/216] =?UTF-8?q?Bump=20version:=200.13.2=20=E2=86=92=20?= =?UTF-8?q?0.14.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .bumpversion.cfg | 2 +- README.md | 4 ++-- docs/source/conf.py | 2 +- docs/source/index.rst | 2 +- pyproject.toml | 2 +- src/aspire/__init__.py | 2 +- src/aspire/config_default.yaml | 2 +- 7 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.bumpversion.cfg b/.bumpversion.cfg index 7986615540..cd7f438681 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.13.2 +current_version = 0.14.0 commit = True tag = True diff --git a/README.md b/README.md index 56ab26d65f..9d8a9d4cf6 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5657281.svg)](https://doi.org/10.5281/zenodo.5657281) [![Downloads](https://static.pepy.tech/badge/aspire/month)](https://pepy.tech/project/aspire) -# ASPIRE - Algorithms for Single Particle Reconstruction - v0.13.2 +# ASPIRE - Algorithms for Single Particle Reconstruction - v0.14.0 The ASPIRE-Python project supersedes [Matlab ASPIRE](https://github.com/PrincetonUniversity/aspire). @@ -20,7 +20,7 @@ For more information about the project, algorithms, and related publications ple Please cite using the following DOI. This DOI represents all versions, and will always resolve to the latest one. ``` -ComputationalCryoEM/ASPIRE-Python: v0.13.2 https://doi.org/10.5281/zenodo.5657281 +ComputationalCryoEM/ASPIRE-Python: v0.14.0 https://doi.org/10.5281/zenodo.5657281 ``` diff --git a/docs/source/conf.py b/docs/source/conf.py index 21e256e48f..69513931b1 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -86,7 +86,7 @@ # built documents. # # The full version, including alpha/beta/rc tags. -release = version = "0.13.2" +release = version = "0.14.0" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. diff --git a/docs/source/index.rst b/docs/source/index.rst index 88f8107752..e841b6a49f 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,4 +1,4 @@ -Aspire v0.13.2 +Aspire v0.14.0 ============== Algorithms for Single Particle Reconstruction diff --git a/pyproject.toml b/pyproject.toml index 25e400dfaa..328e25aad3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "aspire" -version = "0.13.2" +version = "0.14.0" description = "Algorithms for Single Particle Reconstruction" readme = "README.md" # Optional requires-python = ">=3.9" diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py index f332f8a6e0..35223891ca 100644 --- a/src/aspire/__init__.py +++ b/src/aspire/__init__.py @@ -15,7 +15,7 @@ from aspire.exceptions import handle_exception # version in maj.min.bld format -__version__ = "0.13.2" +__version__ = "0.14.0" # Setup `confuse` config diff --git a/src/aspire/config_default.yaml b/src/aspire/config_default.yaml index e1eadbb08b..53b8b7ab87 100644 --- a/src/aspire/config_default.yaml +++ b/src/aspire/config_default.yaml @@ -1,4 +1,4 @@ -version: 0.13.2 +version: 0.14.0 common: # numeric module to use - one of numpy/cupy numeric: numpy From 33e24afae3f5c602e11f6d88c3e3c537ff2ea1e9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 26 Jun 2025 15:47:27 -0400 Subject: [PATCH 208/216] update pipeline demo --- gallery/tutorials/pipeline_demo.py | 131 +++++++++++++++++------------ 1 file changed, 78 insertions(+), 53 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index 77d304b156..d58e17976b 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -22,16 +22,18 @@ # Load 80s Ribosome as a ``Volume`` object. original_vol = emdb_2660() -# Downsample the volume -res = 41 -vol = original_vol.downsample(res) +# During the preprocessing stages of the pipeline we will downsample +# the images to an image size of 64 pixels. Here, we also downsample the +# volume so we can compare to our reconstruction later. +res = 64 +vol_ds = original_vol.downsample(res) # %% # .. note:: # A ``Volume`` can be saved using the ``Volume.save()`` method as follows:: # # fn = f"downsampled_80s_ribosome_size{res}.mrc" -# vol.save(fn, overwrite=True) +# vol_ds.save(fn, overwrite=True) # %% @@ -63,7 +65,7 @@ defocus_ct = 7 ctf_filters = [ - RadialCTFFilter(pixel_size=vol.pixel_size, defocus=d) + RadialCTFFilter(pixel_size=original_vol.pixel_size, defocus=d) for d in np.linspace(defocus_min, defocus_max, defocus_ct) ] @@ -72,8 +74,8 @@ # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ # We feed our ``Volume`` and filters into ``Simulation`` to generate # the dataset of images. When controlled white Gaussian noise is -# desired, ``WhiteNoiseAdder.from_snr()`` can be used to generate a -# simulation data set around a specific SNR. +# desired, ``WhiteNoiseAdder(var=VAR)`` can be used to generate a +# simulation data set around a specific noise variance. # # Alternatively, users can bring their own images using an # ``ArrayImageSource``, or define their own custom noise functions via @@ -84,28 +86,23 @@ from aspire.noise import WhiteNoiseAdder from aspire.source import Simulation -# set parameters -n_imgs = 2500 - -# SNR target for white gaussian noise. -snr = 0.5 - -# %% -# .. note:: -# Note, the SNR value was chosen based on other parameters for this -# quick tutorial, and can be changed to adjust the power of the -# additive noise. - # For this ``Simulation`` we set all 2D offset vectors to zero, # but by default offset vectors will be randomly distributed. src = Simulation( - n=n_imgs, # number of projections - vols=vol, # volume source + n=2500, # number of projections + vols=original_vol, # volume source offsets=0, # Default: images are randomly shifted unique_filters=ctf_filters, - noise_adder=WhiteNoiseAdder.from_snr(snr=snr), # desired SNR + noise_adder=WhiteNoiseAdder(var=0.0002), # desired noise variance ) +# %% +# .. note:: +# The noise variance value above was chosen based on other parameters for this +# quick tutorial, and can be changed to adjust the power of the additive noise. +# Alternatively, an SNR value can be prescribed as follows:: +# +# Simulation(..., noise_adder=WhiteNoiseAdder.from_snr(SNR)) # %% # Several Views of the Projection Images @@ -125,6 +122,24 @@ # with noise and CTF corruption src.images[0:10].show() +# %% +# Image Preprocessing +# ------------------- +# We apply some image preprocessing techniques to prepare the +# the images for denoising via Class Averaging. When processing +# experimental data additional preproccesing methods such as +# noise whitening, contrast inversion, and background normalization +# can be applied in a similar fashion. + +# %% +# Downsampling +# ------------ +# We downsample the images to remove high frequency noise and improve the +# efficiency of subsequent pipeline stages. Metadata such as pixel size is +# scaled appropriately to correspond correctly with the image resolution. + +src = src.downsample(res) +src.images[:10].show() # %% # CTF Correction @@ -132,6 +147,7 @@ # We apply ``phase_flip()`` to correct for CTF effects. src = src.phase_flip() +src.images[:10].show() # %% # Cache @@ -146,39 +162,17 @@ # %% # Class Averaging # --------------- -# We use ``RIRClass2D`` object to classify the images via the -# rotationally invariant representation (RIR) algorithm. Class -# selection is customizable. The classification module also includes a -# set of protocols for selecting a set of images to be used after -# classification. Here we're using the simplest -# ``DebugClassAvgSource`` which internally uses the -# ``TopClassSelector`` to select the first ``n_classes`` images from -# the source. In practice, the selection is done by sorting class -# averages based on some configurable notion of quality (contrast, -# neighbor distance etc). - -from aspire.classification import RIRClass2D - -# set parameters -n_classes = 200 -n_nbor = 6 - -# We will customize our class averaging source. Note that the -# ``fspca_components`` and ``bispectrum_components`` were selected for -# this small tutorial. -rir = RIRClass2D( - src, - fspca_components=40, - bispectrum_components=30, - n_nbor=n_nbor, -) +# For this tutorial we use the ``DebugClassAvgSource`` to generate an ``ImageSource`` +# of class averages. Internally, ``DebugClassAvgSource`` uses the ``RIRClass2D`` +# object to classify the source images via the rotationally invariant representation +# (RIR) algorithm and the ``TopClassSelector`` object to select the first ``n_classes`` +# images from the source. In practice, class selection is commonly done by sorting class +# averages based on some configurable notion of quality (contrast, neighbor distance etc) +# which can be accomplished by providing a custom class selector to ``ClassAverageSource``. from aspire.denoising import DebugClassAvgSource -avgs = DebugClassAvgSource( - src=src, - classifier=rir, -) +avgs = DebugClassAvgSource(src=src) # We'll continue our pipeline using only the first ``n_classes`` from # ``avgs``. The ``cache()`` call is used here to precompute results @@ -187,6 +181,7 @@ # the following ``CLSyncVoting`` algorithm. Outside of demonstration # purposes, where we are repeatedly peeking at various stage results, # such caching can be dropped allowing for more lazy evaluation. +n_classes = 250 avgs = avgs[:n_classes].cache() @@ -222,7 +217,7 @@ # For this low resolution example we will customize the ``CLSyncVoting`` # instance to use fewer theta points ``n_theta`` then the default value of 360. -orient_est = CLSyncVoting(avgs, n_theta=72) +orient_est = CLSyncVoting(avgs) # Instantiate an ``OrientedSource``. oriented_src = OrientedSource(avgs, orient_est) @@ -278,3 +273,33 @@ # For comparison, we view the first 10 source projections. src.projections[0:10].show() + + +# %% +# Fourier Shell Correlation +# ------------------------- +# Additionally, we can compare our reconstruction to the known source volume +# by performing a Fourier shell correlation (FSC). We use a Bayesian Optimal +# Transport Alignment method to align the estimated volume to the source +# volume and compute the FSC. + +# Due to the inherent handedness ambiguity involved in cryo-EM reconstructions +# we will attempt aligning both the estimated volume and a flipped volume +# to the original volume which has been downsampled to the same resolution. +from aspire.utils import Rotation, align_BO + +# Use BOT Alignment to find the best aligning rotation. +align_rot = align_BO(vol_ds, estimated_volume) + +# Align the volume and compute the FSC. +aligned_vol = estimated_volume.rotate(Rotation(align_rot[1])) +vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True) + + +# %% + +# Perform alignment and FSC on a flipped volume. +align_rot = align_BO(vol_ds, estimated_volume.flip()) +aligned_vol = estimated_volume.flip().rotate(Rotation(align_rot[1])) +vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True) + From 092cd3e8bee73723557fb9fbbdb1bcae9289e381 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 26 Jun 2025 15:52:11 -0400 Subject: [PATCH 209/216] tox --- gallery/tutorials/pipeline_demo.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index d58e17976b..4c40a00f8c 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -302,4 +302,3 @@ align_rot = align_BO(vol_ds, estimated_volume.flip()) aligned_vol = estimated_volume.flip().rotate(Rotation(align_rot[1])) vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True) - From de35baa11d9e7e12cc763ab78bc6446fa97e7f34 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 27 Jun 2025 13:54:04 -0400 Subject: [PATCH 210/216] typo/grammar --- gallery/tutorials/pipeline_demo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index 4c40a00f8c..a3cdf4c8d9 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -127,7 +127,7 @@ # ------------------- # We apply some image preprocessing techniques to prepare the # the images for denoising via Class Averaging. When processing -# experimental data additional preproccesing methods such as +# experimental data additional preprocesing methods such as # noise whitening, contrast inversion, and background normalization # can be applied in a similar fashion. @@ -136,7 +136,7 @@ # ------------ # We downsample the images to remove high frequency noise and improve the # efficiency of subsequent pipeline stages. Metadata such as pixel size is -# scaled appropriately to correspond correctly with the image resolution. +# scaled accordingly to correspond correctly with the image resolution. src = src.downsample(res) src.images[:10].show() From 8d565953d6ba0f068678f52854aea822cf42ba58 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 30 Jun 2025 10:19:46 -0400 Subject: [PATCH 211/216] Remove BOT. Align with Q_mat --- gallery/tutorials/pipeline_demo.py | 36 +++++++++++------------------- 1 file changed, 13 insertions(+), 23 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index a3cdf4c8d9..243983a13d 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -211,9 +211,10 @@ from aspire.abinitio import CLSyncVoting from aspire.source import OrientedSource +from aspire.utils import Rotation # Stash true rotations for later comparison -true_rotations = src.rotations[:n_classes] +true_rotations = Rotation(src.rotations[:n_classes]) # For this low resolution example we will customize the ``CLSyncVoting`` # instance to use fewer theta points ``n_theta`` then the default value of 360. @@ -279,26 +280,15 @@ # Fourier Shell Correlation # ------------------------- # Additionally, we can compare our reconstruction to the known source volume -# by performing a Fourier shell correlation (FSC). We use a Bayesian Optimal -# Transport Alignment method to align the estimated volume to the source -# volume and compute the FSC. - -# Due to the inherent handedness ambiguity involved in cryo-EM reconstructions -# we will attempt aligning both the estimated volume and a flipped volume -# to the original volume which has been downsampled to the same resolution. -from aspire.utils import Rotation, align_BO - -# Use BOT Alignment to find the best aligning rotation. -align_rot = align_BO(vol_ds, estimated_volume) - -# Align the volume and compute the FSC. -aligned_vol = estimated_volume.rotate(Rotation(align_rot[1])) -vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True) - - -# %% - -# Perform alignment and FSC on a flipped volume. -align_rot = align_BO(vol_ds, estimated_volume.flip()) -aligned_vol = estimated_volume.flip().rotate(Rotation(align_rot[1])) +# by performing a Fourier shell correlation (FSC). We first find a rotation +# matrix which best aligns the estimated rotations to the ground truth rotations +# using the ``find_registration`` method. We then use that rotation to align +# the reconstructed volume to the ground truth volume. + +# `find_registration` returns the best aligning rotation, `Q`, as well as +# a `flag` which indicates if the volume needs to be reflected. +Q, flag = Rotation(oriented_src.rotations).find_registration(true_rotations) +aligned_vol = estimated_volume.rotate(Rotation(Q.T)) +if flag == 1: + aligned_vol = aligned_vol.flip() vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True) From d4a42c65d2932f84aa4afd780c75dfaacb53e413 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Jul 2025 08:48:35 -0400 Subject: [PATCH 212/216] Add additional preprocessing --- gallery/tutorials/pipeline_demo.py | 37 +++++++++++++++++++++++------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index 243983a13d..92ed9b8282 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -126,10 +126,7 @@ # Image Preprocessing # ------------------- # We apply some image preprocessing techniques to prepare the -# the images for denoising via Class Averaging. When processing -# experimental data additional preprocesing methods such as -# noise whitening, contrast inversion, and background normalization -# can be applied in a similar fashion. +# the images for denoising via Class Averaging. # %% # Downsampling @@ -150,14 +147,38 @@ src.images[:10].show() # %% -# Cache -# ----- +# Normalize Background +# -------------------- +# We apply ``normalize_background()`` to prepare the image class averaging. + +src = src.normalize_background() +src.images[:10].show() + +# %% +# Caching +# ------- # We apply ``cache`` to store the results of the ``ImageSource`` # pipeline up to this point. This is optional, but can provide # benefit when used intently on machines with adequate memory. - +# Since the remaining preprocessing steps, whitening and contrast +# inversion, as well as class averaging require passing over the +# full set of images, caching at this point saves compute time. src = src.cache() -src.images[0:10].show() + +# %% +# Noise Whitening +# --------------- +# We apply ``whiten()`` to estimate and whiten the noise. + +src = src.whiten() +src.images[:10].show() + +# %% +# Contrast Inversion +# ------------------ +# We apply ``invert_contrast()`` to ensure a positive valued signal. + +src = src.invert_contrast() # %% # Class Averaging From 9650b9063c4092e26fb6969e3d62d78fcc6ebd45 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Jul 2025 12:13:17 -0400 Subject: [PATCH 213/216] Comment cleanup --- gallery/tutorials/pipeline_demo.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index 92ed9b8282..904d68fb84 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -131,7 +131,7 @@ # %% # Downsampling # ------------ -# We downsample the images to remove high frequency noise and improve the +# We downsample the images. Reducing the image size will improve the # efficiency of subsequent pipeline stages. Metadata such as pixel size is # scaled accordingly to correspond correctly with the image resolution. @@ -187,9 +187,11 @@ # of class averages. Internally, ``DebugClassAvgSource`` uses the ``RIRClass2D`` # object to classify the source images via the rotationally invariant representation # (RIR) algorithm and the ``TopClassSelector`` object to select the first ``n_classes`` -# images from the source. In practice, class selection is commonly done by sorting class -# averages based on some configurable notion of quality (contrast, neighbor distance etc) -# which can be accomplished by providing a custom class selector to ``ClassAverageSource``. +# images in the original order from the source. In practice, class selection is commonly +# done by sorting class averages based on some configurable notion of quality +# (contrast, neighbor distance etc) which can be accomplished by providing a custom +# class selector to ``ClassAverageSource``, which changes the ordering of the classes +# returned by ``ClassAverageSource``. from aspire.denoising import DebugClassAvgSource @@ -237,8 +239,7 @@ # Stash true rotations for later comparison true_rotations = Rotation(src.rotations[:n_classes]) -# For this low resolution example we will customize the ``CLSyncVoting`` -# instance to use fewer theta points ``n_theta`` then the default value of 360. +# Instantiate a ``CLSyncVoting`` object for estimating orientations. orient_est = CLSyncVoting(avgs) # Instantiate an ``OrientedSource``. From f067b3cb9f154d898abf29f8f55c696bd3cbb6cc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 1 Jul 2025 14:44:31 -0400 Subject: [PATCH 214/216] Cache at every stage. Fix alignment flip/rotate order --- gallery/tutorials/pipeline_demo.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index 904d68fb84..cb7a417a2a 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -88,13 +88,15 @@ # For this ``Simulation`` we set all 2D offset vectors to zero, # but by default offset vectors will be randomly distributed. +# We cache the Simulation to prevent regenerating the projections +# for each preproccesing stage. src = Simulation( n=2500, # number of projections vols=original_vol, # volume source offsets=0, # Default: images are randomly shifted unique_filters=ctf_filters, noise_adder=WhiteNoiseAdder(var=0.0002), # desired noise variance -) +).cache() # %% # .. note:: @@ -135,7 +137,7 @@ # efficiency of subsequent pipeline stages. Metadata such as pixel size is # scaled accordingly to correspond correctly with the image resolution. -src = src.downsample(res) +src = src.downsample(res).cache() src.images[:10].show() # %% @@ -143,7 +145,7 @@ # -------------- # We apply ``phase_flip()`` to correct for CTF effects. -src = src.phase_flip() +src = src.phase_flip().cache() src.images[:10].show() # %% @@ -151,26 +153,15 @@ # -------------------- # We apply ``normalize_background()`` to prepare the image class averaging. -src = src.normalize_background() +src = src.normalize_background().cache() src.images[:10].show() -# %% -# Caching -# ------- -# We apply ``cache`` to store the results of the ``ImageSource`` -# pipeline up to this point. This is optional, but can provide -# benefit when used intently on machines with adequate memory. -# Since the remaining preprocessing steps, whitening and contrast -# inversion, as well as class averaging require passing over the -# full set of images, caching at this point saves compute time. -src = src.cache() - # %% # Noise Whitening # --------------- # We apply ``whiten()`` to estimate and whiten the noise. -src = src.whiten() +src = src.whiten().cache() src.images[:10].show() # %% @@ -178,7 +169,7 @@ # ------------------ # We apply ``invert_contrast()`` to ensure a positive valued signal. -src = src.invert_contrast() +src = src.invert_contrast().cache() # %% # Class Averaging @@ -310,7 +301,10 @@ # `find_registration` returns the best aligning rotation, `Q`, as well as # a `flag` which indicates if the volume needs to be reflected. Q, flag = Rotation(oriented_src.rotations).find_registration(true_rotations) -aligned_vol = estimated_volume.rotate(Rotation(Q.T)) +aligned_vol = estimated_volume if flag == 1: aligned_vol = aligned_vol.flip() +aligned_vol = aligned_vol.rotate(Rotation(Q.T)) + +# Compute the FSC. vol_ds.fsc(aligned_vol, cutoff=0.143, plot=True) From 990d666bdbc74759c528834c183ac4c442c21739 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 2 Jul 2025 08:38:14 -0400 Subject: [PATCH 215/216] cache sim, post preprocessing, post averages --- gallery/tutorials/pipeline_demo.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index cb7a417a2a..df42849cc7 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -89,7 +89,7 @@ # For this ``Simulation`` we set all 2D offset vectors to zero, # but by default offset vectors will be randomly distributed. # We cache the Simulation to prevent regenerating the projections -# for each preproccesing stage. +# for each preprocessing stage. src = Simulation( n=2500, # number of projections vols=original_vol, # volume source @@ -137,7 +137,7 @@ # efficiency of subsequent pipeline stages. Metadata such as pixel size is # scaled accordingly to correspond correctly with the image resolution. -src = src.downsample(res).cache() +src = src.downsample(res) src.images[:10].show() # %% @@ -145,7 +145,7 @@ # -------------- # We apply ``phase_flip()`` to correct for CTF effects. -src = src.phase_flip().cache() +src = src.phase_flip() src.images[:10].show() # %% @@ -153,7 +153,7 @@ # -------------------- # We apply ``normalize_background()`` to prepare the image class averaging. -src = src.normalize_background().cache() +src = src.normalize_background() src.images[:10].show() # %% @@ -161,7 +161,7 @@ # --------------- # We apply ``whiten()`` to estimate and whiten the noise. -src = src.whiten().cache() +src = src.whiten() src.images[:10].show() # %% @@ -169,7 +169,15 @@ # ------------------ # We apply ``invert_contrast()`` to ensure a positive valued signal. -src = src.invert_contrast().cache() +src = src.invert_contrast() + +# %% +# Caching +# ------- +# We apply ``cache`` to store the results of the ``ImageSource`` +# pipeline up to this point. This is optional, but can provide +# benefit when used intently on machines with adequate memory. +src = src.cache() # %% # Class Averaging From 60d1db661c10f0d7167c713a31a622224aa2c75e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 2 Jul 2025 10:52:14 -0400 Subject: [PATCH 216/216] fix sphinx cell ignore. estimate rots in Orientation Estimation block. --- gallery/tutorials/pipeline_demo.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index df42849cc7..ad6143c669 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -14,9 +14,11 @@ # Ribosome, sourced from EMDB: https://www.ebi.ac.uk/emdb/EMD-2660. # This is one of several volume maps that can be downloaded with # ASPIRE's data downloading utility by using the following import. + # sphinx_gallery_start_ignore # flake8: noqa # sphinx_gallery_end_ignore + from aspire.downloader import emdb_2660 # Load 80s Ribosome as a ``Volume`` object. @@ -244,6 +246,8 @@ # Instantiate an ``OrientedSource``. oriented_src = OrientedSource(avgs, orient_est) +# Estimate Rotations. +est_rotations = oriented_src.rotations # %% # Mean Error of Estimated Rotations @@ -255,7 +259,7 @@ from aspire.utils import mean_aligned_angular_distance # Compare with known true rotations -mean_ang_dist = mean_aligned_angular_distance(oriented_src.rotations, true_rotations) +mean_ang_dist = mean_aligned_angular_distance(est_rotations, true_rotations) print(f"Mean aligned angular distance: {mean_ang_dist} degrees") @@ -286,7 +290,7 @@ # Get the first 10 projections from the estimated volume using the # estimated orientations. Recall that ``project`` returns an # ``Image`` instance, which we can peek at with ``show``. -projections_est = estimated_volume.project(oriented_src.rotations[0:10]) +projections_est = estimated_volume.project(est_rotations[0:10]) # We view the first 10 projections of the estimated volume. projections_est.show() @@ -308,7 +312,7 @@ # `find_registration` returns the best aligning rotation, `Q`, as well as # a `flag` which indicates if the volume needs to be reflected. -Q, flag = Rotation(oriented_src.rotations).find_registration(true_rotations) +Q, flag = Rotation(est_rotations).find_registration(true_rotations) aligned_vol = estimated_volume if flag == 1: aligned_vol = aligned_vol.flip()