From 67fa61eda2dea223c532c516c3ea554e9f97a647 Mon Sep 17 00:00:00 2001
From: Libin Lu <llu@flatironinstitute.org>
Date: Tue, 28 May 2024 10:18:25 -0400
Subject: [PATCH] cufinufft python: improve testing of modeord option

---
 python/cufinufft/tests/test_basic.py | 52 ++++++++--------------------
 1 file changed, 15 insertions(+), 37 deletions(-)

diff --git a/python/cufinufft/tests/test_basic.py b/python/cufinufft/tests/test_basic.py
index 6d65ac85f..b66e00f18 100644
--- a/python/cufinufft/tests/test_basic.py
+++ b/python/cufinufft/tests/test_basic.py
@@ -14,6 +14,7 @@
 TOLS = [1e-3, 1e-6]
 OUTPUT_ARGS = [False, True]
 CONTIGUOUS = [False, True]
+MODEORDS = [0, 1]
 
 
 @pytest.mark.parametrize("dtype", DTYPES)
@@ -21,7 +22,8 @@
 @pytest.mark.parametrize("M", MS)
 @pytest.mark.parametrize("tol", TOLS)
 @pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
-def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
+@pytest.mark.parametrize("modeord", MODEORDS)
+def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg, modeord):
     complex_dtype = utils._complex_dtype(dtype)
 
     k, c = utils.type1_problem(dtype, shape, M)
@@ -29,7 +31,7 @@ def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
     k_gpu = to_gpu(k)
     c_gpu = to_gpu(c)
 
-    plan = Plan(1, shape, eps=tol, dtype=complex_dtype)
+    plan = Plan(1, shape, eps=tol, dtype=complex_dtype, modeord=modeord)
 
     # Since k_gpu is an array of shape (dim, M), this will expand to
     # plan.setpts(k_gpu[0], ..., k_gpu[dim]), allowing us to handle all
@@ -43,38 +45,8 @@ def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
         fk_gpu = plan.execute(c_gpu)
 
     fk = to_cpu(fk_gpu)
-
-    utils.verify_type1(k, c, fk, tol)
-
-
-@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("shape", SHAPES)
-@pytest.mark.parametrize("M", MS)
-@pytest.mark.parametrize("tol", TOLS)
-@pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
-def test_type1_modeord(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
-    complex_dtype = utils._complex_dtype(dtype)
-
-    k, c = utils.type1_problem(dtype, shape, M)
-
-    k_gpu = to_gpu(k)
-    c_gpu = to_gpu(c)
-
-    plan = Plan(1, shape, eps=tol, dtype=complex_dtype, modeord=1)
-
-    # Since k_gpu is an array of shape (dim, M), this will expand to
-    # plan.setpts(k_gpu[0], ..., k_gpu[dim]), allowing us to handle all
-    # dimensions with the same call.
-    plan.setpts(*k_gpu)
-
-    if output_arg:
-        fk_gpu = _compat.array_empty_like(c_gpu, shape, dtype=complex_dtype)
-        plan.execute(c_gpu, out=fk_gpu)
-    else:
-        fk_gpu = plan.execute(c_gpu)
-
-    fk = to_cpu(fk_gpu)
-    fk = np.fft.fftshift(fk)
+    if modeord == 1:
+        fk = np.fft.fftshift(fk)
 
     utils.verify_type1(k, c, fk, tol)
 
@@ -85,12 +57,13 @@ def test_type1_modeord(to_gpu, to_cpu, dtype, shape, M, tol, output_arg):
 @pytest.mark.parametrize("tol", TOLS)
 @pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
 @pytest.mark.parametrize("contiguous", CONTIGUOUS)
-def test_type2(to_gpu, to_cpu, dtype, shape, M, tol, output_arg, contiguous):
+@pytest.mark.parametrize("modeord", MODEORDS)
+def test_type2(to_gpu, to_cpu, dtype, shape, M, tol, output_arg, contiguous, modeord):
     complex_dtype = utils._complex_dtype(dtype)
 
     k, fk = utils.type2_problem(dtype, shape, M)
 
-    plan = Plan(2, shape, eps=tol, dtype=complex_dtype)
+    plan = Plan(2, shape, eps=tol, dtype=complex_dtype, modeord=modeord)
 
     check_result = True
 
@@ -113,7 +86,12 @@ def _execute(*args, **kwargs):
             return plan.execute(*args, **kwargs)
 
     k_gpu = to_gpu(k)
-    fk_gpu = to_gpu(fk)
+
+    if modeord == 1:
+        _fk = np.fft.ifftshift(fk)
+    else:
+        _fk = fk
+    fk_gpu = to_gpu(_fk)
 
     plan.setpts(*k_gpu)