TST: Add a fixture to create random PET data for tests

Add a fixture to create random PET data for tests.

Adopt the fixture across PET tests.

Take advantage of the commit to adopt `dataset` as the name of the
created PET dataset variable across tests.

Author: jhlegarreta

pyproject.toml

@@ -223,6 +223,7 @@ env = "PYTHONHASHSEED=0"
 markers = [
     "random_gtab_data: Custom marker for random gtab data tests",
     "random_dwi_data: Custom marker for random dwi data tests",
+    "random_pet_data: Custom marker for random pet data tests",
     "random_uniform_ndim_data: Custom marker for random multi-dimensional data tests",
     "random_uniform_spatial_data: Custom marker for random spatial data tests",
 ]

test/conftest.py

@@ -297,3 +297,31 @@ def setup_random_dwi_data(request, setup_random_gtab_data):
         gradients,
         b0_thres,
     )
+
+
+@pytest.fixture(autouse=True)
+def setup_random_pet_data(request):
+    """Automatically generate random PET data for tests."""
+    marker = request.node.get_closest_marker("random_pet_data")
+
+    n_frames = 5
+    vol_size = (4, 4, 4)
+    midframe = np.arange(n_frames, dtype=np.float32) + 1
+    total_duration = float(n_frames + 1)
+    if marker:
+        n_frames, vol_size, midframe, total_duration = marker.args
+
+    rng = request.node.rng
+
+    pet_dataobj, affine = _generate_random_uniform_spatial_data(
+        request, (*vol_size, n_frames), 0.0, 1.0
+    )
+    brainmask_dataobj = rng.choice([True, False], size=vol_size).astype(np.uint8)
+
+    return (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    )

test/test_pet_data.py

@@ -24,23 +24,24 @@ def test_from_nii_requires_frame_time(setup_random_uniform_spatial_data, tmp_pat
         from_nii(fname)
 
 
-def _create_dataset():
-    rng = np.random.default_rng(12345)
-    data = rng.random((4, 4, 4, 5), dtype=np.float32)
-    affine = np.eye(4, dtype=np.float32)
-    mask = np.ones((4, 4, 4), dtype=bool)
-    midframe = np.array([10, 20, 30, 40, 50], dtype=np.float32)
-    return PET(
-        dataobj=data,
+@pytest.mark.random_pet_data(5, (4, 4, 4), (10, 20, 30, 40, 50), 60.0)
+def test_pet_set_transform_updates_motion_affines(setup_random_pet_data):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
         affine=affine,
-        brainmask=mask,
+        brainmask=brainmask_dataobj,
         midframe=midframe,
-        total_duration=60.0,
+        total_duration=total_duration,
     )
 
-
-def test_pet_set_transform_updates_motion_affines():
-    dataset = _create_dataset()
     idx = 2
     data_before = np.copy(dataset.dataobj[..., idx])
 

test/test_pet_model.py

@@ -5,23 +5,24 @@
 from nifreeze.model.pet import PETModel
 
 
-def _create_dataset():
-    rng = np.random.default_rng(12345)
-    data = rng.random((4, 4, 4, 5), dtype=np.float32)
-    affine = np.eye(4, dtype=np.float32)
-    mask = np.ones((4, 4, 4), dtype=bool)
-    midframe = np.array([10, 20, 30, 40, 50], dtype=np.float32)
-    return PET(
-        dataobj=data,
+@pytest.mark.random_pet_data(5, (4, 4, 4), (10, 20, 30, 40, 50), 60.0)
+def test_petmodel_fit_predict(setup_random_pet_data):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
         affine=affine,
-        brainmask=mask,
+        brainmask=brainmask_dataobj,
         midframe=midframe,
-        total_duration=60.0,
+        total_duration=total_duration,
     )
 
-
-def test_petmodel_fit_predict():
-    dataset = _create_dataset()
     model = PETModel(
         dataset=dataset,
         timepoints=dataset.midframe,
@@ -40,14 +41,46 @@ def test_petmodel_fit_predict():
     assert vol.dtype == dataset.dataobj.dtype
 
 
-def test_petmodel_invalid_init():
-    dataset = _create_dataset()
+@pytest.mark.random_pet_data(5, (4, 4, 4), (10, 20, 30, 40, 50), 60.0)
+def test_petmodel_invalid_init(setup_random_pet_data):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
+        affine=affine,
+        brainmask=brainmask_dataobj,
+        midframe=midframe,
+        total_duration=total_duration,
+    )
+
     with pytest.raises(TypeError):
         PETModel(dataset=dataset)
 
 
-def test_petmodel_time_check():
-    dataset = _create_dataset()
+@pytest.mark.random_pet_data(5, (4, 4, 4), (10, 20, 30, 40, 50), 60.0)
+def test_petmodel_time_check(setup_random_pet_data):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
+        affine=affine,
+        brainmask=brainmask_dataobj,
+        midframe=midframe,
+        total_duration=total_duration,
+    )
+
     bad_times = np.array([0, 10, 20, 30, 50], dtype=np.float32)
     with pytest.raises(ValueError):
         PETModel(dataset=dataset, timepoints=bad_times, xlim=60.0)

test/test_pet_workflow.py

@@ -1,50 +1,83 @@
 import types
 
 import numpy as np
+import pytest
 
 from nifreeze.data.pet import PET
 from nifreeze.estimator import PETMotionEstimator
 
 
-def _pet_dataset(n_frames=3):
-    rng = np.random.default_rng(42)
-    data = rng.random((2, 2, 2, n_frames), dtype=np.float32)
-    affine = np.eye(4, dtype=np.float32)
-    mask = np.ones((2, 2, 2), dtype=bool)
-    midframe = np.arange(n_frames, dtype=np.float32) + 1
-    return PET(
-        dataobj=data,
+@pytest.mark.random_pet_data(4, (2, 2, 2), (1.0, 2.0, 3.0, 4.0), 5.0)
+def test_lofo_split_shapes(setup_random_pet_data, tmp_path):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
         affine=affine,
-        brainmask=mask,
+        brainmask=brainmask_dataobj,
         midframe=midframe,
-        total_duration=float(n_frames + 1),
+        total_duration=total_duration,
     )
 
-
-def test_lofo_split_shapes(tmp_path):
-    ds = _pet_dataset(4)
     idx = 2
-    (train_data, train_times), (test_data, test_time) = ds.lofo_split(idx)
-    assert train_data.shape[-1] == ds.dataobj.shape[-1] - 1
-    np.testing.assert_array_equal(test_data, ds.dataobj[..., idx])
-    np.testing.assert_array_equal(train_times, np.delete(ds.midframe, idx))
-    assert test_time == ds.midframe[idx]
-
+    (train_data, train_times), (test_data, test_time) = dataset.lofo_split(idx)
+    assert train_data.shape[-1] == dataset.dataobj.shape[-1] - 1
+    np.testing.assert_array_equal(test_data, dataset.dataobj[..., idx])
+    np.testing.assert_array_equal(train_times, np.delete(dataset.midframe, idx))
+    assert test_time == dataset.midframe[idx]
+
+
+@pytest.mark.random_pet_data(3, (2, 2, 2), (1.0, 2.0, 3.0), 4.0)
+def test_to_from_filename_roundtrip(setup_random_pet_data):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
+        affine=affine,
+        brainmask=brainmask_dataobj,
+        midframe=midframe,
+        total_duration=total_duration,
+    )
 
-def test_to_from_filename_roundtrip(tmp_path):
-    ds = _pet_dataset(3)
     out_file = tmp_path / "petdata"
-    ds.to_filename(out_file)
+    dataset.to_filename(out_file)
     assert (tmp_path / "petdata.h5").exists()
     loaded = PET.from_filename(tmp_path / "petdata.h5")
-    np.testing.assert_allclose(loaded.dataobj, ds.dataobj)
-    np.testing.assert_allclose(loaded.affine, ds.affine)
-    np.testing.assert_allclose(loaded.midframe, ds.midframe)
-    assert loaded.total_duration == ds.total_duration
-
-
-def test_pet_motion_estimator_run(monkeypatch):
-    ds = _pet_dataset(3)
+    np.testing.assert_allclose(loaded.dataobj, dataset.dataobj)
+    np.testing.assert_allclose(loaded.affine, dataset.affine)
+    np.testing.assert_allclose(loaded.midframe, dataset.midframe)
+    assert loaded.total_duration == dataset.total_duration
+
+
+@pytest.mark.random_pet_data(3, (2, 2, 2), (1.0, 2.0, 3.0), 4.0)
+def test_pet_motion_estimator_run(setup_random_pet_data, monkeypatch):
+    (
+        pet_dataobj,
+        affine,
+        brainmask_dataobj,
+        midframe,
+        total_duration,
+    ) = setup_random_pet_data
+
+    dataset = PET(
+        dataobj=pet_dataobj,
+        affine=affine,
+        brainmask=brainmask_dataobj,
+        midframe=midframe,
+        total_duration=total_duration,
+    )
 
     class DummyModel:
         def __init__(self, _dataset):
@@ -67,7 +100,7 @@ def run(self, cwd=None):
     monkeypatch.setattr("nifreeze.estimator.Registration", DummyRegistration)
 
     estimator = PETMotionEstimator(None)
-    affines = estimator.run(ds)
-    assert len(affines) == len(ds)
+    affines = estimator.run(dataset)
+    assert len(affines) == len(dataset)
     for mat in affines:
         np.testing.assert_array_equal(mat, np.eye(4))