junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/preprocess/warping/space_warper.py

@@ -4,19 +4,49 @@
 # License: AGPL
 
 from collections.abc import Sequence
-from typing import Any, ClassVar, Optional, Union
+from enum import StrEnum
+from typing import Annotated, Any, ClassVar, Literal
 
-from templateflow import api as tflow
+from pydantic import BeforeValidator
 
 from ...api.decorators import register_preprocessor
+from ...data.template_spaces import _get_templateflow_templates
+from ...datagrabber import DataType
 from ...typing import ConditionalDependencies
-from ...utils import logger, raise_error
-from ..base import BasePreprocessor
+from ...utils import ensure_list_or_none, raise_error
+from ..base import BasePreprocessor, logger
 from ._ants_warper import ANTsWarper
 from ._fsl_warper import FSLWarper
 
 
-__all__ = ["SpaceWarper"]
+__all__ = ["SpaceWarper", "SpaceWarpingImpl"]
+
+
+class SpaceWarpingImpl(StrEnum):
+    """Accepted space warping implementations.
+
+    * ``fsl`` : FSL's ``applywarp``
+    * ``ants`` : ANTs' ``antsApplyTransforms``
+    * ``auto`` : Auto-select tool when ``reference="T1w"``
+
+    """
+
+    fsl = "fsl"
+    ants = "ants"
+    auto = "auto"
+
+
+_on = Literal[
+    DataType.T1w,
+    DataType.T2w,
+    DataType.BOLD,
+    DataType.VBM_GM,
+    DataType.VBM_WM,
+    DataType.VBM_CSF,
+    DataType.FALFF,
+    DataType.GCOR,
+    DataType.LCOR,
+]
 
 
 @register_preprocessor
@@ -25,91 +55,67 @@ class SpaceWarper(BasePreprocessor):
 
     Parameters
     ----------
-    using : {"fsl", "ants", "auto"}
-        Implementation to use for warping:
-
-        * "fsl" : Use FSL's ``applywarp``
-        * "ants" : Use ANTs' ``antsApplyTransforms``
-        * "auto" : Auto-select tool when ``reference="T1w"``
-
+    using : :enum:`.SpaceWarpingImpl`
     reference : str
         The data type to use as reference for warping, can be either a data
         type like ``"T1w"`` or a template space like ``"MNI152NLin2009cAsym"``.
         Use ``"T1w"`` for native space warping and named templates for
         template space warping.
-    on : {"T1w", "T2w", "BOLD", "VBM_GM", "VBM_WM", "VBM_CSF", "fALFF", \
-        "GCOR", "LCOR"} or list of the options
-        The data type to warp.
-
-    Raises
-    ------
-    ValueError
-        If ``using`` is invalid or
-        if ``reference`` is invalid.
+    on : {``DataType.T1w``, ``DataType.T2w``, ``DataType.BOLD``, \
+         ``DataType.VBM_GM``, ``DataType.VBM_WM``, ``DataType.VBM_CSF``, \
+         ``DataType.FALFF``, ``DataType.GCOR``, ``DataType.LCOR``} or \
+         list of them
+        The data type(s) to warp.
 
     """
 
     _CONDITIONAL_DEPENDENCIES: ClassVar[ConditionalDependencies] = [
         {
-            "using": "fsl",
-            "depends_on": FSLWarper,
+            "using": SpaceWarpingImpl.fsl,
+            "depends_on": [FSLWarper],
         },
         {
-            "using": "ants",
-            "depends_on": ANTsWarper,
+            "using": SpaceWarpingImpl.ants,
+            "depends_on": [ANTsWarper],
         },
         {
-            "using": "auto",
+            "using": SpaceWarpingImpl.auto,
             "depends_on": [FSLWarper, ANTsWarper],
         },
     ]
-    _VALID_DATA_TYPES: ClassVar[Sequence[str]] = [
-        "T1w",
-        "T2w",
-        "BOLD",
-        "VBM_GM",
-        "VBM_WM",
-        "VBM_CSF",
-        "fALFF",
-        "GCOR",
-        "LCOR",
+    _VALID_DATA_TYPES: ClassVar[Sequence[DataType]] = [
+        DataType.T1w,
+        DataType.T2w,
+        DataType.BOLD,
+        DataType.VBM_GM,
+        DataType.VBM_WM,
+        DataType.VBM_CSF,
+        DataType.FALFF,
+        DataType.GCOR,
+        DataType.LCOR,
     ]
 
-    def __init__(
-        self, using: str, reference: str, on: Union[list[str], str]
-    ) -> None:
-        """Initialize the class."""
-        # Validate `using` parameter
-        valid_using = [dep["using"] for dep in self._CONDITIONAL_DEPENDENCIES]
-        if using not in valid_using:
-            raise_error(
-                f"Invalid value for `using`, should be one of: {valid_using}"
-            )
-        self.using = using
-        self.reference = reference
-        # Set required data types based on reference and
-        # initialize superclass
-        if self.reference == "T1w":  # pragma: no cover
-            required_data_types = [self.reference, "Warp"]
-            # Listify on
-            if not isinstance(on, list):
-                on = [on]
-            # Extend required data types
-            required_data_types.extend(on)
+    using: SpaceWarpingImpl
+    reference: str
+    on: Annotated[
+        _on | list[_on],
+        BeforeValidator(ensure_list_or_none),
+    ]
 
-            super().__init__(
-                on=on,
-                required_data_types=required_data_types,
-            )
-        elif self.reference in tflow.templates():
-            super().__init__(on=on)
-        else:
+    def validate_preprocessor_params(self) -> None:
+        """Run extra logical validation for preprocessor."""
+        # Set required data types based on reference
+        if self.reference == "T1w":  # pragma: no cover
+            # Update required data types
+            self.required_data_types = [DataType.T1w, DataType.Warp]
+            self.required_data_types.extend(self.on)
+        elif self.reference not in _get_templateflow_templates():
             raise_error(f"Unknown reference: {self.reference}")
 
     def preprocess(  # noqa: C901
         self,
         input: dict[str, Any],
-        extra_input: Optional[dict[str, Any]] = None,
+        extra_input: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """Preprocess.
 
@@ -129,20 +135,21 @@ class SpaceWarper(BasePreprocessor):
         ------
         ValueError
             If ``extra_input`` is None when transforming to native space
-            i.e., using ``"T1w"`` as reference.
+            i.e., using ``"T1w"`` as reference or converting from native to
+            template space or
+            if the ``reference`` key is missing from ``input`` when converting
+            from native to template space.
         RuntimeError
             If warper could not be found in ``extra_input`` when
             ``using="auto"`` or converting from native space or
             if the data is in the correct space and does not require
             warping or
-            if FSL is used when ``reference="T1w"``.
+            if FSL or "auto" is used when ``reference!="T1w"``.
 
         """
         logger.info(f"Warping to {self.reference} space using SpaceWarper")
         # Transform to native space
-        if (
-            self.using in ["fsl", "ants", "auto"] and self.reference == "T1w"
-        ):  # pragma: no cover
+        if self.reference == "T1w":  # pragma: no cover
             # Check for extra inputs
             if extra_input is None:
                 raise_error(
@@ -150,41 +157,30 @@ class SpaceWarper(BasePreprocessor):
                     f"`{self.reference}` data types in particular."
                 )
             # Conditional preprocessor
-            if self.using == "fsl":
+            warper = None
+            if self.using == "auto":
+                for entry in extra_input["Warp"]:
+                    if entry["dst"] == "native":
+                        warper = entry["warper"]
+                if warper is None:
+                    raise_error(
+                        klass=RuntimeError, msg="Could not find correct warper"
+                    )
+            else:
+                warper = self.using
+            if warper == "fsl":
                 input = FSLWarper().preprocess(
                     input=input,
                     extra_input=extra_input,
                     reference=self.reference,
                 )
-            elif self.using == "ants":
+            elif warper == "ants":
                 input = ANTsWarper().preprocess(
                     input=input,
                     extra_input=extra_input,
                     reference=self.reference,
                 )
-            elif self.using == "auto":
-                warper = None
-                for entry in extra_input["Warp"]:
-                    if entry["dst"] == "native":
-                        warper = entry["warper"]
-                if warper is None:
-                    raise_error(
-                        klass=RuntimeError, msg="Could not find correct warper"
-                    )
-                if warper == "fsl":
-                    input = FSLWarper().preprocess(
-                        input=input,
-                        extra_input=extra_input,
-                        reference=self.reference,
-                    )
-                elif warper == "ants":
-                    input = ANTsWarper().preprocess(
-                        input=input,
-                        extra_input=extra_input,
-                        reference=self.reference,
-                    )
-        # Transform to template space
-        if self.using in ["fsl", "ants"] and self.reference != "T1w":
+        else:
             input_space = input["space"]
             # Check pre-requirements for space manipulation
             if self.using == "ants" and self.reference == input_space:
@@ -241,12 +237,12 @@ class SpaceWarper(BasePreprocessor):
                     )
             else:
                 # Transform from MNI to MNI template space not possible
-                if self.using == "fsl":
+                if self.using in ["fsl", "auto"]:
                     raise_error(
                         (
                             f"Warping from {input_space} space to "
                             f"{self.reference} space not possible with "
-                            "FSL, use ANTs instead."
+                            f"{self.using}, use ANTs instead."
                         ),
                         klass=RuntimeError,
                     )
@@ -257,5 +253,10 @@ class SpaceWarper(BasePreprocessor):
                         extra_input={},
                         reference=self.reference,
                     )
-
+        logger.debug("Completed warping step")
+        logger.debug("Warped data types: ")
+        for k, v in input.items():
+            if k in ["data", "meta"]:
+                continue
+            logger.debug(f"\t{k}: {v}")
         return input

junifer/preprocess/warping/tests/test_space_warper.py

@@ -8,25 +8,36 @@ import socket
 import pytest
 from numpy.testing import assert_array_equal, assert_raises
 
-from junifer.datagrabber import DataladHCP1200, DMCC13Benchmark
+from junifer.datagrabber import (
+    DataladHCP1200,
+    DataType,
+    DMCC13Benchmark,
+    DMCCPhaseEncoding,
+    DMCCRun,
+    DMCCSession,
+    DMCCTask,
+    HCP1200PhaseEncoding,
+    HCP1200Task,
+)
 from junifer.datareader import DefaultDataReader
 from junifer.pipeline.utils import _check_ants, _check_fsl
-from junifer.preprocess import SpaceWarper
+from junifer.preprocess import SpaceWarper, SpaceWarpingImpl
 from junifer.testing.datagrabbers import PartlyCloudyTestingDataGrabber
 from junifer.typing import DataGrabberLike
 
 
+pytestmark = pytest.mark.external
+
+
 @pytest.mark.parametrize(
     "using, reference, error_type, error_msg",
     [
-        ("jam", "T1w", ValueError, "`using`"),
-        ("ants", "juice", ValueError, "reference"),
-        ("ants", "MNI152NLin2009cAsym", RuntimeError, "remove"),
-        ("fsl", "MNI152NLin2009cAsym", RuntimeError, "ANTs"),
+        (SpaceWarpingImpl.ants, "MNI152NLin2009cAsym", RuntimeError, "remove"),
+        (SpaceWarpingImpl.fsl, "MNI152NLin2009cAsym", RuntimeError, "ANTs"),
     ],
 )
 def test_SpaceWarper_errors(
-    using: str,
+    using: SpaceWarpingImpl,
     reference: str,
     error_type: type[Exception],
     error_msg: str,
@@ -35,7 +46,7 @@ def test_SpaceWarper_errors(
 
     Parameters
     ----------
-    using : str
+    using : SpaceWarpingImpl
         The parametrized implementation method.
     reference : str
         The parametrized reference to use.
@@ -46,14 +57,12 @@ def test_SpaceWarper_errors(
 
     """
     with PartlyCloudyTestingDataGrabber() as dg:
-        # Read data
         element_data = DefaultDataReader().fit_transform(dg["sub-01"])
-        # Preprocess data
         with pytest.raises(error_type, match=error_msg):
             SpaceWarper(
                 using=using,
                 reference=reference,
-                on="BOLD",
+                on=DataType.BOLD,
             ).preprocess(
                 input=element_data["BOLD"],
                 extra_input=element_data,
@@ -65,24 +74,24 @@ def test_SpaceWarper_errors(
     [
         [
             DMCC13Benchmark(
-                types=["BOLD", "T1w", "Warp"],
-                sessions=["ses-wave1bas"],
-                tasks=["Rest"],
-                phase_encodings=["AP"],
-                runs=["1"],
+                types=[DataType.BOLD, DataType.T1w, DataType.Warp],
+                sessions=DMCCSession.Wave1Bas,
+                tasks=DMCCTask.Rest,
+                phase_encodings=DMCCPhaseEncoding.AP,
+                runs=DMCCRun.One,
                 native_t1w=True,
             ),
             ("sub-f9057kp", "ses-wave1bas", "Rest", "AP", "1"),
-            "ants",
+            SpaceWarpingImpl.ants,
         ],
         [
             DataladHCP1200(
-                tasks=["REST1"],
-                phase_encodings=["LR"],
+                tasks=HCP1200Task.REST1,
+                phase_encodings=HCP1200PhaseEncoding.LR,
                 ica_fix=True,
             ),
             ("100206", "REST1", "LR"),
-            "fsl",
+            SpaceWarpingImpl.fsl,
         ],
     ],
 )
@@ -95,7 +104,9 @@ def test_SpaceWarper_errors(
     reason="only for juseless",
 )
 def test_SpaceWarper_native(
-    datagrabber: DataGrabberLike, element: tuple[str, ...], using: str
+    datagrabber: DataGrabberLike,
+    element: tuple[str, ...],
+    using: SpaceWarpingImpl,
 ) -> None:
     """Test SpaceWarper for native space warping.
 
@@ -105,23 +116,20 @@ def test_SpaceWarper_native(
         The parametrized DataGrabber objects.
     element : tuple of str
         The parametrized elements.
-    using : str
+    using : SpaceWarpingImpl
         The parametrized implementation method.
 
     """
     with datagrabber as dg:
-        # Read data
         element_data = DefaultDataReader().fit_transform(dg[element])
-        # Preprocess data
         output = SpaceWarper(
             using=using,
             reference="T1w",
-            on="BOLD",
+            on=DataType.BOLD,
         ).preprocess(
             input=element_data["BOLD"],
             extra_input=element_data,
         )
-        # Check
         assert isinstance(output, dict)
 
 
@@ -130,11 +138,11 @@ def test_SpaceWarper_native(
     [
         [
             DMCC13Benchmark(
-                types=["T1w"],
-                sessions=["ses-wave1bas"],
-                tasks=["Rest"],
-                phase_encodings=["AP"],
-                runs=["1"],
+                types=DataType.T1w,
+                sessions=DMCCSession.Wave1Bas,
+                tasks=DMCCTask.Rest,
+                phase_encodings=DMCCPhaseEncoding.AP,
+                runs=DMCCRun.One,
                 native_t1w=False,
             ),
             ("sub-f9057kp", "ses-wave1bas", "Rest", "AP", "1"),
@@ -142,11 +150,11 @@ def test_SpaceWarper_native(
         ],
         [
             DMCC13Benchmark(
-                types=["T1w"],
-                sessions=["ses-wave1bas"],
-                tasks=["Rest"],
-                phase_encodings=["AP"],
-                runs=["1"],
+                types=DataType.T1w,
+                sessions=DMCCSession.Wave1Bas,
+                tasks=DMCCTask.Rest,
+                phase_encodings=DMCCPhaseEncoding.AP,
+                runs=DMCCRun.One,
                 native_t1w=False,
             ),
             ("sub-f9057kp", "ses-wave1bas", "Rest", "AP", "1"),
@@ -175,19 +183,16 @@ def test_SpaceWarper_multi_mni(
 
     """
     with datagrabber as dg:
-        # Read data
         element_data = DefaultDataReader().fit_transform(dg[element])
         pre_xfm_data = element_data["T1w"]["data"].get_fdata().copy()
-        # Preprocess data
         output = SpaceWarper(
-            using="ants",
+            using=SpaceWarpingImpl.ants,
             reference=space,
-            on=["T1w"],
+            on=DataType.T1w,
         ).preprocess(
             input=element_data["T1w"],
             extra_input=element_data,
         )
-        # Checks
         assert isinstance(output, dict)
         assert output["space"] == space
         with assert_raises(AssertionError):

junifer/preprocess/warping/tests/test_warping_utils.py

@@ -0,0 +1,163 @@
+"""Provide tests for the warping utilities."""
+
+# Authors: Federico Raimondo <f.raimondo@fz-juelich.de>
+# License: AGPL
+
+from pathlib import Path
+
+import nibabel as nib
+import numpy as np
+import pytest
+from numpy.testing import assert_array_equal
+
+from junifer.datagrabber import DataType
+from junifer.datareader import DefaultDataReader
+from junifer.pipeline.utils import _check_ants
+from junifer.preprocess import SpaceWarper, SpaceWarpingImpl
+from junifer.preprocess.warping._utils import warp_in_chunks
+from junifer.testing import config_override
+from junifer.testing.datagrabbers import PartlyCloudyTestingDataGrabber
+
+
+def _copy(in_path: Path, out_path: Path) -> list[str]:
+    """Return a command copying the image, to test the chunks."""
+    return ["cp", str(in_path), str(out_path)]
+
+
+@pytest.mark.parametrize("suffix", [".nii", ".nii.gz"])
+@pytest.mark.parametrize(
+    "n_volumes, chunk_size", [(10, 3), (10, 5), (10, 10), (10, 100), (1, 1)]
+)
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.int16])
+def test_warp_in_chunks(
+    tmp_path: Path,
+    suffix: str,
+    n_volumes: int,
+    chunk_size: int,
+    dtype: type,
+) -> None:
+    """Test warping in chunks gives the same image.
+
+    Parameters
+    ----------
+    tmp_path : pathlib.Path
+        The path to the test directory.
+    suffix : str
+        The parametrized file extension.
+    n_volumes : int
+        The parametrized number of volumes.
+    chunk_size : int
+        The parametrized number of volumes of the chunks.
+    dtype : type
+        The parametrized data type of the image.
+
+    """
+    rng = np.random.default_rng(0)
+    data = rng.normal(size=(4, 5, 6, n_volumes)) * 100
+    img = nib.Nifti1Image(data.astype(dtype), np.diag([2, 2, 2, 1]))
+    if dtype is np.int16:
+        # Scaled data
+        img.header.set_slope_inter(0.5, 10)
+    in_path = tmp_path / f"input{suffix}"
+    out_path = tmp_path / f"output{suffix}"
+    nib.save(img, in_path)
+    with config_override("preprocessing.warping.chunksize", chunk_size):
+        warp_in_chunks(
+            name="cp",
+            input_path=in_path,
+            output_path=out_path,
+            cmd=_copy,
+            tempdir=tmp_path,
+        )
+    out = nib.load(out_path)
+    expected = nib.load(in_path)
+    assert out.shape == expected.shape
+    if dtype is not np.int16:
+        assert out.get_data_dtype() == expected.get_data_dtype()
+    assert_array_equal(out.affine, expected.affine)
+    assert_array_equal(out.get_fdata(), expected.get_fdata())
+    # Only the input and the output are left
+    assert sorted(p.name for p in tmp_path.iterdir()) == sorted(
+        [in_path.name, out_path.name]
+    )
+
+
+def test_warp_in_chunks_3d(tmp_path: Path) -> None:
+    """Test warping a 3D image is not done in chunks.
+
+    Parameters
+    ----------
+    tmp_path : pathlib.Path
+        The path to the test directory.
+
+    """
+    img = nib.Nifti1Image(np.ones((4, 5, 6), dtype=np.float32), np.eye(4))
+    in_path = tmp_path / "input.nii.gz"
+    out_path = tmp_path / "output.nii.gz"
+    nib.save(img, in_path)
+    calls = []
+
+    def cmd(in_path: Path, out_path: Path) -> list[str]:
+        calls.append(in_path)
+        return _copy(in_path, out_path)
+
+    with config_override("preprocessing.warping.chunksize", 1):
+        warp_in_chunks(
+            name="cp",
+            input_path=in_path,
+            output_path=out_path,
+            cmd=cmd,
+            tempdir=tmp_path,
+        )
+    assert calls == [in_path]
+    assert_array_equal(nib.load(out_path).get_fdata(), img.get_fdata())
+
+
+@pytest.mark.parametrize("chunk_size", [0, -1, 2.5, "10"])
+def test_warp_in_chunks_error(tmp_path: Path, chunk_size: object) -> None:
+    """Test an invalid chunk size.
+
+    Parameters
+    ----------
+    tmp_path : pathlib.Path
+        The path to the test directory.
+    chunk_size : object
+        The parametrized invalid chunk size.
+
+    """
+    with config_override("preprocessing.warping.chunksize", chunk_size):
+        with pytest.raises(ValueError, match="must be a positive integer"):
+            warp_in_chunks(
+                name="cp",
+                input_path=tmp_path / "input.nii.gz",
+                output_path=tmp_path / "output.nii.gz",
+                cmd=_copy,
+                tempdir=tmp_path,
+            )
+
+
+@pytest.mark.external
+@pytest.mark.skipif(
+    _check_ants() is False, reason="requires ANTs to be in PATH"
+)
+def test_SpaceWarper_ants_chunks() -> None:
+    """Test warping with ANTs in chunks gives the same data."""
+    outputs = []
+    for chunk_size in (100, 3):
+        with PartlyCloudyTestingDataGrabber(n_timepoints=8) as dg:
+            element_data = DefaultDataReader().fit_transform(dg["sub-01"])
+            with config_override(
+                "preprocessing.warping.chunksize", chunk_size
+            ):
+                output = SpaceWarper(
+                    using=SpaceWarpingImpl.ants,
+                    reference="MNI152NLin6Asym",
+                    on=DataType.BOLD,
+                ).preprocess(
+                    input=element_data["BOLD"],
+                    extra_input=element_data,
+                )
+            assert output["data"].get_data_dtype() == np.float32
+            outputs.append(np.asarray(output["data"].dataobj))
+    assert outputs[0].shape[-1] == 8
+    assert_array_equal(outputs[0], outputs[1])