Change self.L and self.n to self.src.L and self.src.n for estimators (#550)

* changed self.L and self.n to self.src.L and self.src.n for estimators

* extraneous comma

Author: chris-langfield

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/noise/noise.py

@@ -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

src/aspire/reconstruction/estimator.py

@@ -30,9 +30,6 @@ def __init__(self, src, basis, batch_size=512, preconditioner="circulant"):
         self.batch_size = batch_size
         self.preconditioner = preconditioner
 
-        self.L = src.L
-        self.n = src.n
-
         if not self.dtype == self.basis.dtype:
             logger.warning(
                 f"Inconsistent types in {self.dtype} Estimator."
@@ -84,11 +81,11 @@ def src_backward(self):
         :return: The adjoint mapping applied to the images, averaged over the whole dataset and expressed
             as coefficients of `basis`.
         """
-        mean_b = np.zeros((self.L, self.L, self.L), dtype=self.dtype)
+        mean_b = np.zeros((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_mean_b = self.src.im_backward(im, i) / self.n
+            batch_mean_b = self.src.im_backward(im, i) / self.src.n
             mean_b += batch_mean_b.astype(self.dtype)
 
         res = self.basis.evaluate_t(mean_b)

src/aspire/reconstruction/mean.py

@@ -14,17 +14,17 @@
 
 class MeanEstimator(Estimator):
     def compute_kernel(self):
-        _2L = 2 * self.L
+        _2L = 2 * self.src.L
         kernel = np.zeros((_2L, _2L, _2L), dtype=self.dtype)
         sq_filters_f = np.square(evaluate_src_filters_on_grid(self.src))
 
-        for i in range(0, self.n, self.batch_size):
-            _range = np.arange(i, min(self.n, i + self.batch_size), dtype=int)
-            pts_rot = rotated_grids(self.L, self.src.rots[_range, :, :])
+        for i in range(0, self.src.n, self.batch_size):
+            _range = np.arange(i, min(self.src.n, i + self.batch_size), dtype=int)
+            pts_rot = rotated_grids(self.src.L, self.src.rots[_range, :, :])
             weights = sq_filters_f[:, :, _range]
             weights *= self.src.amplitudes[_range] ** 2
 
-            if self.L % 2 == 0:
+            if self.src.L % 2 == 0:
                 weights[0, :, :] = 0
                 weights[:, 0, :] = 0
 
@@ -33,7 +33,7 @@ def compute_kernel(self):
 
             kernel += (
                 1
-                / (self.n * self.L**4)
+                / (self.src.n * self.src.L**4)
                 * anufft(weights, pts_rot[::-1], (_2L, _2L, _2L), real=True)
             )