Simplified type 3. Add test for cuda type 3 nufft. Unified verify_type3 with CPU version gen_coef_ind

Author: MaximEremenko

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(np.random.randint(0, n) for n in n_tr + (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)