py: add type 3 testing for GPU

python/cufinufft/tests/test_basic.py

@@ -107,6 +107,44 @@ def _execute(*args, **kwargs):
         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("output_arg", OUTPUT_ARGS)
+def test_type3(to_gpu, to_cpu, dtype, dim, n_source_pts, n_target_pts, output_arg):
+    if dtype == np.float32 and dim >= 2 and min(n_source_pts, n_target_pts) > 4000:
+        pytest.xfail("Garbage result for larger numbers of pts in single precision type 3")
+        # Strangely, this does not reproduce if we isolate the single case. To
+        # trigger it, we must run many other tests preceding this test case.
+        # So it's related to some global state of the library.
+
+    complex_dtype = utils._complex_dtype(dtype)
+
+    source_pts, source_coefs, target_pts = utils.type3_problem(complex_dtype,
+            dim, n_source_pts, n_target_pts)
+
+    plan = Plan(3, dim, dtype=complex_dtype)
+
+    source_pts_gpu = to_gpu(source_pts)
+    target_pts_gpu = to_gpu(target_pts)
+
+    source_coefs_gpu = to_gpu(source_coefs)
+
+    plan.setpts(*source_pts_gpu, *((None,) * (3 - dim)), *target_pts_gpu)
+
+    if not output_arg:
+        target_coefs_gpu = plan.execute(source_coefs_gpu)
+    else:
+        target_coefs_gpu = _compat.array_empty_like(source_coefs_gpu,
+                n_target_pts, dtype=complex_dtype)
+        plan.execute(source_coefs_gpu, out=target_coefs_gpu)
+
+    target_coefs = to_cpu(target_coefs_gpu)
+
+    utils.verify_type3(source_pts, source_coefs, target_pts, target_coefs, 1e-6)
+
+
 def test_opts(to_gpu, to_cpu, shape=(8, 8, 8), M=32, tol=1e-3):
     dtype = np.float32
 

python/cufinufft/tests/utils.py

@@ -65,6 +65,16 @@ def type2_problem(dtype, shape, M, n_trans=()):
     return k, fk
 
 
+def type3_problem(dtype, dim, n_source_pts, n_target_pts, n_trans=()):
+    real_dtype = _real_dtype(dtype)
+
+    source_pts = gen_nu_pts(n_source_pts, dim=dim).astype(real_dtype)
+    source_coefs = gen_nonuniform_data(n_source_pts, n_trans=n_trans).astype(dtype)
+    target_pts = gen_nu_pts(n_target_pts, dim=dim).astype(real_dtype)
+
+    return source_pts, source_coefs, target_pts
+
+
 def make_grid(shape):
     dim = len(shape)
     shape = shape
@@ -97,6 +107,17 @@ def direct_type2(fk, k, dim):
     return c
 
 
+def direct_type3(source_pts, source_coefs, target_pts, ind):
+    target_pt = target_pts[:, ind[-1]]
+    target_pt = target_pt[:, np.newaxis]
+
+    _source_coef = source_coefs[ind[:-1]]
+
+    target_coef = np.sum(np.exp(1j * np.sum(target_pt * source_pts, axis=0)) * _source_coef)
+
+    return target_coef
+
+
 def verify_type1(k, c, fk, tol):
     dim = fk.ndim - (c.ndim - 1)
 
@@ -128,6 +149,25 @@ def verify_type2(k, fk, c, tol):
     assert type2_rel_err < 25 * tol
 
 
+def verify_type3(source_pts, source_coef, target_pts, target_coef, tol):
+    dim = source_pts.shape[0]
+
+    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),)
+
+    target_est = target_coef[ind]
+    target_true = direct_type3(source_pts, source_coef, target_pts, ind)
+
+    type3_rel_err = np.linalg.norm(target_est - target_true) / np.linalg.norm(target_true)
+
+    assert type3_rel_err < 100 * tol
+
+
 def transfer_funcs(module_name):
     if module_name == "pycuda":
         import pycuda.autoinit # NOQA:401