diff --git a/src/aspire/denoising/denoised_src.py b/src/aspire/denoising/denoised_src.py index df6a6d956f..7ec454cade 100644 --- a/src/aspire/denoising/denoised_src.py +++ b/src/aspire/denoising/denoised_src.py @@ -40,11 +40,13 @@ def _images(self, start=0, num=np.inf, indices=None, batch_size=512): start = indices.min() end = indices.max() - im = np.empty((self.L, self.L, len(indices))) + nimgs = len(indices) + im = np.empty((nimgs, self.L, self.L)) - logger.info(f"Loading {len(indices)} images complete") - for istart in range(start, end, batch_size): + logger.info(f"Loading {nimgs} images complete") + for istart in range(start, end + 1, batch_size): imgs_denoised = self.denoiser.images(istart, batch_size) - im = imgs_denoised.data + iend = min(istart + batch_size, end + 1) + im[istart:iend] = imgs_denoised.data return Image(im) diff --git a/src/aspire/denoising/denoiser_cov2d.py b/src/aspire/denoising/denoiser_cov2d.py index 8e026dd336..124b23d281 100644 --- a/src/aspire/denoising/denoiser_cov2d.py +++ b/src/aspire/denoising/denoiser_cov2d.py @@ -130,9 +130,7 @@ def denoise(self, covar_opt=None, batch_size=512): # Initialize the rotationally invariant covariance matrix of 2D images # A fixed batch size is used to go through each image - self.cov2d = BatchedRotCov2D( - self.src, self.basis, batch_size=batch_size, dtype=self.dtype - ) + self.cov2d = BatchedRotCov2D(self.src, self.basis, batch_size=batch_size) default_opt = { "shrinker": "frobenius_norm", diff --git a/tests/test_covar2d_denoiser.py b/tests/test_covar2d_denoiser.py new file mode 100644 index 0000000000..7386ef7a4f --- /dev/null +++ b/tests/test_covar2d_denoiser.py @@ -0,0 +1,43 @@ +from unittest import TestCase + +import numpy as np + +from aspire.basis.ffb_2d import FFBBasis2D +from aspire.denoising.denoiser_cov2d import DenoiserCov2D +from aspire.operators.filters import RadialCTFFilter, ScalarFilter +from aspire.source.simulation import Simulation + + +class BatchedRotCov2DTestCase(TestCase): + def testMSE(self): + # need larger numbers of images and higher resolution for good MSE + dtype = np.float32 + img_size = 64 + num_imgs = 1024 + noise_var = 0.1848 + noise_filter = ScalarFilter(dim=2, value=noise_var) + filters = [ + RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1) + for d in np.linspace(1.5e4, 2.5e4, 7) + ] + # set simulation object + sim = Simulation( + L=img_size, + n=num_imgs, + unique_filters=filters, + offsets=0.0, + amplitudes=1.0, + dtype=dtype, + noise_filter=noise_filter, + ) + imgs_clean = sim.projections() + + # Specify the fast FB basis method for expending the 2D images + ffbbasis = FFBBasis2D((img_size, img_size), dtype=dtype) + denoiser = DenoiserCov2D(sim, ffbbasis, noise_var) + denoised_src = denoiser.denoise(batch_size=64) + imgs_denoised = denoised_src.images(0, num_imgs) + # Calculate the normalized RMSE of the estimated images. + nrmse_ims = (imgs_denoised - imgs_clean).norm() / imgs_clean.norm() + + self.assertTrue(nrmse_ims < 0.25)