diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index fe91b322ad..3685479732 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -1,6 +1,12 @@ from .commonline_base import CLOrient3D # isort: off +from .commonline_utils import ( + cl_angles_to_ind, + estimate_third_rows, + complete_third_row_to_rot, + estimate_inplane_rotations, +) from .commonline_sdp import CommonlineSDP from .commonline_lud import CommonlineLUD from .commonline_irls import CommonlineIRLS diff --git a/src/aspire/abinitio/commonline_c2.py b/src/aspire/abinitio/commonline_c2.py index 6bf6cf99ed..52b63cc9e8 100644 --- a/src/aspire/abinitio/commonline_c2.py +++ b/src/aspire/abinitio/commonline_c2.py @@ -3,7 +3,11 @@ import numpy as np from scipy.linalg import eigh -from aspire.abinitio import CLSymmetryC3C4 +from aspire.abinitio import ( + CLSymmetryC3C4, + complete_third_row_to_rot, + estimate_third_rows, +) from aspire.utils import J_conjugate, Rotation, all_pairs logger = logging.getLogger(__name__) @@ -224,7 +228,7 @@ def estimate_rotations(self): viis = np.vstack((np.eye(3, dtype=self.dtype),) * self.n_img).reshape( self.n_img, 3, 3 ) - vis = self._estimate_third_rows(vijs, viis) + vis = estimate_third_rows(vijs, viis) logger.info("Estimating in-plane rotations and rotations matrices.") Ris = self._estimate_inplane_rotations(vis, Rijs, Rijgs) @@ -301,7 +305,7 @@ def _estimate_inplane_rotations(self, vis, Rijs, Rijgs): H = np.zeros((self.n_img, self.n_img), dtype=complex) # Step 1: Construct all rotation matrices Ris_tilde whose third rows are equal to # the corresponding third rows vis. - Ris_tilde = self._complete_third_row_to_rot(vis) + Ris_tilde = complete_third_row_to_rot(vis) pairs = all_pairs(self.n_img) for idx, (i, j) in enumerate(pairs): diff --git a/src/aspire/abinitio/commonline_c3_c4.py b/src/aspire/abinitio/commonline_c3_c4.py index 0170d8b88f..af74dc6f6c 100644 --- a/src/aspire/abinitio/commonline_c3_c4.py +++ b/src/aspire/abinitio/commonline_c3_c4.py @@ -3,7 +3,12 @@ import numpy as np from numpy.linalg import eigh, norm, svd -from aspire.abinitio import CLOrient3D, SyncVotingMixin +from aspire.abinitio import ( + CLOrient3D, + SyncVotingMixin, + estimate_inplane_rotations, + estimate_third_rows, +) from aspire.operators import PolarFT from aspire.utils import ( J_conjugate, @@ -12,7 +17,6 @@ all_triplets, anorm, cyclic_rotations, - tqdm, trange, ) from aspire.utils.random import randn @@ -110,10 +114,10 @@ def estimate_rotations(self): vijs, viis = self._global_J_sync(vijs, viis) logger.info("Estimating third rows of rotation matrices.") - vis = self._estimate_third_rows(vijs, viis) + vis = estimate_third_rows(vijs, viis) logger.info("Estimating in-plane rotations and rotations matrices.") - Ris = self._estimate_inplane_rotations(vis) + Ris = estimate_inplane_rotations(self, vis) self.rotations = Ris @@ -209,189 +213,6 @@ def _global_J_sync(self, vijs, viis): viis[i] = vii_J return vijs, viis - def _estimate_third_rows(self, vijs, viis): - """ - Find the third row of each rotation matrix given a collection of matrices - representing the outer products of the third rows from each rotation matrix. - - :param vijs: An (n-choose-2)x3x3 array where each 3x3 slice holds the third rows - outer product of the rotation matrices Ri and Rj. - - :param viis: An n_imgx3x3 array where the i'th 3x3 slice holds the outer product of - the third row of Ri with itself. - - :param order: The underlying molecular symmetry. - - :return: vis, An n_imgx3 matrix whose i'th row is the third row of the rotation matrix Ri. - """ - - n_img = self.n_img - - # Build matrix V whose (i,j)-th block of size 3x3 holds the outer product vij - V = np.zeros((n_img, n_img, 3, 3), dtype=vijs.dtype) - - # All pairs (i,j) where i= 1e-5 - - # If the third row coincides with the z-axis we return the identity matrix. - rots[~mask] = np.eye(3, dtype=r3.dtype) - - # 'norm_12' is non-zero since r3 does not coincide with the z-axis. - norm_12 = np.sqrt(r3[mask, 0] ** 2 + r3[mask, 1] ** 2) - - # Populate 1st rows with vector orthogonal to row 3. - rots[mask, 0, 0] = r3[mask, 1] / norm_12 - rots[mask, 0, 1] = -r3[mask, 0] / norm_12 - - # Populate 2nd rows such that r3 = r1 x r2 - rots[mask, 1, 0] = r3[mask, 0] * r3[mask, 2] / norm_12 - rots[mask, 1, 1] = r3[mask, 1] * r3[mask, 2] / norm_12 - rots[mask, 1, 2] = -norm_12 - - if singleton: - rots = rots.reshape(3, 3) - - return rots - - @staticmethod - def cl_angles_to_ind(cl_angles, n_theta): - thetas = np.arctan2(cl_angles[:, 1], cl_angles[:, 0]) - - # Shift from [-pi,pi] to [0,2*pi). - thetas = np.mod(thetas, 2 * np.pi) - - # linear scale from [0,2*pi) to [0,n_theta). - ind = np.mod(np.round(thetas / (2 * np.pi) * n_theta), n_theta).astype(int) - - # Return scalar for single value. - if ind.size == 1: - ind = ind.flat[0] - - return ind - @staticmethod def g_sync(rots, order, rots_gt): """ diff --git a/src/aspire/abinitio/commonline_cn.py b/src/aspire/abinitio/commonline_cn.py index a1440cadc3..e0d69cb086 100644 --- a/src/aspire/abinitio/commonline_cn.py +++ b/src/aspire/abinitio/commonline_cn.py @@ -3,7 +3,7 @@ import numpy as np from numpy.linalg import norm -from aspire.abinitio import CLSymmetryC3C4 +from aspire.abinitio import CLSymmetryC3C4, cl_angles_to_ind, complete_third_row_to_rot from aspire.operators import PolarFT from aspire.utils import ( J_conjugate, @@ -298,8 +298,8 @@ def relative_rots_to_cl_indices(relative_rots, n_theta): c1s = np.array((-relative_rots[:, 1, 2], relative_rots[:, 0, 2])).T c2s = np.array((relative_rots[:, 2, 1], -relative_rots[:, 2, 0])).T - c1s = CLSymmetryC3C4.cl_angles_to_ind(c1s, n_theta) - c2s = CLSymmetryC3C4.cl_angles_to_ind(c2s, n_theta) + c1s = cl_angles_to_ind(c1s, n_theta) + c2s = cl_angles_to_ind(c2s, n_theta) inds = np.where(c1s >= n_theta // 2) c1s[inds] -= n_theta // 2 @@ -331,7 +331,7 @@ def generate_candidate_rots(n, equator_threshold, order, degree_res, seed): while counter < n: third_row = randn(3) third_row /= anorm(third_row, axes=(-1,)) - Ri_tilde = CLSymmetryC3C4._complete_third_row_to_rot(third_row) + Ri_tilde = complete_third_row_to_rot(third_row) # Exclude candidates that represent equator images. Equator candidates # induce collinear self-common-lines, which always have perfect correlation. diff --git a/src/aspire/abinitio/commonline_sync3n.py b/src/aspire/abinitio/commonline_sync3n.py index 841f9dfceb..25bfd63015 100644 --- a/src/aspire/abinitio/commonline_sync3n.py +++ b/src/aspire/abinitio/commonline_sync3n.py @@ -7,15 +7,7 @@ from scipy.optimize import curve_fit from aspire.abinitio import CLOrient3D, SyncVotingMixin -from aspire.utils import ( - J_conjugate, - Rotation, - all_pairs, - nearest_rotations, - random, - tqdm, - trange, -) +from aspire.utils import J_conjugate, all_pairs, nearest_rotations, random, tqdm, trange from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) @@ -969,38 +961,6 @@ def _estimate_all_Rijs_host(self, clmatrix): return Rijs - def _syncmatrix_ij_vote_3n(self, clmatrix, i, j, k_list, n_theta): - """ - Compute the (i,j) rotation block of the synchronization matrix using voting method - - Given the common lines matrix `clmatrix`, a list of images specified in k_list - and the number of common lines n_theta, find the (i, j) rotation block Rij. - - :param clmatrix: The common lines matrix - :param i: The i image - :param j: The j image - :param k_list: The list of images for the third image for voting algorithm - :param n_theta: The number of points in the theta direction (common lines) - :return: The (i,j) rotation block of the synchronization matrix - """ - alphas, good_k = self._vote_ij(clmatrix, n_theta, i, j, k_list, sync=True) - - angles = np.zeros(3) - - if alphas is not None: - angles[0] = clmatrix[i, j] * 2 * np.pi / n_theta + np.pi / 2 - angles[1] = np.mean(alphas) - angles[2] = -np.pi / 2 - clmatrix[j, i] * 2 * np.pi / n_theta - rot = Rotation.from_euler(angles).matrices - - else: - # This is for the case that images i and j correspond to the same - # viewing direction and differ only by in-plane rotation. - # We set to zero as in the Matlab code. - rot = np.zeros((3, 3)) - - return rot - ####################################### # Secondary Methods for Global J Sync # ####################################### diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py new file mode 100644 index 0000000000..bfd1e5ce5f --- /dev/null +++ b/src/aspire/abinitio/commonline_utils.py @@ -0,0 +1,253 @@ +import logging + +import numpy as np +from numpy.linalg import eigh, norm + +from aspire.operators import PolarFT +from aspire.utils import Rotation, all_pairs, anorm, tqdm + +logger = logging.getLogger(__name__) + + +def estimate_third_rows(vijs, viis): + """ + Find the third row of each rotation matrix given a collection of matrices + representing the outer products of the third rows from each rotation matrix. + + :param vijs: An (n-choose-2)x3x3 array where each 3x3 slice holds the third rows + outer product of the rotation matrices Ri and Rj. + + :param viis: An n_imgx3x3 array where the i'th 3x3 slice holds the outer product of + the third row of Ri with itself. + + :param order: The underlying molecular symmetry. + + :return: vis, An n_imgx3 matrix whose i'th row is the third row of the rotation matrix Ri. + """ + + n_img = viis.shape[0] + + # Build matrix V whose (i,j)-th block of size 3x3 holds the outer product vij + V = np.zeros((n_img, n_img, 3, 3), dtype=vijs.dtype) + + # All pairs (i,j) where i= 1e-5 + + # If the third row coincides with the z-axis we return the identity matrix. + rots[~mask] = np.eye(3, dtype=r3.dtype) + + # 'norm_12' is non-zero since r3 does not coincide with the z-axis. + norm_12 = np.sqrt(r3[mask, 0] ** 2 + r3[mask, 1] ** 2) + + # Populate 1st rows with vector orthogonal to row 3. + rots[mask, 0, 0] = r3[mask, 1] / norm_12 + rots[mask, 0, 1] = -r3[mask, 0] / norm_12 + + # Populate 2nd rows such that r3 = r1 x r2 + rots[mask, 1, 0] = r3[mask, 0] * r3[mask, 2] / norm_12 + rots[mask, 1, 1] = r3[mask, 1] * r3[mask, 2] / norm_12 + rots[mask, 1, 2] = -norm_12 + + if singleton: + rots = rots.reshape(3, 3) + + return rots + + +def cl_angles_to_ind(cl_angles, n_theta): + thetas = np.arctan2(cl_angles[:, 1], cl_angles[:, 0]) + + # Shift from [-pi,pi] to [0,2*pi). + thetas = np.mod(thetas, 2 * np.pi) + + # linear scale from [0,2*pi) to [0,n_theta). + ind = np.mod(np.round(thetas / (2 * np.pi) * n_theta), n_theta).astype(int) + + # Return scalar for single value. + if ind.size == 1: + ind = ind.flat[0] + + return ind diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index 651ba35c5e..b866f1e8ff 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -13,6 +13,38 @@ class SyncVotingMixin(object): which are shared by CLSynVoting and CLSymmetryC3C4 """ + def _syncmatrix_ij_vote_3n(self, clmatrix, i, j, k_list, n_theta): + """ + Compute the (i,j) rotation block of the synchronization matrix using voting method + + Given the common lines matrix `clmatrix`, a list of images specified in k_list + and the number of common lines n_theta, find the (i, j) rotation block Rij. + + :param clmatrix: The common lines matrix + :param i: The i image + :param j: The j image + :param k_list: The list of images for the third image for voting algorithm + :param n_theta: The number of points in the theta direction (common lines) + :return: The (i,j) rotation block of the synchronization matrix + """ + alphas, good_k = self._vote_ij(clmatrix, n_theta, i, j, k_list, sync=True) + + angles = np.zeros(3) + + if alphas is not None: + angles[0] = clmatrix[i, j] * 2 * np.pi / n_theta + np.pi / 2 + angles[1] = np.mean(alphas) + angles[2] = -np.pi / 2 - clmatrix[j, i] * 2 * np.pi / n_theta + rot = Rotation.from_euler(angles).matrices + + else: + # This is for the case that images i and j correspond to the same + # viewing direction and differ only by in-plane rotation. + # We set to zero as in the Matlab code. + rot = np.zeros((3, 3)) + + return rot + def _rotratio_eulerangle_vec(self, clmatrix, i, j, good_k, n_theta): """ Compute the rotation that takes image i to image j diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index ae896ebeb8..ae781d823f 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -1,4 +1,5 @@ from .types import complex_type, real_type, utest_tolerance # isort:skip + from .coor_trans import ( # isort:skip mean_aligned_angular_distance, cart2pol, diff --git a/tests/test_orient_symmetric.py b/tests/test_orient_symmetric.py index 9ab63459df..d7c4c4716f 100644 --- a/tests/test_orient_symmetric.py +++ b/tests/test_orient_symmetric.py @@ -3,7 +3,14 @@ from numpy import pi, random from numpy.linalg import det, norm -from aspire.abinitio import CLSymmetryC2, CLSymmetryC3C4, CLSymmetryCn +from aspire.abinitio import ( + CLSymmetryC2, + CLSymmetryC3C4, + CLSymmetryCn, + cl_angles_to_ind, + complete_third_row_to_rot, + estimate_third_rows, +) from aspire.abinitio.commonline_cn import MeanOuterProductEstimator from aspire.source import Simulation from aspire.utils import ( @@ -479,17 +486,14 @@ def test_global_J_sync(n_img, dtype): @pytest.mark.parametrize("dtype", [np.float32, np.float64]) def test_estimate_third_rows(dtype): - L = 16 n_img = 20 - order = 3 # test not dependent on order - _, orient_est = source_orientation_objs(n_img, L, order, dtype) # Build outer products vijs, viis, and get ground truth third rows. vijs, viis, gt_vis = build_outer_products(n_img, dtype) # Estimate third rows from outer products. # Due to factorization of V, these might be negated third rows. - vis = orient_est._estimate_third_rows(vijs, viis) + vis = estimate_third_rows(vijs, viis) # Check if all-close up to difference of sign ground_truth = np.sign(gt_vis[0, 0]) * gt_vis @@ -507,7 +511,7 @@ def test_complete_third_row(dtype): r3[0] = np.array([0, 0, 1], dtype=dtype) # Generate rotations. - R = CLSymmetryC3C4._complete_third_row_to_rot(r3) + R = complete_third_row_to_rot(r3) # Assert that first rotation is the identity matrix. assert np.allclose(R[0], np.eye(3, dtype=dtype)) @@ -638,10 +642,6 @@ def _gt_cl_c2(n_theta, rots_gt): U = Ri.T @ g @ Rj c1 = np.array([-U[1, 2], U[0, 2]]) c2 = np.array([U[2, 1], -U[2, 0]]) - clmatrix_gt[idx, i, j] = CLSymmetryC3C4.cl_angles_to_ind( - c1[np.newaxis, :], n_theta - ) - clmatrix_gt[idx, j, i] = CLSymmetryC3C4.cl_angles_to_ind( - c2[np.newaxis, :], n_theta - ) + clmatrix_gt[idx, i, j] = cl_angles_to_ind(c1[np.newaxis, :], n_theta) + clmatrix_gt[idx, j, i] = cl_angles_to_ind(c2[np.newaxis, :], n_theta) return clmatrix_gt