Merge branch 'merge' into 'NIFTy_7'

Add finufft

See merge request !600

Dockerfile

@@ -13,6 +13,7 @@ RUN apt-get update && apt-get install -y \
     python3-mpi4py python3-matplotlib \
   # more optional NIFTy dependencies
   && pip3 install ducc0 \
+  && pip3 install finufft \
   && pip3 install jupyter \
   && rm -rf /var/lib/apt/lists/*
 

src/__init__.py

@@ -86,7 +86,7 @@ from .library.light_cone_operator import LightConeOperator
 from .library.wiener_filter_curvature import WienerFilterCurvature
 from .library.adjust_variances import (make_adjust_variances_hamiltonian,
                                        do_adjust_variances)
-from .library.gridder import Gridder
+from .library.nft import Gridder, FinuFFT
 from .library.correlated_fields import CorrelatedFieldMaker
 from .library.correlated_fields_simple import SimpleCorrelatedField
 

src/library/gridder.py → src/library/nft.py

@@ -11,23 +11,40 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright(C) 2019 Max-Planck-Society
+# Copyright(C) 2019-2021 Max-Planck-Society
 #
 # NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
 
 import numpy as np
+from scipy.constants import speed_of_light
+from ducc0.wgridder import ms2dirty, dirty2ms
 
 from ..domain_tuple import DomainTuple
 from ..domains.rg_space import RGSpace
 from ..domains.unstructured_domain import UnstructuredDomain
 from ..operators.linear_operator import LinearOperator
 from ..sugar import makeDomain, makeField
+from ..fft import nthreads
 
 
 class Gridder(LinearOperator):
-    def __init__(self, target, uv, eps=2e-10, nthreads=1):
+    """Compute non-uniform 2D FFT with ducc.
+
+    Parameters
+    ----------
+    target : Domain, tuple of domains or DomainTuple.
+        Target domain, must be a single two-dimensional RGSpace.
+    uv : np.ndarray
+        Coordinates of the data-points, shape (n, 2).
+    eps : float
+        Requested precision, defaults to 2e-10.
+    """
+    def __init__(self, target, uv, eps=2e-10):
         self._capability = self.TIMES | self.ADJOINT_TIMES
         self._target = makeDomain(target)
+        for ii in [0, 1]:
+            if target.shape[ii] % 2 != 0:
+                raise ValueError("even number of pixels is required for gridding operation")
         if (len(self._target) != 1 or not isinstance(self._target[0], RGSpace)
                 or not len(self._target.shape) == 2):
             raise ValueError("need target with exactly one 2D RGSpace")
@@ -35,30 +52,71 @@ class Gridder(LinearOperator):
             raise ValueError("uv must be a 2D array")
         if uv.shape[1] != 2:
             raise ValueError("second dimension of uv must have length 2")
-        self._domain = DomainTuple.make(
-            UnstructuredDomain((uv.shape[0])))
+        self._domain = DomainTuple.make(UnstructuredDomain((uv.shape[0])))
         # wasteful hack to adjust to shape required by ducc0.wgridder
         self._uvw = np.empty((uv.shape[0], 3), dtype=np.float64)
         self._uvw[:, 0:2] = uv
         self._uvw[:, 2] = 0.
         self._eps = float(eps)
-        self._nthreads = int(nthreads)
-
 
     def apply(self, x, mode):
         self._check_input(x, mode)
-        speedOfLight = 299792458.
-        freq = np.array([speedOfLight])
+        freq = np.array([speed_of_light])
         x = x.val
         nxdirty, nydirty = self._target[0].shape
         dstx, dsty = self._target[0].distances
-        from ducc0.wgridder import ms2dirty, dirty2ms
         if mode == self.TIMES:
             res = ms2dirty(self._uvw, freq, x.reshape((-1,1)), None, nxdirty,
                            nydirty, dstx, dsty, 0, 0,
-                           self._eps, False, self._nthreads, 0)
+                           self._eps, False, nthreads(), 0)
         else:
             res = dirty2ms(self._uvw, freq, x, None, dstx, dsty, 0, 0,
-                           self._eps, False, self._nthreads, 0)
+                           self._eps, False, nthreads(), 0)
             res = res.reshape((-1,))
         return makeField(self._tgt(mode), res)
+
+
+class FinuFFT(LinearOperator):
+    """Compute non-uniform 1D, 2D and 3D FFTs with finufft.
+
+    Parameters
+    ----------
+    target : Domain, tuple of domains or DomainTuple.
+        Target domain, must be an RGSpace with one to three dimensions.
+    pos : np.ndarray
+        Coordinates of the data-points, shape (n, ndim).
+    eps: float
+        Requested precision, defaults to 2e-10.
+    """
+    def __init__(self, target, pos, eps=2e-10):
+        import finufft
+        self._capability = self.TIMES | self.ADJOINT_TIMES
+        self._target = makeDomain(target)
+        if not isinstance(self._target[0], RGSpace):
+            raise TypeError("target needs to be an RGSpace")
+        if len(self._target.shape) > 3:
+            raise ValueError("Only 1D, 2D and 3D FFTs are supported by finufft")
+        self._domain = DomainTuple.make(UnstructuredDomain((pos.shape[0])))
+        self._eps = float(eps)
+        dst = np.array(self._target[0].distances)
+        pos = (2*np.pi*pos*dst) % (2*np.pi)
+        self._eps = float(eps)
+        if len(target.shape) > 1:
+            self._pos = [pos[:, k] for k in range(pos.shape[1])]
+            s = 'nufft' + str(pos.shape[1]) + 'd'
+        else:
+            self._pos = [pos.ravel()]
+            s = 'nufft1d'
+        self._f = getattr(finufft, s+'1')
+        self._fadj = getattr(finufft, s+'2')
+
+    def apply(self, x, mode):
+        self._check_input(x, mode)
+        if mode == self.TIMES:
+            res = self._f(*self._pos, c=x.val_rw(),
+                          n_modes=self._target[0].shape, eps=self._eps).real
+        else:
+            res = self._fadj(*self._pos, f=x.val, eps=self._eps)
+            if res.ndim == 0:
+                res = np.array([res])
+        return makeField(self._tgt(mode), res)

src/library/special_distributions.py

@@ -136,7 +136,6 @@ class UniformOperator(Operator):
         The domain on which the field shall be defined. This is at the same
         time the domain and the target of the operator.
     loc: float
-
     scale: float
 
     """

src/operators/simple_linear_operators.py

@@ -141,7 +141,7 @@ class FieldAdapter(LinearOperator):
 
     Parameters
     ----------
-    tgt : Domain, tuple of Domain, DomainTuple, dict or MultiDomain:
+    tgt : Domain, tuple of domains, DomainTuple, dict or MultiDomain:
         If this is a Domain, tuple of Domain or DomainTuple, this will be the
         operator's target, and its domain will be a MultiDomain consisting of
         its domain with the supplied `name`

test/test_operators/test_interpolators.py

+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+# Copyright(C) 2013-2021 Max-Planck-Society
+#
+# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
+
+import numpy as np
+import pytest
+
+import nifty7 as ift
+
+pmp = pytest.mark.parametrize
+
+
+def test_grid_points():
+    res = 64
+    vol = 2
+    sp = ift.RGSpace([res, res], [vol / res, vol / res])
+    mg = np.mgrid[(slice(0, res),) * 2]
+    mg = np.array(list(map(np.ravel, mg)))
+
+    dist = [list(sp.distances)]
+    dist = np.array(dist).reshape(-1, 1)
+
+    sampling_points = dist * mg
+    R = ift.LinearInterpolator(sp, sampling_points)
+
+    ift.extra.check_linear_operator(R, atol=1e-7, rtol=1e-7)
+    inp = ift.from_random(R.domain)
+    out = R(inp).val
+    np.testing.assert_allclose(out, inp.val.reshape(-1), rtol=1e-7)

test/test_operators/test_nft.py

@@ -11,7 +11,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright(C) 2019 Max-Planck-Society
+# Copyright(C) 2019-2021 Max-Planck-Society
 #
 # NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
 
@@ -29,6 +29,13 @@ def _l2error(a, b):
     return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))
 
 
+def _finufft_available():
+    try:
+        import finufft
+    except ImportError:
+        pytest.skip()
+
+
 @pmp('eps', [1e-2, 1e-4, 1e-7, 1e-10, 1e-11, 1e-12, 2e-13])
 @pmp('nxdirty', [32, 128])
 @pmp('nydirty', [32, 48, 128])
@@ -101,3 +108,109 @@ def test_build(nxdirty, nydirty, N, eps):
     # We set rtol=eps here, because the gridder operator only guarantees
     # adjointness to this accuracy.
     ift.extra.check_linear_operator(RF, cmplx, flt, only_r_linear=True, rtol=eps)
+
+
+@pmp('eps', [1e-2, 1e-4, 1e-7, 1e-10, 1e-11, 1e-12, 2e-13])
+@pmp('nxdirty', [32, 128])
+@pmp('N', [1, 10, 100])
+def test_finu1d(nxdirty, N, eps):
+    _finufft_available()
+    pos = ift.random.current_rng().random((N)) - 0.5
+    vis = (ift.random.current_rng().standard_normal(N)
+           + 1j*ift.random.current_rng().standard_normal(N))
+
+    if N > 2:
+        pos[-1] = 0
+        pos[-2] = 1e-5
+    # Nifty
+    dom = ift.RGSpace((nxdirty), distances=0.2)
+    dstx = dom.distances
+    pos = pos / dstx
+    Op = ift.FinuFFT(dom, pos=pos, eps=eps)
+    vis2 = ift.makeField(ift.UnstructuredDomain(vis.shape), vis)
+    pynu = Op(vis2).val
+    # DFT
+    x = -nxdirty/2 + np.arange(nxdirty)
+
+    dft = pynu*0
+    for i in range(N):
+        dft += (vis[i]*np.exp(2j*np.pi*(x*pos[i]*dstx))).real
+    ift.myassert(_l2error(dft, pynu) < eps*10)
+
+
+@pmp('eps', [1e-2, 1e-4, 1e-7, 1e-10, 1e-11, 1e-12, 2e-13])
+@pmp('nxdirty', [32, 128])
+@pmp('nydirty', [32, 48, 128])
+@pmp('N', [1, 10, 100])
+def test_finu2d(nxdirty, nydirty, N, eps):
+    _finufft_available()
+    uv = ift.random.current_rng().random((N, 2)) - 0.5
+    vis = (ift.random.current_rng().standard_normal(N)
+           + 1j*ift.random.current_rng().standard_normal(N))
+
+    if N > 2:
+        uv[-1] = 0
+        uv[-2] = 1e-5
+    # Nifty
+    dom = ift.RGSpace((nxdirty, nydirty), distances=(0.2, 1.12))
+    dstx, dsty = dom.distances
+    uv[:, 0] = uv[:, 0]/dstx
+    uv[:, 1] = uv[:, 1]/dsty
+    Op = ift.FinuFFT(dom, pos=uv, eps=eps)
+    vis2 = ift.makeField(ift.UnstructuredDomain(vis.shape), vis)
+
+    pynu = Op(vis2).val
+    # DFT
+    x, y = np.meshgrid(
+        *[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]], indexing='ij')
+    dft = pynu*0.
+    for i in range(N):
+        dft += (
+            vis[i]*np.exp(2j*np.pi*(x*uv[i, 0]*dstx + y*uv[i, 1]*dsty))).real
+    ift.myassert(_l2error(dft, pynu) < eps*10)
+
+
+@pmp('eps', [1e-2, 1e-4, 1e-7, 1e-11])
+@pmp('nxdirty', [32, 128])
+@pmp('nydirty', [32, 48])
+@pmp('nzdirty', [32, 54])
+@pmp('N', [1, 10])
+def test_finu3d(nxdirty, nydirty, nzdirty, N, eps):
+    _finufft_available()
+    pos = ift.random.current_rng().random((N, 3)) - 0.5
+    vis = (ift.random.current_rng().standard_normal(N)
+           + 1j*ift.random.current_rng().standard_normal(N))
+    # Nifty
+    dom = ift.RGSpace((nxdirty, nydirty, nzdirty), distances=(0.2, 1.12, 0.7))
+    dstx, dsty, dstz = dom.distances
+    pos[:, 0] = pos[:, 0]/dstx
+    pos[:, 1] = pos[:, 1]/dsty
+    pos[:, 2] = pos[:, 2]/dstz
+    Op = ift.FinuFFT(dom, pos=pos, eps=eps)
+    vis2 = ift.makeField(ift.UnstructuredDomain(vis.shape), vis)
+
+    pynu = Op(vis2).val
+    # DFT
+    x, y, z = np.meshgrid(
+        *[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty, nzdirty]], indexing='ij')
+    dft = pynu*0.
+    for i in range(N):
+        dft += (
+            vis[i]*np.exp(2j*np.pi*(x*pos[i, 0]*dstx + y*pos[i, 1]*dsty + z*pos[i, 2]*dstz))).real
+    ift.myassert(_l2error(dft, pynu) < eps*10)
+
+
+@pmp('eps', [1e-2, 1e-6, 2e-13])
+@pmp('space', [ift.RGSpace(128),
+               ift.RGSpace([32, 64]),
+               ift.RGSpace([4, 27, 32])])
+@pmp('N', [1, 10, 100])
+def test_build_finufft(space, N, eps):
+    _finufft_available()
+    pos = ift.random.current_rng().random((N, len(space.shape))) - 0.5
+    RF = ift.FinuFFT(space, pos=pos, eps=eps)
+    flt = np.float64
+    cmplx = np.complex128
+    # We set rtol=eps here, because the gridder operator only guarantees
+    # adjointness to this accuracy.
+    ift.extra.check_linear_operator(RF, cmplx, flt, only_r_linear=True, rtol=eps)