Skip to content

Commit

Permalink
python: improve testing of modeord option
Browse files Browse the repository at this point in the history
Integrate it into the matrix of tests for type 1 and add it for the type
2 tests as well.
  • Loading branch information
janden committed May 28, 2024
1 parent 180029d commit dfe4465
Showing 1 changed file with 17 additions and 24 deletions.
41 changes: 17 additions & 24 deletions python/finufft/test/test_finufft_plan.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,18 @@
N_PTS = [10, 11]
DTYPES = [np.complex64, np.complex128]
OUTPUT_ARGS = [False, True]
MODEORDS = [0, 1]


@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("n_pts", N_PTS)
@pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
def test_finufft1_plan(dtype, shape, n_pts, output_arg):
@pytest.mark.parametrize("modeord", MODEORDS)
def test_finufft1_plan(dtype, shape, n_pts, output_arg, modeord):
pts, coefs = utils.type1_problem(dtype, shape, n_pts)

plan = Plan(1, shape, dtype=dtype)
plan = Plan(1, shape, dtype=dtype, modeord=modeord)

plan.setpts(*pts)

Expand All @@ -31,25 +33,34 @@ def test_finufft1_plan(dtype, shape, n_pts, output_arg):
sig = np.empty(shape, dtype=dtype)
plan.execute(coefs, out=sig)

if modeord == 1:
sig = np.fft.fftshift(sig)

utils.verify_type1(pts, coefs, shape, sig, 1e-6)


@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("n_pts", N_PTS)
@pytest.mark.parametrize("output_arg", OUTPUT_ARGS)
def test_finufft2_plan(dtype, shape, n_pts, output_arg):
@pytest.mark.parametrize("modeord", MODEORDS)
def test_finufft2_plan(dtype, shape, n_pts, output_arg, modeord):
pts, sig = utils.type2_problem(dtype, shape, n_pts)

plan = Plan(2, shape, dtype=dtype)
plan = Plan(2, shape, dtype=dtype, modeord=modeord)

plan.setpts(*pts)

if modeord == 1:
_sig = np.fft.ifftshift(sig)
else:
_sig = sig

if not output_arg:
coefs = plan.execute(sig)
coefs = plan.execute(_sig)
else:
coefs = np.empty(n_pts, dtype=dtype)
plan.execute(sig, out=coefs)
plan.execute(_sig, out=coefs)

utils.verify_type2(pts, sig, coefs, 1e-6)

Expand All @@ -76,24 +87,6 @@ def test_finufft3_plan(dtype, dim, n_source_pts, n_target_pts, output_arg):
utils.verify_type3(source_pts, source_coefs, target_pts, target_coefs, 1e-6)


def test_finufft_plan_modeord():
dtype = "complex64"
shape = (8, 8)
n_pts = 17

plan = Plan(1, shape, dtype=dtype, modeord=1)

pts, coefs = utils.type1_problem(dtype, shape, n_pts)

plan.setpts(*pts)

sig = plan.execute(coefs)

sig = np.fft.fftshift(sig)

utils.verify_type1(pts, coefs, shape, sig, 1e-6)


def test_finufft_plan_errors():
with pytest.raises(RuntimeError, match="must be single or double"):
Plan(1, (8, 8), dtype="uint32")
Expand Down

0 comments on commit dfe4465

Please sign in to comment.