Merge pull request #151 from ComputationalCryoEM/orientation

Include the Dirac and Polar Fourier basis Classes

Author: junchaoxia

src/aspire/basis/__init__.py

@@ -152,46 +152,46 @@ def mat_evaluate_t(self, X):
         """
         return mdim_mat_fun_conj(X, len(self.sz), 1, self.evaluate_t)
 
-    def expand(self, v):
+    def expand(self, x):
         """
-        Expand array in basis
+        Obtain coefficients in the basis from those in standard coordinate basis
 
-        This is a similar function to `evaluate_t` but with more accuracy by
-         using the cg optimizing of linear equation, Ax=b.
+        This is a similar function to evaluate_t but with more accuracy by using
+        the cg optimizing of linear equation, Ax=b.
 
-        If `v` is a matrix of size `basis.ct`-by-..., `B` is the change-of-basis
-        matrix of this basis, and `x` is a matrix of size `self.sz`-by-...,
-        the function calculates  v = (B' * B)^(-1) * B' * x, where the rows
-        of `B` and columns of `x` are read as vectorized arrays.
-
-        :param v: An array whose first few dimensions are to be expanded in this basis.
-            These dimensions must equal `self.sz`.
-        :return: The coefficients of `v` expanded in this basis. If more than
-            one array of size `self.sz` is found in `v`, the second and higher
-            dimensions of the return value correspond to those higher dimensions of `v`.
+        :param x: An array whose first two or three dimensions are to be expanded
+            the desired basis. These dimensions must equal `self.sz`.
+        :return : The coefficients of `v` expanded in the desired basis.
+            The first dimension of `v` is with size of `count` and the
+            second and higher dimensions of the return value correspond to
+            those higher dimensions of `x`.
 
         """
-        ensure(v.shape[:self.ndim] == self.sz, f'First {self.ndim} dimensions of v must match {self.sz}.')
+        # ensure the first dimensions with size of self.sz
+        x, sz_roll = unroll_dim(x, self.ndim + 1)
+        ensure(x.shape[:self.ndim] == self.sz,
+               f'First {self.ndim} dimensions of x must match {self.sz}.')
 
-        v, sz_roll = unroll_dim(v, self.ndim + 1)
-        b = self.evaluate_t(v)
-        operator = LinearOperator(
-            shape=(self.count, self.count),
-            matvec=lambda x: self.evaluate_t(self.evaluate(x))
-        )
+        operator = LinearOperator(shape=(self.count, self.count),
+                                  matvec=lambda v: self.evaluate_t(self.evaluate(v)))
 
         # TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
-        tol = 10 * np.finfo(v.dtype).eps
+        tol = 10*np.finfo(x.dtype).eps
         logger.info('Expanding array in basis')
-        v, info = cg(operator, b, tol=tol)
 
-        v = v[..., np.newaxis]
+        # number of image samples
+        n_data = np.size(x, self.ndim)
+        v = np.zeros((self.count, n_data), dtype=x.dtype)
 
-        if info != 0:
-            raise RuntimeError('Unable to converge!')
+        for isample in range(0, n_data):
+            b = self.evaluate_t(x[..., isample])
+            # TODO: need check the initial condition x0 can improve the results or not.
+            v[..., isample], info = cg(operator, b, tol=tol)
+            if info != 0:
+                raise RuntimeError('Unable to converge!')
 
+        # return v coefficients with the first dimension of self.count
         v = roll_dim(v, sz_roll)
-
         return v
 
     def expand_t(self, v):

src/aspire/basis/dirac.py

@@ -1,5 +1,74 @@
+import logging
+import numpy as np
+
 from aspire.basis import Basis
+from aspire.utils.matrix import roll_dim, unroll_dim
+from aspire.utils.matlab_compat import m_flatten, m_reshape
+
+logger = logging.getLogger(__name__)
 
 
 class DiracBasis(Basis):
-    pass
+    """
+    Define a derived class for Dirac basis
+    """
+    def __init__(self, sz, mask=None):
+        """
+        Initialize an object for Dirac basis
+        :param sz: The shape of the vectors for which to define the basis.
+        :param mask: A boolean _mask of size sz indicating which coordinates
+            to include in the basis (default np.full(sz, True)).
+        """
+        if mask is None:
+            mask = np.full(sz, True)
+        self._mask = m_flatten(mask)
+
+        super().__init__(sz)
+
+    def _build(self):
+        """
+        Build the internal data structure to Dirac basis
+        """
+        logger.info('Expanding object in a Dirac basis.')
+        self.count = np.sum(self._mask)
+        self._sz_prod = self.nres ** self.ndim
+
+    def evaluate(self, v):
+        """
+        Evaluate coefficients in standard coordinate basis from those in Dirac basis
+
+        :param v: A coefficient vector (or an array of coefficient vectors) to
+            be evaluated. The first dimension must equal `self.count`.
+        :return: The evaluation of the coefficient vector(s) `v` for this basis.
+            This is an array whose first dimensions equal `self.sz` and the remaining
+            dimensions correspond to dimensions two and higher of `v`.
+        """
+        v, sz_roll = unroll_dim(v, 2)
+        x = np.zeros(shape=(self._sz_prod,) + v.shape[1:])
+        x[self._mask, ...] = v
+        x = m_reshape(x, self.sz + x.shape[1:])
+        x = roll_dim(x, sz_roll)
+
+        return x
+
+    def evaluate_t(self, x):
+        """
+        Evaluate coefficient in Dirac basis from those in standard coordinate basis
+
+        :param x: The coefficient array to be evaluated. The first dimensions
+            must equal `self.sz`.
+        :return: The evaluation of the coefficient array `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`.
+        """
+        x, sz_roll = unroll_dim(x, self.ndim + 1)
+        x = m_reshape(x, new_shape=(self._sz_prod,) + x.shape[self.ndim:])
+        v = np.zeros(shape=(self.count, ) + x.shape[1:])
+        v = x[self._mask, ...]
+        v = roll_dim(v, sz_roll)
+
+        return v
+
+    def expand(self, x):
+        return self.evaluate_t(x)

src/aspire/basis/ffb_2d.py

@@ -3,11 +3,8 @@
 from numpy import pi
 from scipy.special import jv
 from scipy.fftpack import ifft, fft
-from scipy.sparse.linalg import LinearOperator, cg
 
 from aspire.nfft import anufft3, nufft3
-
-from aspire.utils import ensure
 from aspire.utils.matrix import roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_reshape
 from aspire.basis.basis_utils import lgwt
@@ -246,45 +243,3 @@ def evaluate_t(self, x):
         # return v coefficients with the first dimension of self.count
         v = roll_dim(v, sz_roll)
         return v
-
-    def expand(self, x):
-        """
-        Obtain coefficients in FB basis from those in standard 2D coordinate basis
-
-        This is a similar function to evaluate_t but with more accuracy by using
-        the cg optimizing of linear equation, Ax=b.
-
-        :param x: An array whose first two dimensions are to be expanded in FB basis.
-             These dimensions must equal `self.sz`.
-        :return : The coefficients of `v` expanded in FB basis. The first dimension
-            of `v` is with size of `count` and the second and higher dimensions
-            of the return value correspond to those higher dimensions of `x`.
-
-        """
-        # ensure the first two dimensions with size of self.sz
-        x, sz_roll = unroll_dim(x, self.ndim + 1)
-        x = m_reshape(x, (self.sz[0], self.sz[1], -1))
-        ensure(x.shape[:self.ndim] == self.sz,
-               f'First {self.ndim} dimensions of x must match {self.sz}.')
-
-        operator = LinearOperator(shape=(self.count, self.count),
-                                  matvec=lambda v: self.evaluate_t(self.evaluate(v)))
-
-        # TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
-        tol = 10*np.finfo(x.dtype).eps
-        logger.info('Expanding array in basis')
-
-        # number of image samples
-        n_data = np.size(x, self.ndim)
-        v = np.zeros((self.count, n_data), dtype=x.dtype)
-
-        for isample in range(0, n_data):
-            b = self.evaluate_t(x[..., isample])
-            # TODO: need check the initial condition x0 can improve the results or not.
-            v[..., isample], info = cg(operator, b, tol=tol)
-            if info != 0:
-                raise RuntimeError('Unable to converge!')
-
-        # return v coefficients with the first dimension of self.count
-        v = roll_dim(v, sz_roll)
-        return v

src/aspire/basis/ffb_3d.py

@@ -1,10 +1,8 @@
 import logging
 import numpy as np
 from numpy import pi
-from scipy.sparse.linalg import LinearOperator, cg
 
 from aspire.nfft import anufft3, nufft3
-from aspire.utils import ensure
 from aspire.utils.matrix import roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_flatten, m_reshape
 from aspire.basis.basis_utils import sph_bessel, norm_assoc_legendre, lgwt
@@ -339,48 +337,3 @@ def evaluate_t(self, x):
             v[ind, :] = v_ell
         v = roll_dim(v, sz_roll)
         return v
-
-    def expand(self, x):
-
-        """
-        Obtain expansion coefficients in FB basis from those in standard 3D coordinate basis
-
-        This is a similar function to evaluate_t but with more accuracy by using
-        the cg optimizing of linear equation, Ax=b.
-
-        :param x: An array whose first three dimensions are to be expanded in FB basis.
-             These dimensions must equal `self.sz`.
-        :return : The coefficients of `v` expanded in FB basis. The first dimension
-            of `v` is with size of `count` and the second and higher dimensions
-            of the return value correspond to those higher dimensions of `x`.
-
-        """
-        # TODO: this is function could be move to base class if all standard and fast versions of 2d and 3d are using
-        #       the same data structures of x and v.
-        # ensure the first three dimensions with size of self.sz
-        x, sz_roll = unroll_dim(x, self.ndim + 1)
-        x = m_reshape(x, (self.sz[0], self.sz[1], self.sz[2], -1))
-
-        ensure(x.shape[:self.ndim] == self.sz, f'First {self.ndim} dimensions of x must match {self.sz}.')
-
-        operator = LinearOperator(shape=(self.count, self.count),
-                                  matvec=lambda v: self.evaluate_t(self.evaluate(v)))
-
-        # TODO: (from MATLAB implementation) - Check that this tolerance make sense for multiple columns in v
-        tol = 10*np.finfo(x.dtype).eps
-        logger.info('Expanding array in basis')
-
-        # number of image samples
-        n_data = np.size(x, self.ndim)
-        v = np.zeros((self.count, n_data), dtype=x.dtype)
-
-        for isample in range(0, n_data):
-            b = self.evaluate_t(x[..., isample])
-            # TODO: need check the initial condition x0 can improve the results or not.
-            v[..., isample], info = cg(operator, b, tol=tol)
-            if info != 0:
-                raise RuntimeError('Unable to converge!')
-
-        # return v coefficients with the first dimension of self.count
-        v = roll_dim(v, sz_roll)
-        return v

src/aspire/basis/polar_2d.py

@@ -0,0 +1,130 @@
+import logging
+import numpy as np
+import finufftpy
+
+from aspire.utils import ensure
+from aspire.utils.matrix import roll_dim, unroll_dim
+from aspire.utils.matlab_compat import m_reshape
+from aspire.basis import Basis
+
+logger = logging.getLogger(__name__)
+
+
+class PolarBasis2D(Basis):
+    """
+    Define a derived class for polar Fourier representation for 2D images
+    """
+
+    def __init__(self, size, nrad=None, ntheta=None):
+        """
+        Initialize an object for the 2D polar Fourier grid class
+
+        :param size: The shape of the vectors for which to define the grid.
+            Currently only square images are supported.
+        :param nrad: The number of points in the radial dimension.
+        :param ntheta: The number of points in the angular dimension.
+        """
+
+        ndim = len(size)
+        ensure(ndim == 2, 'Only two-dimensional grids are supported.')
+        ensure(len(set(size)) == 1, 'Only square domains are supported.')
+
+        self.nrad = nrad
+        if nrad is None:
+            self.nrad = self.nres // 2
+
+        self.ntheta = ntheta
+        if ntheta is None:
+            # try to use the same number as Fast FB basis
+            self.ntheta = 8 * self.nrad
+
+        super().__init__(size)
+
+    def _build(self):
+        """
+        Build the internal data structure to 2D polar Fourier grid
+        """
+        logger.info('Represent 2D image in a polar Fourier grid')
+
+        self.count = self.nrad * self.ntheta
+        self._sz_prod = self.sz[0] * self.sz[1]
+
+        # precompute the basis functions in 2D grids
+        self.freqs = self._precomp()
+
+    def _precomp(self):
+        """
+        Precomute the polar Fourier grid
+        """
+        omega0 = 2 * np.pi / (2 * self.nrad - 1)
+        dtheta = 2 * np.pi / self.ntheta
+
+        freqs = np.zeros((2, self.nrad * self.ntheta))
+        for i in range(self.ntheta):
+            freqs[0, i * self.nrad: (i + 1) * self.nrad] = np.arange(self.nrad) * np.sin(i * dtheta)
+            freqs[1, i * self.nrad: (i + 1) * self.nrad] = np.arange(self.nrad) * np.cos(i * dtheta)
+
+        freqs *= omega0
+        return freqs
+
+    def evaluate(self, v):
+        """
+        Evaluate coefficients in standard 2D coordinate basis from those in polar Fourier basis
+
+        :param v: A coefficient vector (or an array of coefficient vectors)
+            in polar Fourier basis to be evaluated. The first dimension must equal to
+            `self.count`.
+        :return x: The evaluation of the coefficient vector(s) `x` in standard 2D
+            coordinate basis. This is an array whose first two dimensions equal `self.sz`
+            and the remaining dimensions correspond to dimensions two and higher of `v`.
+        """
+        v, sz_roll = unroll_dim(v, 2)
+        nimgs = v.shape[1]
+
+        half_size = self.nrad * self.ntheta // 2
+
+        v = m_reshape(v, (self.nrad, self.ntheta, nimgs))
+
+        v = (v[:, :self.ntheta // 2, :]
+             + v[:, self.ntheta // 2:, :].conj())
+
+        v = m_reshape(v, (half_size, nimgs))
+
+        # finufftpy require it to be aligned in fortran order
+        x = np.empty((self._sz_prod, nimgs), dtype='complex128', order='F')
+        finufftpy.nufft2d1many(self.freqs[0, :half_size],
+                               self.freqs[1, :half_size],
+                               v, 1, 1e-15, self.sz[0], self.sz[1], x)
+        x = m_reshape(x, (self.sz[0], self.sz[1], nimgs))
+        x = x.real
+        # return coefficients whose first two dimensions equal to self.sz
+        x = roll_dim(x, sz_roll)
+
+        return x
+
+    def evaluate_t(self, x):
+        """
+        Evaluate coefficient in polar Fourier grid from those in standard 2D coordinate basis
+
+        :param x: The coefficient array in the standard 2D coordinate basis to be
+            evaluated. The first two dimensions must equal `self.sz`.
+        :return v: The evaluation of the coefficient array `v` in the polar Fourier grid.
+            This is an array of vectors whose first dimension is `self.count` and
+            whose remaining dimensions correspond to higher dimensions of `x`.
+        """
+        # ensure the first two dimensions with size of self.sz
+        x, sz_roll = unroll_dim(x, self.ndim + 1)
+        nimgs = x.shape[2]
+
+        # finufftpy require it to be aligned in fortran order
+        half_size = self.nrad * self.ntheta // 2
+        pf = np.empty((half_size, nimgs), dtype='complex128', order='F')
+        finufftpy.nufft2d2many(self.freqs[0, :half_size], self.freqs[1, :half_size], pf, 1, 1e-15, x)
+        pf = m_reshape(pf, (self.nrad, self.ntheta // 2, nimgs))
+        v = np.concatenate((pf, pf.conj()), axis=1)
+
+        # return v coefficients with the first dimension size of self.count
+        v = m_reshape(v, (self.nrad * self.ntheta, nimgs))
+        v = roll_dim(v, sz_roll)
+
+        return v