Python simple cuda

python/cufinufft/cufinufft/__init__.py

@@ -1,11 +1,12 @@
 from cufinufft._plan import Plan
 
-from cufinufft._simple import (nufft1d1, nufft1d2, nufft2d1, nufft2d2,
-                               nufft3d1, nufft3d2)
+from cufinufft._simple import (nufft1d1, nufft1d2, nufft1d3,
+                               nufft2d1, nufft2d2, nufft2d3,
+                               nufft3d1, nufft3d2, nufft3d3)
 
-__all__ = ["nufft1d1", "nufft1d2",
-           "nufft2d1", "nufft2d2",
-           "nufft3d1", "nufft3d2",
+__all__ = ["nufft1d1", "nufft1d2", "nufft1d3",
+           "nufft2d1", "nufft2d2", "nufft2d3",
+           "nufft3d1", "nufft3d2", "nufft3d3",
            "Plan"]
 
 __version__ = '2.4.0beta1'

python/cufinufft/cufinufft/_plan.py

@@ -332,7 +332,7 @@ def execute(self, data, out=None):
 
         _data, data_shape = _ensure_array_shape(_data, "data", req_data_shape,
                                                 allow_reshape=True)
-        if self._type == 1:
+        if self._type in [1, 3]:
             batch_shape = data_shape[:-1]
         else:
             batch_shape = data_shape[:-self._dim]

python/cufinufft/cufinufft/_simple.py

@@ -1,28 +1,37 @@
 from cufinufft import Plan, _compat
 
 def nufft1d1(x, data, n_modes=None, out=None, eps=1e-6, isign=1, **kwargs):
-    return _invoke_plan(1, 1, x, None, None, data, out, isign, eps, n_modes,
+    return _invoke_plan(1, 1, x, None, None, data, None, None, None, out, isign, eps, n_modes,
             kwargs)
 
 def nufft1d2(x, data, out=None, eps=1e-6, isign=-1, **kwargs):
-    return _invoke_plan(1, 2, x, None, None, data, out, isign, eps, None,
+    return _invoke_plan(1, 2, x, None, None, data, None, None, None, out, isign, eps, None,
             kwargs)
 
 def nufft2d1(x, y, data, n_modes=None, out=None, eps=1e-6, isign=1, **kwargs):
-    return _invoke_plan(2, 1, x, y, None, data, out, isign, eps, n_modes,
+    return _invoke_plan(2, 1, x, y, None, data, None, None, None, out, isign, eps, n_modes,
             kwargs)
 
 def nufft2d2(x, y, data, out=None, eps=1e-6, isign=-1, **kwargs):
-    return _invoke_plan(2, 2, x, y, None, data, out, isign, eps, None, kwargs)
+    return _invoke_plan(2, 2, x, y, None, data, None, None, None, out, isign, eps, None, kwargs)
 
 def nufft3d1(x, y, z, data, n_modes=None, out=None, eps=1e-6, isign=1,
         **kwargs):
-    return _invoke_plan(3, 1, x, y, z, data, out, isign, eps, n_modes, kwargs)
+    return _invoke_plan(3, 1, x, y, z, data, None, None, None, out, isign, eps, n_modes, kwargs)
 
 def nufft3d2(x, y, z, data, out=None, eps=1e-6, isign=-1, **kwargs):
-    return _invoke_plan(3, 2, x, y, z, data, out, isign, eps, None, kwargs)
+    return _invoke_plan(3, 2, x, y, z, data, None, None, None, out, isign, eps, None, kwargs)
 
-def _invoke_plan(dim, nufft_type, x, y, z, data, out, isign, eps,
+def nufft3d3(x, y, z, data, s, t, u, out=None, eps=1e-6, isign=1, **kwargs):
+    return _invoke_plan(3, 3, x, y, z, data, s, t, u, out, isign, eps, None,  kwargs)
+
+def nufft2d3(x, y, data, s, t, out=None, eps=1e-6, isign=1, **kwargs):
+    return _invoke_plan(2, 3, x, y, None, data, s, t, None, out, isign, eps, None, kwargs)
+
+def nufft1d3(x, data, s, out=None, eps=1e-6, isign=1, **kwargs):
+    return _invoke_plan(1, 3, x, None, None, data, s, None, None, out, isign, eps, None,  kwargs)
+
+def _invoke_plan(dim, nufft_type, x, y, z, data, s, t, u,  out, isign, eps,
         n_modes=None, kwargs=None):
     dtype = _compat.get_array_dtype(data)
 
@@ -32,10 +41,13 @@ def _invoke_plan(dim, nufft_type, x, y, z, data, out, isign, eps,
         n_modes = out.shape[-dim:]
     if nufft_type == 2:
         n_modes = data.shape[-dim:]
+
+    if nufft_type == 3:
+        plan = Plan(nufft_type, dim, n_trans = n_trans,  eps=eps, isign=isign, dtype=dtype, **kwargs)
+    else:
+        plan = Plan(nufft_type, n_modes, n_trans, eps, isign, dtype, **kwargs)
 
-    plan = Plan(nufft_type, n_modes, n_trans, eps, isign, dtype, **kwargs)
-
-    plan.setpts(x, y, z)
+    plan.setpts(x, y, z, s, t, u)
 
     if out is None:
         out = plan.execute(data)
@@ -46,7 +58,7 @@ def _invoke_plan(dim, nufft_type, x, y, z, data, out, isign, eps,
 
 
 def _get_ntrans(dim, nufft_type, data):
-    if nufft_type == 1:
+    if ((nufft_type == 1) | (nufft_type == 3)):
         expect_dim = 1
     else:
         expect_dim = dim

python/cufinufft/tests/test_simple.py

@@ -84,3 +84,39 @@ def test_simple_type2(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg)
     c = to_cpu(c_gpu)
 
     utils.verify_type2(k, fk, c, tol)
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("dim", list(set(len(shape) for shape in SHAPES)))
+@pytest.mark.parametrize("n_source_pts", MS)
+@pytest.mark.parametrize("n_target_pts", MS)
+@pytest.mark.parametrize("n_trans", N_TRANS)
+@pytest.mark.parametrize("tol", TOLS)
+@pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
+def test_cufinufft3_simple(to_gpu, to_cpu, dtype, dim, n_source_pts, n_target_pts, n_trans, tol, output_arg):
+    complex_dtype = utils._complex_dtype(dtype)
+
+    fun = {1: cufinufft.nufft1d3,
+           2: cufinufft.nufft2d3,
+           3: cufinufft.nufft3d3}[dim]
+
+    source_pts, source_coefs, target_pts = utils.type3_problem(
+        complex_dtype, dim, n_source_pts, n_target_pts, n_trans
+    )
+
+
+    source_pts_gpu = to_gpu(source_pts) 
+    source_coefs_gpu = to_gpu(source_coefs)
+    target_pts_gpu = to_gpu(target_pts)
+
+    if output_arg:
+        target_coefs_gpu = _compat.array_empty_like(
+            source_coefs_gpu, n_trans + (n_target_pts,), dtype=complex_dtype)
+
+        fun(*source_pts_gpu, source_coefs_gpu, *target_pts_gpu, out=target_coefs_gpu, eps=tol)
+    else: 
+        target_coefs_gpu = fun(*source_pts_gpu, source_coefs_gpu, *target_pts_gpu, eps=tol)
+
+    target_coefs = to_cpu(target_coefs_gpu)
+
+    utils.verify_type3(source_pts, source_coefs, target_pts, target_coefs, tol)

python/cufinufft/tests/utils.py

@@ -22,6 +22,9 @@ def _real_dtype(complex_dtype):
 
     return real_dtype
 
+def gen_coef_ind(n_pts, n_tr):
+    ind = tuple(n-1 for n in n_tr + (int(0.1789 *n_pts),))
+    return ind
 
 def gen_nu_pts(M, dim=3, seed=0):
     np.random.seed(seed)
@@ -155,10 +158,8 @@ def verify_type3(source_pts, source_coef, target_pts, target_coef, tol):
     n_source_pts = source_pts.shape[-1]
     n_target_pts = target_pts.shape[-1]
     n_tr = source_coef.shape[:-1]
-
     assert target_coef.shape == n_tr + (n_target_pts,)
-
-    ind = (int(0.1789 * n_target_pts),)
+    ind = gen_coef_ind(n_target_pts, n_tr)
 
     target_est = target_coef[ind]
     target_true = direct_type3(source_pts, source_coef, target_pts, ind)