enh: add test checking stacked models

Related: #12.

Author: oesteban

src/nifreeze/estimator.py

@@ -117,6 +117,7 @@ def run(self, dataset: DatasetT, **kwargs) -> Self:
             )
 
         kwargs["num_threads"] = kwargs.pop("omp_nthreads", None) or kwargs.pop("num_threads", None)
+        kwargs = self._align_kwargs | kwargs
 
         dataset_length = len(dataset)
         with TemporaryDirectory() as tmp_dir:

src/nifreeze/registration/ants.py

@@ -97,7 +97,7 @@ def _prepare_registration_data(
     affine: np.ndarray,
     vol_idx: int,
     dirname: Path | str,
-    clip: str | None = None,
+    clip: str | bool | None = None,
     init_affine: np.ndarray | None = None,
 ) -> tuple[Path, Path, Path | None]:
     """
@@ -129,20 +129,21 @@ def _prepare_registration_data(
         An initialization affine (for second and further estimators).
 
     """
-    clip = clip or "none"
+
     predicted_path = Path(dirname) / f"predicted_{vol_idx:05d}.nii.gz"
     sample_path = Path(dirname) / f"sample_{vol_idx:05d}.nii.gz"
+
     _to_nifti(
         sample,
         affine,
         sample_path,
-        clip=clip.lower() in ("sample", "both"),
+        clip=str(clip).lower() in ("sample", "both", "true"),
     )
     _to_nifti(
         predicted,
         affine,
         predicted_path,
-        clip=clip.lower() in ("predicted", "both"),
+        clip=str(clip).lower() in ("predicted", "both", "true"),
     )
 
     init_path = None
@@ -232,6 +233,9 @@ def generate_command(
     movingmask_path: str | Path | list[str] | None = None,
     init_affine: str | Path | None = None,
     default: str = "b0-to-b0_level0",
+    terminal_output: str | None = None,
+    num_threads: int | None = None,
+    environ: dict | None = None,
     **kwargs,
 ) -> Registration:
     """
@@ -251,6 +255,12 @@ def generate_command(
         Initial affine transformation.
     default : :obj:`str`, optional
         Default settings configuration.
+    terminal_output : :obj:`str`, optional
+        Redirect terminal output (Nipype configuration)
+    environ : :obj:`dict`, optional
+        Add environment variables to the execution.
+    num_threads : :obj:`int`, optional
+        Set the number of threads for ANTs' execution.
     **kwargs : :obj:`dict`
         Additional parameters for ANTs registration.
 
@@ -413,11 +423,17 @@ def generate_command(
         settings["initial_moving_transform"] = str(init_affine)
 
     # Generate command line with nipype and return
-    return Registration(
+    reg_iface = Registration(
         fixed_image=str(Path(fixed_path).absolute()),
         moving_image=str(Path(moving_path).absolute()),
+        terminal_output=terminal_output,
+        environ=environ or {},
         **settings,
     )
+    if num_threads:
+        reg_iface.inputs.num_threads = num_threads
+
+    return reg_iface
 
 
 def _run_registration(
@@ -451,10 +467,11 @@ def _run_registration(
     """
 
     align_kwargs = kwargs.copy()
-    environ = align_kwargs.pop("environ", {})
+    environ = align_kwargs.pop("environ", None)
     num_threads = align_kwargs.pop("num_threads", None)
 
     if (seed := align_kwargs.pop("seed", None)) is not None:
+        environ = environ or {}
         environ["ANTS_RANDOM_SEED"] = str(seed)
 
     if "ants_config" in kwargs:
@@ -463,12 +480,11 @@ def _run_registration(
     registration = generate_command(
         fixed_path,
         moving_path,
-        terminal_output="file",
         environ=environ,
+        terminal_output="file_split",
+        num_threads=num_threads,
         **align_kwargs,
     )
-    if num_threads:
-        registration.inputs.num_threads = num_threads
 
     (dirname / f"cmd-{vol_idx:05d}.sh").write_text(registration.cmdline)
 

test/conftest.py

@@ -26,9 +26,12 @@
 from pathlib import Path
 
 import nibabel as nb
+import nitransforms as nt
 import numpy as np
 import pytest
 
+from nifreeze.data.dmri import DWI
+
 test_data_env = os.getenv("TEST_DATA_HOME", str(Path.home() / "nifreeze-tests"))
 test_output_dir = os.getenv("TEST_OUTPUT_DIR")
 test_workdir = os.getenv("TEST_WORK_DIR")
@@ -54,19 +57,19 @@ def doctest_imports(doctest_namespace):
     doctest_namespace["repodata"] = _datadir
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def outdir():
     """Determine if test artifacts should be stored somewhere or deleted."""
     return None if test_output_dir is None else Path(test_output_dir)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def datadir():
     """Return a data path outside the package's structure (i.e., large datasets)."""
     return Path(test_data_env)
 
 
-@pytest.fixture
+@pytest.fixture(scope="session")
 def repodata():
     """Return the path to this repository's test data folder."""
     return _datadir
@@ -80,6 +83,59 @@ def pytest_addoption(parser):
     )
 
 
+@pytest.fixture(scope="session")
+def motion_data(tmp_path_factory, datadir):
+    # Temporary directory for session-scoped fixtures
+    tmp_path = tmp_path_factory.mktemp("motion_test_data")
+
+    dwdata = DWI.from_filename(datadir / "dwi.h5")
+    b0nii = nb.Nifti1Image(dwdata.bzero, dwdata.affine, None)
+    masknii = nb.Nifti1Image(dwdata.brainmask.astype("uint8"), dwdata.affine, None)
+
+    # Generate a list of large-yet-plausible bulk-head motion
+    xfms = nt.linear.LinearTransformsMapping(
+        [
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=0.03, z=0.005), (0.8, 0.2, 0.2)),
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=0.02, z=0.005), (0.8, 0.2, 0.2)),
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=0.02, z=0.02), (0.4, 0.2, 0.2)),
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=-0.02, z=0.02), (0.4, 0.2, 0.2)),
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=-0.02, z=0.002), (0.0, 0.2, 0.2)),
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(y=-0.02, z=0.002), (0.0, 0.2, 0.2)),
+            nb.affines.from_matvec(nb.eulerangles.euler2mat(y=-0.01, z=0.002), (0.0, 0.4, 0.2)),
+        ],
+        reference=b0nii,
+    )
+
+    # Induce motion into dataset (i.e., apply the inverse transforms)
+    moved_nii = (~xfms).apply(b0nii, reference=b0nii)
+
+    # Save the moved dataset for debugging or further processing
+    moved_path = tmp_path / "test.nii.gz"
+    ground_truth_path = tmp_path / "ground_truth.nii.gz"
+    moved_nii.to_filename(moved_path)
+    xfms.apply(moved_nii).to_filename(ground_truth_path)
+
+    # Wrap into dataset object
+    dwi_motion = DWI(
+        dataobj=np.asanyarray(moved_nii.dataobj),
+        affine=b0nii.affine,
+        bzero=dwdata.bzero,
+        gradients=dwdata.gradients[..., : len(xfms)],
+        brainmask=dwdata.brainmask,
+    )
+
+    # Return data as a dictionary (or any format that makes sense for your tests)
+    return {
+        "b0nii": b0nii,
+        "masknii": masknii,
+        "moved_nii": moved_nii,
+        "xfms": xfms,
+        "moved_path": moved_path,
+        "ground_truth_path": ground_truth_path,
+        "moved_nifreeze": dwi_motion,
+    }
+
+
 @pytest.hookimpl(trylast=True)
 def pytest_sessionfinish(session, exitstatus):
     have_werrors = os.getenv("NIFREEZE_WERRORS", False)

test/test_integration.py

@@ -24,54 +24,22 @@
 
 from os import cpu_count
 
-import nibabel as nb
 import nitransforms as nt
-import numpy as np
 
-from nifreeze.data.dmri import DWI
 from nifreeze.estimator import Estimator
 from nifreeze.model.base import TrivialModel
 from nifreeze.registration.utils import displacements_within_mask
 
 
-def test_proximity_estimator_trivial_model(datadir, tmp_path):
+def test_proximity_estimator_trivial_model(motion_data, tmp_path):
     """Check the proximity of transforms estimated by the estimator with a trivial B0 model."""
 
-    dwdata = DWI.from_filename(datadir / "dwi.h5")
-    b0nii = nb.Nifti1Image(dwdata.bzero, dwdata.affine, None)
-    masknii = nb.Nifti1Image(dwdata.brainmask.astype(np.uint8), dwdata.affine, None)
+    b0nii = motion_data["b0nii"]
+    moved_nii = motion_data["moved_nii"]
+    xfms = motion_data["xfms"]
+    dwi_motion = motion_data["moved_nifreeze"]
 
-    # Generate a list of large-yet-plausible bulk-head motion.
-    xfms = nt.linear.LinearTransformsMapping(
-        [
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=0.03, z=0.005), (0.8, 0.2, 0.2)),
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=0.02, z=0.005), (0.8, 0.2, 0.2)),
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=0.02, z=0.02), (0.4, 0.2, 0.2)),
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=-0.02, z=0.02), (0.4, 0.2, 0.2)),
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(x=-0.02, z=0.002), (0.0, 0.2, 0.2)),
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(y=-0.02, z=0.002), (0.0, 0.2, 0.2)),
-            nb.affines.from_matvec(nb.eulerangles.euler2mat(y=-0.01, z=0.002), (0.0, 0.4, 0.2)),
-        ],
-        reference=b0nii,
-    )
-
-    # Induce motion into dataset (i.e., apply the inverse transforms)
-    moved_nii = (~xfms).apply(b0nii, reference=b0nii)
-
-    # Uncomment to see the moved dataset
-    moved_nii.to_filename(tmp_path / "test.nii.gz")
-    xfms.apply(moved_nii).to_filename(tmp_path / "ground_truth.nii.gz")
-
-    # Wrap into dataset object
-    dwi_motion = DWI(
-        dataobj=moved_nii.dataobj,
-        affine=b0nii.affine,
-        bzero=dwdata.bzero,
-        gradients=dwdata.gradients[..., : len(xfms)],
-        brainmask=dwdata.brainmask,
-    )
-
-    model = TrivialModel(dwdata)
+    model = TrivialModel(dwi_motion)
     estimator = Estimator(model)
     estimator.run(
         dwi_motion,
@@ -89,9 +57,29 @@ def test_proximity_estimator_trivial_model(datadir, tmp_path):
     for i, est in enumerate(dwi_motion.motion_affines):
         assert (
             displacements_within_mask(
-                masknii,
+                motion_data["masknii"],
                 nt.linear.Affine(est),
                 xfms[i],
             ).max()
             < 0.25
         )
+
+
+def test_stacked_estimators(motion_data):
+    """Check that models can be stacked."""
+
+    # Wrap into dataset object
+    dmri_dataset = motion_data["moved_nifreeze"]
+
+    estimator1 = Estimator(
+        TrivialModel(dmri_dataset),
+        ants_config="dwi-to-dwi_level0.json",
+        clip=False,
+    )
+    estimator2 = Estimator(
+        TrivialModel(dmri_dataset),
+        prev=estimator1,
+        clip=False,
+    )
+
+    estimator2.run(dmri_dataset)