Merge remote-tracking branch 'origin/develop' into ctf_estimator_result

Author: chris-langfield

gallery/tutorials/class_averaging.py

@@ -114,6 +114,7 @@
     large_pca_implementation="legacy",
     nn_implementation="legacy",
     bispectrum_implementation="legacy",
+    num_procs=1,  # Change to "auto" if your machine has many processors
 )
 
 classes, reflections, dists = rir.classify()
@@ -167,6 +168,7 @@
     large_pca_implementation="legacy",
     nn_implementation="sklearn",
     bispectrum_implementation="legacy",
+    num_procs=1,  # Change to "auto" if your machine has many processors
 )
 
 classes, reflections, dists = noisy_rir.classify()

setup.py

@@ -30,10 +30,13 @@ def read(fname):
         "matplotlib>=3.2.0",
         "mrcfile",
         "numpy==1.21.5",
+        "packaging",
         "pandas==1.3.5",
+        "psutil",
         "pyfftw",
         "PyWavelets",
         "pillow",
+        "ray",
         "scipy==1.7.3",
         "scikit-learn",
         "scikit-image",

src/aspire/basis/__init__.py

@@ -7,7 +7,6 @@
 
 # isort: on
 
-from .dirac import DiracBasis
 from .fb_2d import FBBasis2D
 from .fb_3d import FBBasis3D
 from .ffb_2d import FFBBasis2D

src/aspire/basis/basis.py

@@ -74,24 +74,55 @@ def evaluate(self, v):
         Evaluate coefficient vector in basis
 
         :param v: A coefficient vector (or an array of coefficient vectors)
-            to be evaluated. The first dimension must equal `self.count`.
+            to be evaluated. The first dimension must correspond to the number of
+            coefficient vectors, while the second must correspond to `self.count`
         :return: The evaluation of the coefficient vector(s) `v` for this basis.
-            This is an array whose first dimensions equal `self.z` and the
-            remaining dimensions correspond to dimensions two and higher of `v`.
+            This is an Image or a Volume object containing one image/volume for each
+            coefficient vector, and of size `self.sz`.
         """
+        if v.dtype != self.dtype:
+            logger.warning(
+                f"{self.__class__.__name__}::evaluate"
+                f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
+            )
+
+        if self.ndim == 2:
+            return Image(self._evaluate(v))
+        elif self.ndim == 3:
+            return Volume(self._evaluate(v))
+
+    def _evaluate(self, v):
         raise NotImplementedError("subclasses must implement this")
 
     def evaluate_t(self, v):
         """
         Evaluate coefficient in dual basis
 
-        :param v: The coefficient array to be evaluated. The first dimensions
-            must equal `self.sz`.
-        :return: The evaluation of the coefficient array `v` in the dual
+        :param v: An Image or Volume object whose size matches `self.sz`.
+        :return: The evaluation of the Image or Volume object `v` in the dual
             basis of `basis`.
-            This is an array of vectors whose first dimension equals `self.count`
-            and whose remaining dimensions correspond to higher dimensions of `v`.
+            This is an array of vectors whose first dimension equals the number of
+            images/volumes in `v`. and whose second dimension is `self.count`.
         """
+        if v.dtype != self.dtype:
+            logger.warning(
+                f"{self.__class__.__name__}::evaluate_t"
+                f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
+            )
+
+        if not isinstance(v, Image) and not isinstance(v, Volume):
+            if self.ndim == 2:
+                _class = Image
+            elif self.ndim == 3:
+                _class = Volume
+            logger.warning(
+                f"{self.__class__.__name__}::evaluate_t"
+                f" passed numpy array instead of {_class}."
+            )
+            v = _class(v)
+        return self._evaluate_t(v)
+
+    def _evaluate_t(self, v):
         raise NotImplementedError("Subclasses should implement this")
 
     def mat_evaluate(self, V):
@@ -104,7 +135,7 @@ def mat_evaluate(self, V):
             -`self.sz` corresponding to the evaluation of `V` in
             this basis.
         """
-        return mdim_mat_fun_conj(V, 1, len(self.sz), self.evaluate)
+        return mdim_mat_fun_conj(V, 1, len(self.sz), self._evaluate)
 
     def mat_evaluate_t(self, X):
         """
@@ -120,7 +151,7 @@ def mat_evaluate_t(self, X):
             function calculates V = B' * X * B, where the rows of `B`, rows
             of 'X', and columns of `X` are read as vectorized arrays.
         """
-        return mdim_mat_fun_conj(X, len(self.sz), 1, self.evaluate_t)
+        return mdim_mat_fun_conj(X, len(self.sz), 1, self._evaluate_t)
 
     def expand(self, x):
         """

src/aspire/basis/dirac.py

@@ -5,7 +5,6 @@
 
 from aspire.basis import FBBasisMixin, SteerableBasis2D
 from aspire.basis.basis_utils import unique_coords_nd
-from aspire.image import Image
 from aspire.utils import complex_type, real_type, roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_flatten, m_reshape
 
@@ -189,7 +188,7 @@ def basis_norm_2d(self, ell, k):
 
         return rad_norm, ang_norm
 
-    def evaluate(self, v):
+    def _evaluate(self, v):
         """
         Evaluate coefficients in standard 2D coordinate basis from those in FB basis
 
@@ -199,13 +198,6 @@ def evaluate(self, v):
             This is an array whose last dimensions equal `self.sz` and the remaining
             dimensions correspond to first dimensions of `v`.
         """
-
-        if v.dtype != self.dtype:
-            logger.warning(
-                f"{self.__class__.__name__}::evaluate"
-                f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
-            )
-
         # Transpose here once, instead of several times below  #RCOPT
         v = v.reshape(-1, self.count).T
 
@@ -242,7 +234,7 @@ def evaluate(self, v):
 
         return x
 
-    def evaluate_t(self, v):
+    def _evaluate_t(self, v):
         """
         Evaluate coefficient in FB basis from those in standard 2D coordinate basis
 
@@ -253,17 +245,7 @@ def evaluate_t(self, v):
              `self.count` and whose first dimensions correspond to
              first dimensions of `v`.
         """
-
-        if v.dtype != self.dtype:
-            logger.warning(
-                f"{self.__class__.__name__}::evaluate_t"
-                f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
-            )
-
-        if isinstance(v, Image):
-            v = v.asnumpy()
-
-        v = v.T  # RCOPT
+        v = v.asnumpy().T  # RCOPT
 
         x, sz_roll = unroll_dim(v, self.ndim + 1)
         x = m_reshape(

src/aspire/basis/fb_2d.py

@@ -6,6 +6,7 @@
 from aspire.basis.basis_utils import real_sph_harmonic, sph_bessel, unique_coords_nd
 from aspire.utils import roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_flatten, m_reshape
+from aspire.volume import Volume
 
 logger = logging.getLogger(__name__)
 
@@ -140,7 +141,7 @@ def basis_norm_3d(self, ell, k):
             * np.sqrt((self.nres / 2) ** 3)
         )
 
-    def evaluate(self, v):
+    def _evaluate(self, v):
         """
         Evaluate coefficients in standard 3D coordinate basis from those in FB basis
         :param v: A coefficient vector (or an array of coefficient vectors) to
@@ -186,7 +187,7 @@ def evaluate(self, v):
 
         return x.T
 
-    def evaluate_t(self, v):
+    def _evaluate_t(self, v):
         """
         Evaluate coefficient in FB basis from those in standard 3D coordinate basis
 
@@ -197,7 +198,10 @@ def evaluate_t(self, v):
             equals `self.count` and whose remaining dimensions correspond
             to higher dimensions of `v`.
         """
-
+        # v may be a Volume object or a 7D array passed from Basis.mat_evaluate_t
+        # making this check important
+        if isinstance(v, Volume):
+            v = v.asnumpy()
         v = v.T
         x, sz_roll = unroll_dim(v, self.ndim + 1)
         x = m_reshape(

src/aspire/basis/fb_3d.py

@@ -6,7 +6,6 @@
 
 from aspire.basis import FBBasis2D
 from aspire.basis.basis_utils import lgwt
-from aspire.image import Image
 from aspire.nufft import anufft, nufft
 from aspire.numeric import fft, xp
 from aspire.utils import complex_type
@@ -102,7 +101,7 @@ def get_radial(self):
         """
         return self._precomp["radial"]
 
-    def evaluate(self, v):
+    def _evaluate(self, v):
         """
         Evaluate coefficients in standard 2D coordinate basis from those in FB basis
 
@@ -112,13 +111,6 @@ def evaluate(self, v):
             coordinate basis. This is Image instance with resolution of `self.sz`
             and the first dimension correspond to remaining dimension of `v`.
         """
-
-        if v.dtype != self.dtype:
-            logger.debug(
-                f"{self.__class__.__name__}::evaluate"
-                f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
-            )
-
         sz_roll = v.shape[:-1]
         v = v.reshape(-1, self.count)
 
@@ -187,9 +179,9 @@ def evaluate(self, v):
         # Return X as Image instance with the last two dimensions as *self.sz
         x = x.reshape((*sz_roll, *self.sz))
 
-        return Image(x)
+        return x
 
-    def evaluate_t(self, x):
+    def _evaluate_t(self, x):
         """
         Evaluate coefficient in FB basis from those in standard 2D coordinate basis
 
@@ -199,20 +191,6 @@ def evaluate_t(self, x):
             This is an array of vectors whose last dimension equals `self.count`
             and whose first dimension correspond to `x.n_images`.
         """
-
-        if x.dtype != self.dtype:
-            logger.warning(
-                f"{self.__class__.__name__}::evaluate_t"
-                f" Inconsistent dtypes v: {x.dtype} self: {self.dtype}"
-            )
-
-        if not isinstance(x, Image):
-            logger.warning(
-                f"{self.__class__.__name__}::evaluate_t"
-                " passed numpy array instead of Image."
-            )
-            x = Image(x)
-
         # get information on polar grids from precomputed data
         n_theta = np.size(self._precomp["freqs"], 2)
         n_r = np.size(self._precomp["freqs"], 1)

src/aspire/basis/ffb_2d.py

@@ -153,7 +153,7 @@ def _precomp(self):
             "fourier_pts": fourier_pts,
         }
 
-    def evaluate(self, v):
+    def _evaluate(self, v):
         """
         Evaluate coefficients in standard 3D coordinate basis from those in 3D FB basis
 
@@ -278,7 +278,7 @@ def evaluate(self, v):
         x = x.reshape((*sz_roll, *self.sz))
         return x
 
-    def evaluate_t(self, x):
+    def _evaluate_t(self, x):
         """
         Evaluate coefficient in FB basis from those in standard 3D coordinate basis
 
@@ -289,6 +289,8 @@ def evaluate_t(self, x):
             `self.count` and whose remaining dimensions correspond to higher
             dimensions of `x`.
         """
+        x = x.asnumpy()
+
         # roll dimensions
         sz_roll = x.shape[:-3]
         x = x.reshape((-1, *self.sz))