Merge branch 'develop' into ds_logical_tests

Author: chris-langfield

gallery/tutorials/lecture_feature_demo.py

@@ -42,8 +42,8 @@
 from aspire.noise import AnisotropicNoiseEstimator, WhiteNoiseEstimator
 from aspire.operators import FunctionFilter, RadialCTFFilter, ScalarFilter
 from aspire.source import RelionSource, Simulation
-from aspire.utils import Rotation
-from aspire.utils.coor_trans import (
+from aspire.utils import (
+    Rotation,
     get_aligned_rotations,
     get_rots_mse,
     register_rotations,

gallery/tutorials/orient3d_simulation.py

@@ -15,11 +15,7 @@
 from aspire.abinitio import CLSyncVoting
 from aspire.operators import RadialCTFFilter
 from aspire.source.simulation import Simulation
-from aspire.utils.coor_trans import (
-    get_aligned_rotations,
-    get_rots_mse,
-    register_rotations,
-)
+from aspire.utils import get_aligned_rotations, get_rots_mse, register_rotations
 from aspire.volume import Volume
 
 logger = logging.getLogger(__name__)

src/aspire/abinitio/commonline_base.py

@@ -6,7 +6,7 @@
 
 from aspire.abinitio.orientation_src import OrientEstSource
 from aspire.basis import PolarBasis2D
-from aspire.utils.coor_trans import common_line_from_rots
+from aspire.utils import common_line_from_rots
 from aspire.utils.random import choice
 
 logger = logging.getLogger(__name__)

src/aspire/basis/basis_utils.py

@@ -10,7 +10,7 @@
 from numpy.polynomial.legendre import leggauss
 from scipy.special import jn, jv, sph_harm
 
-from aspire.utils.coor_trans import grid_2d, grid_3d
+from aspire.utils import grid_2d, grid_3d
 
 logger = logging.getLogger(__name__)
 

src/aspire/classification/averager2d.py

@@ -466,6 +466,8 @@ def __init__(
         if self.alignment_src.dtype != src.dtype:
             raise RuntimeError("Currently `alignment_src.dtype` must equal `src.dtype`")
 
+        self.mask = grid_2d(src.L, normalized=False)["r"] < src.L // 2
+
         super().__init__(composite_basis, src, composite_basis, dtype=dtype)
 
     def _phase_cross_correlation(self, img0, img1):
@@ -547,7 +549,7 @@ def _reddychatterji(self, images, class_k, reflection_k):
         # Result arrays
         M = len(images)
         rotations_k = np.zeros(M, dtype=self.dtype)
-        correlations_k = np.zeros(M, dtype=self.dtype)
+        correlations_k = np.full(M, -np.inf, dtype=self.dtype)
         shifts_k = np.zeros((M, 2), dtype=int)
 
         # De-Mean, note images is mutated and should be a `copy`.
@@ -655,8 +657,8 @@ def _reddychatterji(self, images, class_k, reflection_k):
             # Hack
             regis_img_estimated = rotate(regis_img, r)
             regis_img_rotated_p180 = rotate(regis_img, r + 180)
-            da = np.dot(fixed_img.flatten(), regis_img_estimated.flatten())
-            db = np.dot(fixed_img.flatten(), regis_img_rotated_p180.flatten())
+            da = np.dot(fixed_img[self.mask], regis_img_estimated[self.mask])
+            db = np.dot(fixed_img[self.mask], regis_img_rotated_p180[self.mask])
             if db > da:
                 regis_img_estimated = regis_img_rotated_p180
                 r += 180
@@ -702,7 +704,7 @@ def _reddychatterji(self, images, class_k, reflection_k):
                 shift = None  # For logger line
 
             # Estimated `corr` metric
-            corr = np.dot(fixed_img.flatten(), regis_img_estimated.flatten())
+            corr = np.dot(fixed_img[self.mask], regis_img_estimated[self.mask])
             correlations_k[m] = corr
 
             logger.debug(
@@ -962,9 +964,7 @@ def align(self, classes, reflections, basis_coefficients):
         L = self.alignment_src.L
 
         # Instantiate matrices for inner loop, and best results.
-        _rotations = np.zeros(classes.shape, dtype=self.dtype)
         rotations = np.zeros(classes.shape, dtype=self.dtype)
-        _correlations = np.zeros(classes.shape, dtype=self.dtype)
         correlations = np.ones(classes.shape, dtype=self.dtype) * -np.inf
         shifts = np.zeros((*classes.shape, 2), dtype=int)
 
@@ -986,16 +986,16 @@ def align(self, classes, reflections, basis_coefficients):
                 # Don't shift the base image
                 images[1:] = Image(unshifted_images[1:]).shift(s).asnumpy()
 
-                rotations[k], _, correlations[k] = self._reddychatterji(
+                _rotations, _, _correlations = self._reddychatterji(
                     images, classes[k], reflections[k]
                 )
 
                 # Where corr has improved
                 #  update our rolling best results with this loop.
-                improved = _correlations > correlations
-                correlations = np.where(improved, _correlations, correlations)
-                rotations = np.where(improved, _rotations, rotations)
-                shifts = np.where(improved[..., np.newaxis], s, shifts)
+                improved = _correlations > correlations[k]
+                correlations[k] = np.where(improved, _correlations, correlations[k])
+                rotations[k] = np.where(improved, _rotations, rotations[k])
+                shifts[k] = np.where(improved[..., np.newaxis], s, shifts[k])
                 logger.debug(f"Shift {s} has improved {np.sum(improved)} results")
 
         return rotations, shifts, correlations

src/aspire/covariance/covar.py

@@ -42,9 +42,9 @@ def __getattr__(self, name):
 
     def compute_kernel(self):
         # TODO: Most of this stuff is duplicated in MeanEstimator - move up the hierarchy?
-        n = self.n
-        L = self.L
-        _2L = 2 * self.L
+        n = self.src.n
+        L = self.src.L
+        _2L = 2 * self.src.L
 
         kernel = np.zeros((_2L, _2L, _2L, _2L, _2L, _2L), dtype=self.dtype)
         sq_filters_f = np.square(evaluate_src_filters_on_grid(self.src))
@@ -168,22 +168,23 @@ def src_backward(self, mean_vol, noise_variance, shrink_method=None):
         contribution and expressed as coefficients of `basis`.
         """
         covar_b = np.zeros(
-            (self.L, self.L, self.L, self.L, self.L, self.L), dtype=self.dtype
+            (self.src.L, self.src.L, self.src.L, self.src.L, self.src.L, self.src.L),
+            dtype=self.dtype,
         )
 
-        for i in range(0, self.n, self.batch_size):
+        for i in range(0, self.src.n, self.batch_size):
             im = self.src.images(i, self.batch_size)
             batch_n = im.n_images
             im_centered = im - self.src.vol_forward(mean_vol, i, self.batch_size)
 
             im_centered_b = np.zeros(
-                (batch_n, self.L, self.L, self.L), dtype=self.dtype
+                (batch_n, self.src.L, self.src.L, self.src.L), dtype=self.dtype
             )
             for j in range(batch_n):
                 im_centered_b[j] = self.src.im_backward(Image(im_centered[j]), i + j)
             im_centered_b = Volume(im_centered_b).to_vec()
 
-            covar_b += vecmat_to_volmat(im_centered_b.T @ im_centered_b) / self.n
+            covar_b += vecmat_to_volmat(im_centered_b.T @ im_centered_b) / self.src.n
 
         covar_b_coeff = self.basis.mat_evaluate_t(covar_b)
         return self._shrink(covar_b_coeff, noise_variance, shrink_method)

src/aspire/ctf/ctf_estimator.py

@@ -21,8 +21,7 @@
 from aspire.numeric import fft
 from aspire.operators import voltage_to_wavelength
 from aspire.storage import StarFile
-from aspire.utils import abs2, complex_type
-from aspire.utils.coor_trans import grid_1d, grid_2d
+from aspire.utils import abs2, complex_type, grid_1d, grid_2d
 
 logger = logging.getLogger(__name__)
 

src/aspire/denoising/adaptive_support.py

@@ -5,7 +5,7 @@
 from aspire.noise import WhiteNoiseEstimator
 from aspire.numeric import fft
 from aspire.source import ImageSource
-from aspire.utils.coor_trans import grid_2d
+from aspire.utils import grid_2d
 
 logger = logging.getLogger(__name__)
 

src/aspire/image/image.py

@@ -9,7 +9,7 @@
 import aspire.volume
 from aspire.nufft import anufft
 from aspire.numeric import fft, xp
-from aspire.utils.coor_trans import grid_2d
+from aspire.utils import grid_2d
 from aspire.utils.matrix import anorm
 
 logger = logging.getLogger(__name__)

src/aspire/noise/noise.py

@@ -4,7 +4,7 @@
 
 from aspire.numeric import fft, xp
 from aspire.operators import ArrayFilter, ScalarFilter
-from aspire.utils.coor_trans import grid_2d
+from aspire.utils import grid_2d
 
 logger = logging.getLogger(__name__)
 
@@ -27,8 +27,6 @@ def __init__(self, src, bgRadius=1, batchSize=512):
 
         self.src = src
         self.dtype = self.src.dtype
-        self.L = src.L
-        self.n = src.n
         self.bgRadius = bgRadius
         self.batchSize = batchSize
 
@@ -72,16 +70,16 @@ def _estimate_noise_variance(self):
         TODO: How's this initial estimate of variance different from the 'estimate' method?
         """
         # Run estimate using saved parameters
-        g2d = grid_2d(self.L, indexing="yx", dtype=self.dtype)
+        g2d = grid_2d(self.src.L, indexing="yx", dtype=self.dtype)
         mask = g2d["r"] >= self.bgRadius
 
         first_moment = 0
         second_moment = 0
-        for i in range(0, self.n, self.batchSize):
+        for i in range(0, self.src.n, self.batchSize):
             images = self.src.images(start=i, num=self.batchSize).asnumpy()
             images_masked = images * mask
 
-            _denominator = self.n * np.sum(mask)
+            _denominator = self.src.n * np.sum(mask)
             first_moment += np.sum(images_masked) / _denominator
             second_moment += np.sum(np.abs(images_masked**2)) / _denominator
         return second_moment - first_moment**2
@@ -100,7 +98,7 @@ def estimate(self):
         # AnisotropicNoiseEstimator.filter is an ArrayFilter.
         #   We average the variance over all frequencies,
 
-        return np.mean(self.filter.evaluate_grid(self.L))
+        return np.mean(self.filter.evaluate_grid(self.src.L))
 
     def _create_filter(self, noise_psd=None):
         """
@@ -117,21 +115,21 @@ def estimate_noise_psd(self):
         TODO: How's this initial estimate of variance different from the 'estimate' method?
         """
         # Run estimate using saved parameters
-        g2d = grid_2d(self.L, indexing="yx", dtype=self.dtype)
+        g2d = grid_2d(self.src.L, indexing="yx", dtype=self.dtype)
         mask = g2d["r"] >= self.bgRadius
 
         mean_est = 0
-        noise_psd_est = np.zeros((self.L, self.L)).astype(self.src.dtype)
-        for i in range(0, self.n, self.batchSize):
+        noise_psd_est = np.zeros((self.src.L, self.src.L)).astype(self.src.dtype)
+        for i in range(0, self.src.n, self.batchSize):
             images = self.src.images(i, self.batchSize).asnumpy()
             images_masked = images * mask
 
-            _denominator = self.n * np.sum(mask)
+            _denominator = self.src.n * np.sum(mask)
             mean_est += np.sum(images_masked) / _denominator
             im_masked_f = xp.asnumpy(fft.centered_fft2(xp.asarray(images_masked)))
             noise_psd_est += np.sum(np.abs(im_masked_f**2), axis=0) / _denominator
 
-        mid = self.L // 2
+        mid = self.src.L // 2
         noise_psd_est[mid, mid] -= mean_est**2
 
         return noise_psd_est