diff --git a/CHANGELOG b/CHANGELOG index 66cf704ff..fb1d9ec0d 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -3,6 +3,7 @@ If not stated, FINUFFT is assumed (cuFINUFFT <=1.3 is listed separately). Master (working towards v2.5.0), 6/24/25 +* Update Python cufinufft unit tests to use complex dtypes (Andén, #705). * Python version of 2D Poisson solve tutorial (Julius Herb, #700). * Cached the optimal thread number (# physical cores) to reduce system call overhead in repeated small transforms (YuWei-CH, #697, fixing #693). diff --git a/python/cufinufft/tests/test_array_ordering.py b/python/cufinufft/tests/test_array_ordering.py index 0fba8f8f5..1cfa1c752 100644 --- a/python/cufinufft/tests/test_array_ordering.py +++ b/python/cufinufft/tests/test_array_ordering.py @@ -7,19 +7,17 @@ import utils -def test_type1_ordering(to_gpu, to_cpu, dtype=np.float32, shape=(16, 16, 16), M=4096, tol=1e-3): - complex_dtype = utils._complex_dtype(dtype) - +def test_type1_ordering(to_gpu, to_cpu, dtype=np.complex64, shape=(16, 16, 16), M=4096, tol=1e-3): 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) + plan = Plan(1, shape, eps=tol, dtype=dtype) plan.setpts(*k_gpu) - out = np.empty(shape, dtype=complex_dtype, order="F") + out = np.empty(shape, dtype=dtype, order="F") out_gpu = to_gpu(out) diff --git a/python/cufinufft/tests/test_basic.py b/python/cufinufft/tests/test_basic.py index 01645781b..b95741b7b 100644 --- a/python/cufinufft/tests/test_basic.py +++ b/python/cufinufft/tests/test_basic.py @@ -8,7 +8,7 @@ # NOTE: Tests below fail for tolerance 1e-4 (error executing plan). -DTYPES = [np.float32, np.float64] +DTYPES = [np.complex64, np.complex128] SHAPES = [(16,), (16, 16), (16, 16, 16), (19,), (17, 19), (17, 19, 24)] MS = [256, 1024, 4096] TOLS = [1e-3, 1e-6] @@ -24,14 +24,12 @@ @pytest.mark.parametrize("output_arg", OUTPUT_ARGS) @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) k_gpu = to_gpu(k) c_gpu = to_gpu(c) - plan = Plan(1, shape, eps=tol, dtype=complex_dtype, modeord=modeord) + plan = Plan(1, shape, eps=tol, dtype=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 @@ -39,7 +37,7 @@ def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg, modeord): plan.setpts(*k_gpu) if output_arg: - fk_gpu = _compat.array_empty_like(c_gpu, shape, dtype=complex_dtype) + fk_gpu = _compat.array_empty_like(c_gpu, shape, dtype=dtype) plan.execute(c_gpu, out=fk_gpu) else: fk_gpu = plan.execute(c_gpu) @@ -59,11 +57,9 @@ def test_type1(to_gpu, to_cpu, dtype, shape, M, tol, output_arg, modeord): @pytest.mark.parametrize("contiguous", 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, modeord=modeord) + plan = Plan(2, shape, eps=tol, dtype=dtype, modeord=modeord) check_result = True @@ -96,7 +92,7 @@ def _execute(*args, **kwargs): plan.setpts(*k_gpu) if output_arg: - c_gpu = _compat.array_empty_like(fk_gpu, (M,), dtype=complex_dtype) + c_gpu = _compat.array_empty_like(fk_gpu, (M,), dtype=dtype) _execute(fk_gpu, out=c_gpu) else: c_gpu = _execute(fk_gpu) @@ -119,12 +115,10 @@ def test_type3(to_gpu, to_cpu, dtype, dim, n_source_pts, n_target_pts, output_ar # 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, + source_pts, source_coefs, target_pts = utils.type3_problem(dtype, dim, n_source_pts, n_target_pts) - plan = Plan(3, dim, dtype=complex_dtype) + plan = Plan(3, dim, dtype=dtype) source_pts_gpu = to_gpu(source_pts) target_pts_gpu = to_gpu(target_pts) @@ -137,7 +131,7 @@ def test_type3(to_gpu, to_cpu, dtype, dim, n_source_pts, n_target_pts, output_ar 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) + n_target_pts, dtype=dtype) plan.execute(source_coefs_gpu, out=target_coefs_gpu) target_coefs = to_cpu(target_coefs_gpu) @@ -146,17 +140,15 @@ def test_type3(to_gpu, to_cpu, dtype, dim, n_source_pts, n_target_pts, output_ar def test_opts(to_gpu, to_cpu, shape=(8, 8, 8), M=32, tol=1e-3): - dtype = np.float32 - - complex_dtype = utils._complex_dtype(dtype) + dtype = np.complex64 k, c = utils.type1_problem(dtype, shape, M) k_gpu = to_gpu(k) c_gpu = to_gpu(c) - fk_gpu = _compat.array_empty_like(c_gpu, shape, dtype=complex_dtype) + fk_gpu = _compat.array_empty_like(c_gpu, shape, dtype=dtype) - plan = Plan(1, shape, eps=tol, dtype=complex_dtype, gpu_sort=False, + plan = Plan(1, shape, eps=tol, dtype=dtype, gpu_sort=False, gpu_maxsubprobsize=10) plan.setpts(k_gpu[0], k_gpu[1], k_gpu[2]) diff --git a/python/cufinufft/tests/test_simple.py b/python/cufinufft/tests/test_simple.py index e81c066a4..b71c4b048 100644 --- a/python/cufinufft/tests/test_simple.py +++ b/python/cufinufft/tests/test_simple.py @@ -7,7 +7,7 @@ import utils -DTYPES = [np.float32, np.float64] +DTYPES = [np.complex64, np.complex128] SHAPES = [(16,), (16, 16), (16, 16, 16)] N_TRANS = [(), (1,), (2,)] MS = [256, 1024, 4096] @@ -21,9 +21,6 @@ @pytest.mark.parametrize("tol", TOLS) @pytest.mark.parametrize("output_arg", OUTPUT_ARGS) def test_simple_type1(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg): - real_dtype = dtype - complex_dtype = utils._complex_dtype(dtype) - dim = len(shape) # Select which function to call based on dimension. @@ -41,7 +38,7 @@ def test_simple_type1(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg) # batch, (1, N1, ...) for batch of size one, and (n, N1, ...) for # batch of size n. fk_gpu = _compat.array_empty_like(c_gpu, n_trans + shape, - dtype=complex_dtype) + dtype=dtype) fun(*k_gpu, c_gpu, out=fk_gpu, eps=tol) else: @@ -59,9 +56,6 @@ def test_simple_type1(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg) @pytest.mark.parametrize("tol", TOLS) @pytest.mark.parametrize("output_arg", OUTPUT_ARGS) def test_simple_type2(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg): - real_dtype = dtype - complex_dtype = utils._complex_dtype(dtype) - dim = len(shape) fun = {1: cufinufft.nufft1d2, @@ -75,7 +69,7 @@ def test_simple_type2(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg) if output_arg: c_gpu = _compat.array_empty_like(fk_gpu, n_trans + (M,), - dtype=complex_dtype) + dtype=dtype) fun(*k_gpu, fk_gpu, eps=tol, out=c_gpu) else: @@ -94,30 +88,28 @@ def test_simple_type2(to_gpu, to_cpu, dtype, shape, n_trans, M, tol, output_arg) @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 + dtype, dim, n_source_pts, n_target_pts, n_trans ) - source_pts_gpu = to_gpu(source_pts) + 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) - + source_coefs_gpu, n_trans + (n_target_pts,), dtype=dtype) + fun(*source_pts_gpu, source_coefs_gpu, *target_pts_gpu, out=target_coefs_gpu, eps=tol) - else: + 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) - \ No newline at end of file diff --git a/python/cufinufft/tests/utils.py b/python/cufinufft/tests/utils.py index cd0d4350e..923052a7c 100644 --- a/python/cufinufft/tests/utils.py +++ b/python/cufinufft/tests/utils.py @@ -46,21 +46,21 @@ def gen_nonuniform_data(M, seed=0, n_trans=()): def type1_problem(dtype, shape, M, n_trans=()): - complex_dtype = _complex_dtype(dtype) + real_dtype = _real_dtype(dtype) dim = len(shape) - k = gen_nu_pts(M, dim=dim).astype(dtype) - c = gen_nonuniform_data(M, n_trans=n_trans).astype(complex_dtype) + k = gen_nu_pts(M, dim=dim).astype(real_dtype) + c = gen_nonuniform_data(M, n_trans=n_trans).astype(dtype) return k, c def type2_problem(dtype, shape, M, n_trans=()): - complex_dtype = _complex_dtype(dtype) + real_dtype = _real_dtype(dtype) dim = len(shape) - k = gen_nu_pts(M, dim=dim).astype(dtype) - fk = gen_uniform_data(n_trans + shape).astype(complex_dtype) + k = gen_nu_pts(M, dim=dim).astype(real_dtype) + fk = gen_uniform_data(n_trans + shape).astype(dtype) return k, fk