Commit 52520d0f authored by Martin Reinecke's avatar Martin Reinecke
Browse files

Merge branch 'merge' into 'NIFTy_7'

Add finufft

See merge request !600
parents 02fd701c 2f0f8636
Pipeline #94957 passed with stages
in 24 minutes and 57 seconds
......@@ -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/*
......
......@@ -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
......
......@@ -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)
......@@ -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
"""
......
......@@ -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`
......
# 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)
......@@ -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)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment