Merge branch 'develop' into denoiser_test

Author: junchaoxia

.travis.yml

@@ -1,21 +1,38 @@
-sudo: required
 language: python
-matrix:
+
+# manylinux2010 wheels are not found on xenial
+dist: bionic
+
+os:
+  - linux
+
+arch:
+  - amd64
+
+stages:
+  - check
+  - test
+  - deploy
+
+jobs:
   include:
     - python: 3.6
-      env:
-        - TOXENV=clean,py36-stable,coveralls,report
-        - TOXENV=py36-dev
+      env: TOXENV=clean,py36-stable,coveralls,report
+    - python: 3.6
+      env: TOXENV=py36-dev
     - python: 3.7
-      env:
-        - TOXENV=py37-stable
-        - TOXENV=py37-dev
+      env: TOXENV=py37-stable
+    - python: 3.7
+      env: TOXENV=py37-dev
+    - python: 3.8
+      env: TOXENV=py38-stable
     - python: 3.8
-      env:
-        - TOXENV=py38-stable
-        - TOXENV=py38-dev
-    - env: TOXENV=check
-    - env: TOXENV=docs
+      env: TOXENV=py38-dev
+    - python: 3.6
+      env: TOXENV=check
+    - python: 3.6
+      env: TOXENV=docs
+
 install:
   - pip install -U tox tox-travis
 

setup.py

@@ -24,7 +24,7 @@ def read(fname):
     author_email="[email protected]",
     install_requires=[
         "click",
-        "finufftpy",
+        "finufft",
         "importlib_resources>=1.0.2",
         "joblib",
         "jupyter",

src/aspire/basis/ffb_2d.py

@@ -91,7 +91,7 @@ def _precomp(self):
             np.sin(np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
             (1, n_theta),
         )
-        freqs = np.vstack((freqs_x[np.newaxis, ...], freqs_y[np.newaxis, ...]))
+        freqs = np.vstack((freqs_y[np.newaxis, ...], freqs_x[np.newaxis, ...]))
 
         return {"gl_nodes": r, "gl_weights": w, "radial": radial, "freqs": freqs}
 

src/aspire/basis/ffb_3d.py

@@ -138,9 +138,9 @@ def _precomp(self):
             * pi
             * np.vstack(
                 (
-                    fourier_x[np.newaxis, ...],
-                    fourier_y[np.newaxis, ...],
                     fourier_z[np.newaxis, ...],
+                    fourier_y[np.newaxis, ...],
+                    fourier_x[np.newaxis, ...],
                 )
             )
         )
@@ -269,7 +269,7 @@ def evaluate(self, v):
         pf = m_reshape(pf, (n_theta * n_phi * n_r, n_data))
 
         # perform inverse non-uniformly FFT transformation back to 3D rectangular coordinates
-        freqs = m_reshape(self._precomp["fourier_pts"], (3, n_r * n_theta * n_phi, -1))
+        freqs = m_reshape(self._precomp["fourier_pts"], (3, n_r * n_theta * n_phi))
         x = np.zeros((n_data, self.sz[0], self.sz[1], self.sz[2]), dtype=v.dtype)
         for isample in range(0, n_data):
             x[isample] = np.real(anufft(pf[:, isample], freqs, self.sz))

src/aspire/basis/fpswf_2d.py

@@ -87,7 +87,7 @@ def _precomp(self):
         self.num_angular_pts = f
 
         # pre computing variables for forward
-        us_fft_pts = np.column_stack((self.quad_rule_pts_x, self.quad_rule_pts_y))
+        us_fft_pts = np.column_stack((self.quad_rule_pts_y, self.quad_rule_pts_x))
         us_fft_pts = self.bandlimit / (self.rcut * np.pi * 2) * us_fft_pts  # for pynfft
         (
             blk_r,

src/aspire/basis/polar_2d.py

@@ -64,10 +64,10 @@ def _precomp(self):
         for i in range(self.ntheta // 2):
             freqs[0, i * self.nrad : (i + 1) * self.nrad] = np.arange(
                 self.nrad
-            ) * np.sin(i * dtheta)
+            ) * np.cos(i * dtheta)
             freqs[1, i * self.nrad : (i + 1) * self.nrad] = np.arange(
                 self.nrad
-            ) * np.cos(i * dtheta)
+            ) * np.sin(i * dtheta)
 
         freqs *= omega0
         return freqs

src/aspire/nufft/__init__.py

@@ -62,7 +62,7 @@ def _try_backend(backend):
 
         elif backend == "finufft":
             try:
-                from finufftpy import nufft1d1  # noqa: F401
+                from finufft import Plan  # noqa: F401
 
                 from aspire.nufft.finufft import FinufftPlan
 

src/aspire/nufft/cufinufft.py

@@ -82,12 +82,8 @@ def __init__(self, sz, fourier_pts, epsilon=1e-15, ntransforms=1, **kwargs):
         #   is tied to instance, instead of this method.
         self.fourier_pts_gpu = gpuarray.to_gpu(self.fourier_pts)
 
-        # Note, cufinufft expects column major signal data,
-        #   but we don't want to transpose large arrays.
-        #   We instead reverse the references to the points axes.
-        #   This is identical to what is done for FINUFFT.
-        self._transform_plan.set_pts(self.num_pts, *self.fourier_pts_gpu[::-1])
-        self._adjoint_plan.set_pts(self.num_pts, *self.fourier_pts_gpu[::-1])
+        self._transform_plan.set_pts(self.num_pts, *self.fourier_pts_gpu)
+        self._adjoint_plan.set_pts(self.num_pts, *self.fourier_pts_gpu)
 
     def transform(self, signal):
         """

src/aspire/nufft/finufft.py

@@ -1,6 +1,6 @@
 import logging
 
-import finufftpy
+import finufft
 import numpy as np
 
 from aspire.nufft import Plan
@@ -10,84 +10,73 @@
 
 
 class FinufftPlan(Plan):
-    def __init__(self, sz, fourier_pts, epsilon=1e-15, ntransforms=1, **kwargs):
+    def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
         """
         A plan for non-uniform FFT in 2D or 3D.
 
-        :param sz: A tuple indicating the geometry of the signal
-        :param fourier_pts: The points in Fourier space where the Fourier transform is to be calculated,
-            arranged as a dimension-by-K array. These need to be in the range [-pi, pi] in each dimension.
-        :param epsilon: The desired precision of the NUFFT
-        :param ntransforms: Optional integer indicating if you would like to compute a batch of `ntransforms`
-        transforms.  Implies vol_f.shape is (..., `ntransforms`). Defaults to 0 which disables batching.
+        :param sz: A tuple indicating the geometry of the signal.
+        :param fourier_pts: The points in Fourier space where the Fourier
+        transform is to be calculated, arranged as a dimension-by-K array.
+        These need to be in the range [-pi, pi] in each dimension.
+        :param epsilon: The desired precision of the NUFFT.
+        :param ntransforms: Optional integer indicating if you would like
+        to compute a batch of `ntransforms`.
+        transforms.  Implies vol_f.shape is (`ntransforms`, ...).
         """
 
         self.ntransforms = ntransforms
-        manystr = ""
-        if self.ntransforms > 1:
-            manystr = "many"
 
         self.sz = sz
         self.dim = len(sz)
 
         self.dtype = fourier_pts.dtype
 
-        # TODO: Currently/historically finufftpy is hardcoded as doubles only (inside the binding layer).
-        #       This has been changed in their GuruV2 work,
-        #         and an updated package should be released and integrated very soon.
-        #       The following casting business is only to facilitate the transition.
-        #       I don't want to send one precision in and get a different one out.
-        #         We have enough code that does that already.
-        #         (Potentially anything that used this, for example).
-        #       I would error, but I know a bunch of code ASPIRE wants to work,
-        #         the cov2d tutorial for example, would fail and require hacks to run
-        #         if I was strict about dtypes rn.
-        #       I would rather deal with that in other, more targeted PRs.
-        #       This preserves the legacy behavior of admitting singles,
-        #         but I will correct it slightly, to return the precision given as input.
-        #       This approach should contain the hacks here to a single place on the edge,
-        #         instead of spread through the code. ASPIRE code should focus on being
-        #         internally consistent.
-        #       Admittedly not ideal, but ignoring these problems wasn't sustainable.
-
-        self.cast_output = False
-        if self.dtype != np.float64:
-            logger.debug(
-                "This version of finufftpy is hardcoded to doubles internally"
-                "  casting input to doubles, results cast back to singles."
-            )
-            self.cast_output = True
-            self.dtype = np.float64
-
         self.complex_dtype = complex_type(self.dtype)
 
-        # TODO: Things get messed up unless we ensure a 'C' ordering here - investigate why
-        self.fourier_pts = np.asarray(
-            np.mod(fourier_pts + np.pi, 2 * np.pi) - np.pi, order="C", dtype=self.dtype
+        self.fourier_pts = np.ascontiguousarray(
+            np.mod(fourier_pts + np.pi, 2 * np.pi) - np.pi
         )
+
         self.num_pts = fourier_pts.shape[1]
-        self.epsilon = epsilon
 
-        # Get a handle on the appropriate 1d/2d/3d forward transform function in finufftpy
-        self.transform_function = getattr(finufftpy, f"nufft{self.dim}d2{manystr}")
+        self.epsilon = max(epsilon, np.finfo(self.dtype).eps)
+        if self.epsilon != epsilon:
+            logger.debug(
+                f"FinufftPlan adjusted eps={self.epsilon}" f" from requested {epsilon}."
+            )
+
+        self._transform_plan = finufft.Plan(
+            nufft_type=2,
+            n_modes_or_dim=self.sz,
+            eps=self.epsilon,
+            n_trans=self.ntransforms,
+            dtype=self.dtype,
+        )
+
+        self._adjoint_plan = finufft.Plan(
+            nufft_type=1,
+            n_modes_or_dim=self.sz,
+            eps=self.epsilon,
+            n_trans=self.ntransforms,
+            dtype=self.dtype,
+        )
 
-        # Get a handle on the appropriate 1d/2d/3d adjoint function in finufftpy
-        self.adjoint_function = getattr(finufftpy, f"nufft{self.dim}d1{manystr}")
+        self._transform_plan.setpts(*self.fourier_pts)
+        self._adjoint_plan.setpts(*self.fourier_pts)
 
     def transform(self, signal):
         """
         Compute the NUFFT transform using this plan instance.
 
         :param signal: Signal to be transformed. For a single transform,
         this should be a a 1, 2, or 3D array matching the plan `sz`.
-        For a batch, signal should have shape `(*sz, ntransforms)`.
+        For a batch, signal should have shape `(ntransforms, *sz)`.
 
         :returns: Transformed signal of shape `num_pts` or
         `(ntransforms, num_pts)`.
         """
 
-        sig_shape = signal.shape
-        res_shape = self.num_pts
+        sig_frame_shape = signal.shape
         # Note, there is a corner case for ntransforms == 1.
         if self.ntransforms > 1 or (
             self.ntransforms == 1 and len(signal.shape) == self.dim + 1
@@ -103,37 +92,18 @@ def transform(self, signal):
                 f" should match ntransforms {self.ntransforms}.",
             )
 
-            sig_shape = signal.shape[1:]
-            res_shape = (self.ntransforms, self.num_pts)
+            sig_frame_shape = signal.shape[1:]
+
+            # finufft expects signal.ndim == dim for ntransforms = 1.
+            if self.ntransforms == 1:
+                signal = signal.reshape(self.sz)
 
         ensure(
-            sig_shape == self.sz,
+            sig_frame_shape == self.sz,
             f"Signal frame to be transformed must have shape {self.sz}",
         )
 
-        epsilon = max(self.epsilon, np.finfo(signal.dtype).eps)
-
-        # Forward transform functions in finufftpy have signatures of the form:
-        # (x, y, z, c, isign, eps, f, ...)
-        # (x, y     c, isign, eps, f, ...)
-        # (x,       c, isign, eps, f, ...)
-        # Where f is a Fortran-order ndarray of the appropriate dimensions
-        # We form these function signatures here by tuple-unpacking
-
-        result = np.zeros(res_shape, dtype=self.complex_dtype)
-
-        result_code = self.transform_function(
-            *self.fourier_pts,
-            result,
-            -1,
-            epsilon,
-            signal.T,  # RCOPT, currently F ordered, should change in gv2
-        )
-
-        if result_code != 0:
-            raise RuntimeError(f"FINufft transform failed. Result code {result_code}")
-        if self.cast_output:
-            result = result.astype(np.complex64)
+        result = self._transform_plan.execute(signal)
 
         return result
 
@@ -145,19 +115,9 @@ def adjoint(self, signal):
         this should be a a 1D array of len `num_pts`.
         For a batch, signal should have shape `(ntransforms, num_pts)`.
 
-        :returns: Transformed signal `(sz)` or `(sz, ntransforms)`.
+        :returns: Transformed signal `(sz)` or `(ntransforms, sz)`.
         """
 
-        epsilon = max(self.epsilon, np.finfo(signal.dtype).eps)
-
-        # Adjoint functions in finufftpy have signatures of the form:
-        # (x, y, z, c, isign, eps, ms, mt, mu, f, ...)
-        # (x, y     c, isign, eps, ms, mt      f, ...)
-        # (x,       c, isign, eps, ms,         f, ...)
-        # Where f is a Fortran-order ndarray of the appropriate dimensions
-        # We form these function signatures here by tuple-unpacking
-
-        res_shape = self.sz
         # Note, there is a corner case for ntransforms == 1.
         if self.ntransforms > 1 or (self.ntransforms == 1 and len(signal.shape) == 2):
             ensure(
@@ -170,31 +130,11 @@ def adjoint(self, signal):
                 "For multiple transforms, signal stack length"
                 f" should match ntransforms {self.ntransforms}.",
             )
-            res_shape = (
-                self.ntransforms,
-                *self.sz,
-            )
 
-        result = np.zeros(res_shape, dtype=self.complex_dtype)
-
-        # FINUFFT is F order at this time. The bindings
-        # will pickup the fact `signal` is C_Contiguous,
-        # and transpose the data; we just need to transpose
-        # the indices.  I think the next release addresses this.
-        #   Note in the 2020 hackathon this was changed directly in FFB,
-        #   which worked because GPU arrays just need the pointer anyway...
-        #   This is a quirk of this version of FINUFFT, and
-        #   so probably belongs here at the edge,
-        #   away from other implementations.
-        signal = signal.reshape(signal.shape[::-1])
-
-        result_code = self.adjoint_function(
-            *self.fourier_pts, signal, 1, epsilon, *self.sz, result
-        )
-        if result_code != 0:
-            raise RuntimeError(f"FINufft adjoint failed. Result code {result_code}")
+            # finufft is expecting flat array for 1D case.
+            if self.ntransforms == 1:
+                signal = signal.reshape(self.num_pts)
 
-        if self.cast_output:
-            result = result.astype(np.complex64)
+        result = self._adjoint_plan.execute(signal)
 
         return result

src/aspire/utils/types.py

@@ -57,30 +57,6 @@ def complex_type(realtype):
     return complextype
 
 
-def dtype_legacy_string(dtype):
-    """
-    Get legacy dtype string from corresponding numpy dtype.
-
-    :param dtype: Numpy dtype
-    :return: string
-    """
-
-    dtype_to_string_map = {
-        "float32": "single",
-        "float64": "double",
-        "complex128": "complex",
-    }
-
-    dtype_str = dtype_to_string_map.get(str(dtype))
-
-    if not dtype_str:
-        msg = f"Corresponding dtype {str(dtype)} is not defined."
-        logger.error(msg)
-        raise TypeError(msg)
-
-    return dtype_str
-
-
 def utest_tolerance(dtype):
     """
     Return ASPIRE tolerance for unit tests based on `dtype`.