Skip to content

Uses complex dtypes in cufinufft py tests #705

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jun 27, 2025
Merged

Uses complex dtypes in cufinufft py tests #705

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
27fe749
Uses complex dtypes in cufinufft py tests
janden Jun 24, 2025
1a13402
Fix test_array_ordering test
janden Jun 24, 2025
c9e6a27
Fix dtype in type3 tests [py cufinufft]
janden Jun 27, 2025
21bd0d1
Remove real_dtype in test_simple [py cufinufft]
janden Jun 27, 2025
b803959
Update CHANGELOG for dtype change [py cufinufft]
janden Jun 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
8 changes: 3 additions & 5 deletions python/cufinufft/tests/test_array_ordering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
30 changes: 11 additions & 19 deletions python/cufinufft/tests/test_basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -24,22 +24,20 @@
@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
# dimensions with the same call.
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)
Expand All @@ -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

Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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])
Expand Down
26 changes: 9 additions & 17 deletions python/cufinufft/tests/test_simple.py
View file
Open in desktop
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems jenkins failed on the newly merged test_cufinufft3_simple test https://jenkins.flatironinstitute.org/blue/organizations/jenkins/finufft/detail/PR-705/4/pipeline
I guess we need to merge master to this branch do the change for test_cufinufft3_simple too.
And real_dtype seems not used anywhere.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess we need to merge master to this branch do the change for test_cufinufft3_simple too.

Ok I rebaed on top of master and made the changes. Let's see how the tests do.

And real_dtype seems not used anywhere.

I don't quite understand. It's used to generate the correct dtypes for the frequencies, no? See here:

finufft/python/cufinufft/tests/utils.py

Line 52 in c9e6a27

k = gen_nu_pts(M, dim=dim).astype(real_dtype)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great! Thanks, all the jenkins tests pass. I think it's good to merge.
I just happen to see the two variables real_dtype defined in file test_simple.py are not used.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah I see. I've removed them now. Thanks!

Will wait for tests to finish and then merge.

Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand All @@ -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.
Expand All @@ -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:
Expand All @@ -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,
Expand All @@ -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:
Expand All @@ -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)

12 changes: 6 additions & 6 deletions python/cufinufft/tests/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading