junifer 0.0.5.dev240__py3-none-any.whl → 0.0.6__py3-none-any.whl

junifer/data/{tests → masks/tests}/test_masks.py

@@ -7,7 +7,7 @@
 
 import socket
 from pathlib import Path
-from typing import Callable, Dict, List, Optional, Union
+from typing import Callable, Optional, Union
 
 import nibabel as nib
 import numpy as np
@@ -20,15 +20,11 @@ from nilearn.masking import (
 )
 from numpy.testing import assert_array_almost_equal, assert_array_equal
 
-from junifer.data.masks import (
-    _available_masks,
+from junifer.data import MaskRegistry
+from junifer.data.masks import compute_brain_mask
+from junifer.data.masks._masks import (
     _load_ukb_mask,
     _load_vickery_patil_mask,
-    compute_brain_mask,
-    get_mask,
-    list_masks,
-    load_mask,
-    register_mask,
 )
 from junifer.datagrabber import DMCC13Benchmark
 from junifer.datareader import DefaultDataReader
@@ -68,10 +64,9 @@ def test_compute_brain_mask(mask_type: str, threshold: float) -> None:
         element_data = DefaultDataReader().fit_transform(dg["sub-01"])
         mask = compute_brain_mask(
             target_data=element_data["BOLD"],
-            extra_input=None,
             mask_type=mask_type,
         )
-        assert isinstance(mask, nib.Nifti1Image)
+        assert isinstance(mask, nib.nifti1.Nifti1Image)
 
 
 @pytest.mark.skipif(
@@ -111,13 +106,13 @@ def test_compute_brain_mask_for_native(mask_type: str) -> None:
             extra_input=None,
             mask_type=mask_type,
         )
-        assert isinstance(mask, nib.Nifti1Image)
+        assert isinstance(mask, nib.nifti1.Nifti1Image)
 
 
-def test_register_mask_built_in_check() -> None:
+def test_register_built_in_check() -> None:
     """Test mask registration check for built-in masks."""
     with pytest.raises(ValueError, match=r"built-in mask"):
-        register_mask(
+        MaskRegistry().register(
             name="GM_prob0.2",
             mask_path="testmask.nii.gz",
             space="MNI",
@@ -125,39 +120,38 @@ def test_register_mask_built_in_check() -> None:
         )
 
 
-def test_list_masks_incorrect() -> None:
+def test_list_incorrect() -> None:
     """Test incorrect information check for list masks."""
-    masks = list_masks()
-    assert "testmask" not in masks
+    assert "testmask" not in MaskRegistry().list
 
 
-def test_register_mask_already_registered() -> None:
+def test_register_already_registered() -> None:
     """Test mask registration check for already registered."""
     # Register custom mask
-    register_mask(
+    MaskRegistry().register(
         name="testmask",
         mask_path="testmask.nii.gz",
         space="MNI",
     )
-    out = load_mask("testmask", path_only=True)
+    out = MaskRegistry().load("testmask", path_only=True)
     assert out[1] is not None
     assert out[1].name == "testmask.nii.gz"
 
     # Try registering again
     with pytest.raises(ValueError, match=r"already registered."):
-        register_mask(
+        MaskRegistry().register(
             name="testmask",
             mask_path="testmask.nii.gz",
             space="MNI",
         )
-    register_mask(
+    MaskRegistry().register(
         name="testmask",
         mask_path="testmask2.nii.gz",
         space="MNI",
         overwrite=True,
     )
 
-    out = load_mask("testmask", path_only=True)
+    out = MaskRegistry().load("testmask", path_only=True)
     assert out[1] is not None
     assert out[1].name == "testmask2.nii.gz"
 
@@ -170,7 +164,7 @@ def test_register_mask_already_registered() -> None:
         ("testmask_3", Path("testmask_3.nii.gz"), "MNI", True),
     ],
 )
-def test_register_mask(
+def test_register(
     name: str,
     mask_path: str,
     space: str,
@@ -191,17 +185,16 @@ def test_register_mask(
 
     """
     # Register custom mask
-    register_mask(
+    MaskRegistry().register(
         name=name,
         mask_path=mask_path,
         space=space,
         overwrite=overwrite,
     )
     # List available mask and check registration
-    masks = list_masks()
-    assert name in masks
+    assert name in MaskRegistry().list
     # Load registered mask
-    _, fname, mask_space = load_mask(name=name, path_only=True)
+    _, fname, mask_space = MaskRegistry().load(name=name, path_only=True)
     # Check values for registered mask
     assert fname is not None
     assert fname.name == f"{name}.nii.gz"
@@ -216,7 +209,7 @@ def test_register_mask(
         "UKB_15K_GM",
     ],
 )
-def test_list_masks_correct(mask_name: str) -> None:
+def test_list_correct(mask_name: str) -> None:
     """Test correct information check for list masks.
 
     Parameters
@@ -225,14 +218,13 @@ def test_list_masks_correct(mask_name: str) -> None:
         The parametrized mask name.
 
     """
-    masks = list_masks()
-    assert mask_name in masks
+    assert mask_name in MaskRegistry().list
 
 
-def test_load_mask_incorrect() -> None:
+def test_load_incorrect() -> None:
     """Test loading of invalid masks."""
     with pytest.raises(ValueError, match=r"not found"):
-        load_mask("wrongmask")
+        MaskRegistry().load("wrongmask")
 
 
 @pytest.mark.parametrize(
@@ -261,7 +253,7 @@ def test_load_mask_incorrect() -> None:
 def test_vickery_patil(
     name: str,
     resolution: Optional[float],
-    pixdim: List[float],
+    pixdim: list[float],
     fname: str,
 ) -> None:
     """Test Vickery-Patil mask.
@@ -278,7 +270,7 @@ def test_vickery_patil(
         The parametrized name of the mask file.
 
     """
-    mask, mask_fname, space = load_mask(name, resolution=resolution)
+    mask, mask_fname, space = MaskRegistry().load(name, resolution=resolution)
     assert_array_almost_equal(
         mask.header["pixdim"][1:4], pixdim  # type: ignore
     )
@@ -295,7 +287,7 @@ def test_vickery_patil_error() -> None:
 
 def test_ukb() -> None:
     """Test UKB mask."""
-    mask, mask_fname, space = load_mask("UKB_15K_GM", resolution=2.0)
+    mask, mask_fname, space = MaskRegistry().load("UKB_15K_GM", resolution=2.0)
     assert_array_almost_equal(mask.header["pixdim"][1:4], 2.0)  # type: ignore
     assert space == "MNI152NLin6Asym"
     assert mask_fname is not None
@@ -308,18 +300,20 @@ def test_ukb_error() -> None:
         _load_ukb_mask(name="wrong")
 
 
-def test_get_mask() -> None:
-    """Test the get_mask function."""
+def test_get() -> None:
+    """Test tailored mask fetch."""
     with OasisVBMTestingDataGrabber() as dg:
         element_data = DefaultDataReader().fit_transform(dg["sub-01"])
         vbm_gm = element_data["VBM_GM"]
         vbm_gm_img = vbm_gm["data"]
-        mask = get_mask(masks="compute_brain_mask", target_data=vbm_gm)
+        mask = MaskRegistry().get(
+            masks="compute_brain_mask", target_data=vbm_gm
+        )
 
         assert mask.shape == vbm_gm_img.shape
         assert_array_equal(mask.affine, vbm_gm_img.affine)
 
-        raw_mask_callable, _, _ = load_mask(
+        raw_mask_callable, _, _ = MaskRegistry().load(
             "compute_brain_mask", resolution=1.5
         )
         raw_mask_img = raw_mask_callable(vbm_gm)  # type: ignore
@@ -338,7 +332,7 @@ def test_mask_callable() -> None:
     def ident(x):
         return x
 
-    _available_masks["identity"] = {
+    MaskRegistry()._registry["identity"] = {
         "family": "Callable",
         "func": ident,
         "space": "MNI152Lin",
@@ -347,44 +341,48 @@ def test_mask_callable() -> None:
         element_data = DefaultDataReader().fit_transform(dg["sub-01"])
         vbm_gm = element_data["VBM_GM"]
         vbm_gm_img = vbm_gm["data"]
-        mask = get_mask(masks="identity", target_data=vbm_gm)
+        mask = MaskRegistry().get(masks="identity", target_data=vbm_gm)
 
         assert_array_equal(mask.get_fdata(), vbm_gm_img.get_fdata())
 
-    del _available_masks["identity"]
+    del MaskRegistry()._registry["identity"]
 
 
-def test_get_mask_errors() -> None:
-    """Test passing wrong parameters to get_mask."""
+def test_get_errors() -> None:
+    """Test passing wrong parameters to fetch mask."""
     with OasisVBMTestingDataGrabber() as dg:
         element_data = DefaultDataReader().fit_transform(dg["sub-01"])
         vbm_gm = element_data["VBM_GM"]
         # Test wrong masks definitions (more than one key per dict)
         with pytest.raises(ValueError, match=r"only one key"):
-            get_mask(masks={"GM_prob0.2": {}, "Other": {}}, target_data=vbm_gm)
+            MaskRegistry().get(
+                masks={"GM_prob0.2": {}, "Other": {}}, target_data=vbm_gm
+            )
 
         # Test wrong masks definitions (pass paramaeters to non-callable mask)
         with pytest.raises(ValueError, match=r"callable params"):
-            get_mask(masks={"GM_prob0.2": {"param": 1}}, target_data=vbm_gm)
+            MaskRegistry().get(
+                masks={"GM_prob0.2": {"param": 1}}, target_data=vbm_gm
+            )
 
         # Pass only parameters to the intersection function
         with pytest.raises(
             ValueError, match=r" At least one mask is required."
         ):
-            get_mask(masks={"threshold": 1}, target_data=vbm_gm)
+            MaskRegistry().get(masks={"threshold": 1}, target_data=vbm_gm)
 
         # Pass parameters to the intersection function when only one mask
         with pytest.raises(
             ValueError, match=r"parameters to the intersection"
         ):
-            get_mask(
+            MaskRegistry().get(
                 masks=["compute_brain_mask", {"threshold": 1}],
                 target_data=vbm_gm,
             )
 
         # Test "inherited" masks error
         with pytest.raises(ValueError, match=r"provide `mask`"):
-            get_mask(masks="inherit", target_data=vbm_gm)
+            MaskRegistry().get(masks="inherit", target_data=vbm_gm)
 
 
 @pytest.mark.parametrize(
@@ -397,7 +395,7 @@ def test_get_mask_errors() -> None:
 def test_nilearn_compute_masks(
     mask_name: str,
     function: Callable,
-    params: Union[Dict, None],
+    params: Union[dict, None],
     resample: bool,
 ) -> None:
     """Test using nilearn compute mask functions.
@@ -425,7 +423,7 @@ def test_nilearn_compute_masks(
         else:
             mask_spec = {mask_name: params}
 
-        mask = get_mask(masks=mask_spec, target_data=bold)
+        mask = MaskRegistry().get(masks=mask_spec, target_data=bold)
 
         assert_array_equal(mask.affine, bold_img.affine)
 
@@ -443,15 +441,15 @@ def test_nilearn_compute_masks(
         assert_array_equal(mask.get_fdata(), ni_mask.get_fdata())
 
 
-def test_get_mask_inherit() -> None:
-    """Test using the inherit mask functionality."""
+def test_get_inherit() -> None:
+    """Test mask fetch using the inherit mask functionality."""
     with SPMAuditoryTestingDataGrabber() as dg:
         element_data = DefaultDataReader().fit_transform(dg["sub001"])
         # Compute brain mask using nilearn
         gm_mask = compute_brain_mask(element_data["BOLD"], threshold=0.2)
 
         # Get mask using the compute_brain_mask function
-        mask1 = get_mask(
+        mask1 = MaskRegistry().get(
             masks={"compute_brain_mask": {"threshold": 0.2}},
             target_data=element_data["BOLD"],
         )
@@ -463,7 +461,7 @@ def test_get_mask_inherit() -> None:
             "data": gm_mask,
             "space": element_data["BOLD"]["space"],
         }
-        mask2 = get_mask(
+        mask2 = MaskRegistry().get(
             masks="inherit",
             target_data=bold_dict,
         )
@@ -479,8 +477,8 @@ def test_get_mask_inherit() -> None:
         (["compute_brain_mask", "compute_epi_mask"], {}),
     ],
 )
-def test_get_mask_multiple(
-    masks: Union[str, Dict, List[Union[Dict, str]]], params: Dict
+def test_get_multiple(
+    masks: Union[str, dict, list[Union[dict, str]]], params: dict
 ) -> None:
     """Test getting multiple masks.
 
@@ -505,7 +503,7 @@ def test_get_mask_multiple(
         target_img = element_data["BOLD"]["data"]
         resolution = np.min(target_img.header.get_zooms()[:3])
 
-        computed = get_mask(
+        computed = MaskRegistry().get(
             masks=junifer_masks, target_data=element_data["BOLD"]
         )
 
@@ -516,16 +514,18 @@ def test_get_mask_multiple(
         mask_funcs = [
             x
             for x in masks_names
-            if _available_masks[x]["family"] == "Callable"
+            if MaskRegistry()._registry[x]["family"] == "Callable"
         ]
         mask_files = [
             x
             for x in masks_names
-            if _available_masks[x]["family"] != "Callable"
+            if MaskRegistry()._registry[x]["family"] != "Callable"
         ]
 
         mask_imgs = [
-            load_mask(t_mask, path_only=False, resolution=resolution)[0]
+            MaskRegistry().load(
+                t_mask, path_only=False, resolution=resolution
+            )[0]
             for t_mask in mask_files
         ]
 
@@ -533,10 +533,14 @@ def test_get_mask_multiple(
             # Bypass for custom mask
             if t_func == "compute_brain_mask":
                 mask_imgs.append(
-                    _available_masks[t_func]["func"](element_data["BOLD"])
+                    MaskRegistry()._registry[t_func]["func"](
+                        element_data["BOLD"]
+                    )
                 )
             else:
-                mask_imgs.append(_available_masks[t_func]["func"](target_img))
+                mask_imgs.append(
+                    MaskRegistry()._registry[t_func]["func"](target_img)
+                )
 
         mask_imgs = [
             resample_to_img(

junifer/data/parcellations/__init__.py

@@ -0,0 +1,9 @@
+"""Parcellations."""
+
+# Authors: Synchon Mandal <s.mandal@fz-juelich.de>
+# License: AGPL
+
+import lazy_loader as lazy
+
+
+__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__)

junifer/data/parcellations/__init__.pyi

@@ -0,0 +1,6 @@
+__all__ = [
+    "ParcellationRegistry",
+    "merge_parcellations",
+]
+
+from ._parcellations import ParcellationRegistry, merge_parcellations

junifer/data/parcellations/_ants_parcellation_warper.py

@@ -0,0 +1,166 @@
+"""Provide class for parcellation space warping via ANTs."""
+
+# Authors: Synchon Mandal <s.mandal@fz-juelich.de>
+# License: AGPL
+
+import uuid
+from typing import TYPE_CHECKING, Any, Optional
+
+import nibabel as nib
+
+from ...pipeline import WorkDirManager
+from ...utils import logger, raise_error, run_ext_cmd
+from ..template_spaces import get_template, get_xfm
+
+
+if TYPE_CHECKING:
+    from nibabel.nifti1 import Nifti1Image
+
+
+__all__ = ["ANTsParcellationWarper"]
+
+
+class ANTsParcellationWarper:
+    """Class for parcellation space warping via ANTs.
+
+    This class uses ANTs ``antsApplyTransforms`` for transformation.
+
+    """
+
+    def warp(
+        self,
+        parcellation_name: str,
+        parcellation_img: "Nifti1Image",
+        src: str,
+        dst: str,
+        target_data: dict[str, Any],
+        warp_data: Optional[dict[str, Any]],
+    ) -> "Nifti1Image":
+        """Warp ``parcellation_img`` to correct space.
+
+        Parameters
+        ----------
+        parcellation_name : str
+            The name of the parcellation.
+        parcellation_img : nibabel.nifti1.Nifti1Image
+            The parcellation image to transform.
+        src : str
+            The data type or template space to warp from.
+            It should be empty string if ``dst="T1w"``.
+        dst : str
+            The data type or template space to warp to.
+            `"T1w"` is the only allowed data type and it uses the resampled T1w
+            found in ``target_data.reference``. The ``"reference"``
+            key is added if the :class:`.SpaceWarper` is used or if the
+            data is provided in native space.
+        target_data : dict
+            The corresponding item of the data object to which the parcellation
+            will be applied.
+        warp_data : dict or None
+            The warp data item of the data object. The value is unused if
+            ``dst!="T1w"``.
+
+        Returns
+        -------
+        nibabel.nifti1.Nifti1Image
+            The transformed parcellation image.
+
+        Raises
+        ------
+        ValueError
+            If ``warp_data`` is None when ``dst="T1w"``.
+
+        """
+        # Create element-scoped tempdir so that warped parcellation is
+        # available later as nibabel stores file path reference for
+        # loading on computation
+        prefix = (
+            f"ants_parcellation_warper_{parcellation_name}"
+            f"{'' if not src else f'_from_{src}'}_to_{dst}_"
+            f"{uuid.uuid1()}"
+        )
+        element_tempdir = WorkDirManager().get_element_tempdir(
+            prefix=prefix,
+        )
+
+        # Native space warping
+        if dst == "native":
+            # Warp data check
+            if warp_data is None:
+                raise_error("No `warp_data` provided")
+            if "reference" not in target_data:
+                raise_error("No `reference` provided")
+            if "path" not in target_data["reference"]:
+                raise_error("No `path` provided in `reference`")
+
+            logger.debug("Using ANTs for parcellation transformation")
+
+            # Save existing parcellation image to a tempfile
+            prewarp_parcellation_path = (
+                element_tempdir / "prewarp_parcellation.nii.gz"
+            )
+            nib.save(parcellation_img, prewarp_parcellation_path)
+
+            # Create a tempfile for warped output
+            warped_parcellation_path = (
+                element_tempdir / "parcellation_warped.nii.gz"
+            )
+            # Set antsApplyTransforms command
+            apply_transforms_cmd = [
+                "antsApplyTransforms",
+                "-d 3",
+                "-e 3",
+                "-n 'GenericLabel[NearestNeighbor]'",
+                f"-i {prewarp_parcellation_path.resolve()}",
+                # use resampled reference
+                f"-r {target_data['reference']['path'].resolve()}",
+                f"-t {warp_data['path'].resolve()}",
+                f"-o {warped_parcellation_path.resolve()}",
+            ]
+            # Call antsApplyTransforms
+            run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
+
+        # Template space warping
+        else:
+            logger.debug(
+                f"Using ANTs to warp parcellation from {src} to {dst}"
+            )
+
+            # Get xfm file
+            xfm_file_path = get_xfm(src=src, dst=dst)
+            # Get template space image
+            template_space_img = get_template(
+                space=dst,
+                target_img=parcellation_img,
+                extra_input=None,
+            )
+            # Save template to a tempfile
+            template_space_img_path = element_tempdir / f"{dst}_T1w.nii.gz"
+            nib.save(template_space_img, template_space_img_path)
+
+            # Save existing parcellation image to a tempfile
+            prewarp_parcellation_path = (
+                element_tempdir / "prewarp_parcellation.nii.gz"
+            )
+            nib.save(parcellation_img, prewarp_parcellation_path)
+
+            # Create a tempfile for warped output
+            warped_parcellation_path = (
+                element_tempdir / "parcellation_warped.nii.gz"
+            )
+            # Set antsApplyTransforms command
+            apply_transforms_cmd = [
+                "antsApplyTransforms",
+                "-d 3",
+                "-e 3",
+                "-n 'GenericLabel[NearestNeighbor]'",
+                f"-i {prewarp_parcellation_path.resolve()}",
+                f"-r {template_space_img_path.resolve()}",
+                f"-t {xfm_file_path.resolve()}",
+                f"-o {warped_parcellation_path.resolve()}",
+            ]
+            # Call antsApplyTransforms
+            run_ext_cmd(name="antsApplyTransforms", cmd=apply_transforms_cmd)
+
+        # Load nifti
+        return nib.load(warped_parcellation_path)

junifer/data/parcellations/_fsl_parcellation_warper.py

@@ -0,0 +1,89 @@
+"""Provide class for parcellation space warping via FSL FLIRT."""
+
+# Authors: Synchon Mandal <s.mandal@fz-juelich.de>
+# License: AGPL
+
+import uuid
+from typing import TYPE_CHECKING, Any
+
+import nibabel as nib
+
+from ...pipeline import WorkDirManager
+from ...utils import logger, run_ext_cmd
+
+
+if TYPE_CHECKING:
+    from nibabel.nifti1 import Nifti1Image
+
+
+__all__ = ["FSLParcellationWarper"]
+
+
+class FSLParcellationWarper:
+    """Class for parcellation space warping via FSL FLIRT.
+
+    This class uses FSL FLIRT's ``applywarp`` for transformation.
+
+    """
+
+    def warp(
+        self,
+        parcellation_name: str,
+        parcellation_img: "Nifti1Image",
+        target_data: dict[str, Any],
+        warp_data: dict[str, Any],
+    ) -> "Nifti1Image":
+        """Warp ``parcellation_img`` to correct space.
+
+        Parameters
+        ----------
+        parcellation_name : str
+            The name of the parcellation.
+        parcellation_img : nibabel.nifti1.Nifti1Image
+            The parcellation image to transform.
+        target_data : dict
+            The corresponding item of the data object to which the parcellation
+            will be applied.
+        warp_data : dict
+            The warp data item of the data object.
+
+        Returns
+        -------
+        nibabel.nifti1.Nifti1Image
+            The transformed parcellation image.
+
+        """
+        logger.debug("Using FSL for parcellation transformation")
+
+        # Create element-scoped tempdir so that warped parcellation is
+        # available later as nibabel stores file path reference for
+        # loading on computation
+        element_tempdir = WorkDirManager().get_element_tempdir(
+            prefix=f"fsl_parcellation_warper_{parcellation_name}_{uuid.uuid1()}"
+        )
+
+        # Save existing parcellation image to a tempfile
+        prewarp_parcellation_path = (
+            element_tempdir / "prewarp_parcellation.nii.gz"
+        )
+        nib.save(parcellation_img, prewarp_parcellation_path)
+
+        # Create a tempfile for warped output
+        warped_parcellation_path = (
+            element_tempdir / "parcellation_warped.nii.gz"
+        )
+        # Set applywarp command
+        applywarp_cmd = [
+            "applywarp",
+            "--interp=nn",
+            f"-i {prewarp_parcellation_path.resolve()}",
+            # use resampled reference
+            f"-r {target_data['reference']['path'].resolve()}",
+            f"-w {warp_data['path'].resolve()}",
+            f"-o {warped_parcellation_path.resolve()}",
+        ]
+        # Call applywarp
+        run_ext_cmd(name="applywarp", cmd=applywarp_cmd)
+
+        # Load nifti
+        return nib.load(warped_parcellation_path)