From fbcf50c144b3b28b8a45b62b72ec35f93e986357 Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Fri, 15 Nov 2024 17:01:48 +0100
Subject: [PATCH 1/8] refactor: improve interface and logic for
 compute_brain_mask

---
 junifer/data/masks/_masks.py           | 67 ++++++++++++++------------
 junifer/data/masks/tests/test_masks.py |  1 -
 2 files changed, 37 insertions(+), 31 deletions(-)

diff --git a/junifer/data/masks/_masks.py b/junifer/data/masks/_masks.py
index 86bd9cb10..01860a306 100644
--- a/junifer/data/masks/_masks.py
+++ b/junifer/data/masks/_masks.py
@@ -44,23 +44,24 @@ _masks_path = Path(__file__).parent
 
 def compute_brain_mask(
     target_data: dict[str, Any],
-    extra_input: Optional[dict[str, Any]] = None,
+    warp_data: Optional[dict[str, Any]] = None,
     mask_type: str = "brain",
     threshold: float = 0.5,
 ) -> "Nifti1Image":
     """Compute the whole-brain, grey-matter or white-matter mask.
 
     This mask is calculated using the template space and resolution as found
-    in the ``target_data``.
+    in the ``target_data``. If target space is native, then the template is
+    warped to native and then thresholded.
 
     Parameters
     ----------
     target_data : dict
         The corresponding item of the data object for which mask will be
         loaded.
-    extra_input : dict, optional
-        The other fields in the data object. Useful for accessing other data
-        types (default None).
+    warp_data : dict or None, optional
+        The warp data item of the data object. Needs to be provided if
+        ``target_data`` is in native space (default None).
     mask_type : {"brain", "gm", "wm"}, optional
         Type of mask to be computed:
 
@@ -81,7 +82,7 @@ def compute_brain_mask(
     ------
     ValueError
         If ``mask_type`` is invalid or
-        if ``extra_input`` is None when ``target_data``'s space is native.
+        if ``warp_data`` is None when ``target_data``'s space is native.
 
     """
     logger.debug(f"Computing {mask_type} mask")
@@ -90,39 +91,45 @@ def compute_brain_mask(
         raise_error(f"Unknown mask type: {mask_type}")
 
     # Check pre-requirements for space manipulation
-    target_space = target_data["space"]
-    # Set target standard space to target space
-    target_std_space = target_space
-    # Extra data type requirement check if target space is native
-    if target_space == "native":
-        # Check for extra inputs
-        if extra_input is None:
-            raise_error(
-                "No extra input provided, requires `Warp` "
-                "data type to infer target template space."
-            )
-        # Set target standard space to warp file space source
-        for entry in extra_input["Warp"]:
-            if entry["dst"] == "native":
-                target_std_space = entry["src"]
+    if target_data["space"] == "native":
+        # Warp data check
+        if warp_data is None:
+            raise_error("No `warp_data` provided")
+        # Set space to fetch template using
+        target_std_space = warp_data["src"]
+    else:
+        # Set space to fetch template using
+        target_std_space = target_data["space"]
 
-    target_img = target_data["data"]
     # Fetch template in closest resolution
     template = get_template(
         space=target_std_space,
-        target_img=target_img,
-        extra_input=extra_input,
+        target_img=target_data["data"],
+        extra_input=None,
         template_type=mask_type,
     )
-    # Resample template to target image
-    resampled_template = resample_to_img(
-        source_img=template, target_img=target_img
-    )
 
-    # Threshold and get mask
+    # Resample and warp template if target space is native
+    if target_data["space"] == "native":
+        resampled_template = ANTsMaskWarper().warp(
+            mask_name=f"template_{target_std_space}_for_compute_brain_mask",
+            # use template here
+            mask_img=template,
+            src=target_std_space,
+            dst="native",
+            target_data=target_data,
+            warp_data=warp_data,
+        )
+    # Resample template to target image
+    else:
+        resampled_template = resample_to_img(
+            source_img=template, target_img=target_data["data"]
+        )
+
+    # Threshold resampled template and get mask
     mask = (get_data(resampled_template) >= threshold).astype("int8")
 
-    return new_img_like(target_img, mask)  # type: ignore
+    return new_img_like(target_data["data"], mask)  # type: ignore
 
 
 class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
diff --git a/junifer/data/masks/tests/test_masks.py b/junifer/data/masks/tests/test_masks.py
index 65e0e8a43..a51380062 100644
--- a/junifer/data/masks/tests/test_masks.py
+++ b/junifer/data/masks/tests/test_masks.py
@@ -64,7 +64,6 @@ 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.nifti1.Nifti1Image)
-- 
2.52.0


From b47c00c5b4732ea6654680629e3aba7bedabf699 Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Fri, 15 Nov 2024 17:02:24 +0100
Subject: [PATCH 2/8] refactor: improve space warping logic for
 MaskRegistry.get

---
 junifer/data/masks/_masks.py | 123 +++++++++++++++++++----------------
 1 file changed, 67 insertions(+), 56 deletions(-)

diff --git a/junifer/data/masks/_masks.py b/junifer/data/masks/_masks.py
index 01860a306..71da67f8f 100644
--- a/junifer/data/masks/_masks.py
+++ b/junifer/data/masks/_masks.py
@@ -393,6 +393,9 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             # Set target standard space to warp file space source
             target_std_space = warper_spec["src"]
         else:
+            # Set warper_spec so that compute_brain_mask does not fail when
+            # target space is non-native
+            warper_spec = None
             # Set target standard space to target space
             target_std_space = target_space
 
@@ -405,31 +408,33 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             masks = [masks]
 
         # Check that masks passed as dicts have only one key
-        invalid_elements = [
+        invalid_mask_specs = [
             x for x in masks if isinstance(x, dict) and len(x) != 1
         ]
-        if len(invalid_elements) > 0:
+        if invalid_mask_specs:
             raise_error(
                 "Each of the masks dictionary must have only one key, "
                 "the name of the mask. The following dictionaries are "
-                f"invalid: {invalid_elements}"
+                f"invalid: {invalid_mask_specs}"
             )
 
-        # Check params for the intersection function
+        # Store params for nilearn.masking.intersect_mask()
         intersect_params = {}
-        true_masks = []
+        # Store all mask specs for further operations
+        mask_specs = []
         for t_mask in masks:
             if isinstance(t_mask, dict):
+                # Get params to pass to nilearn.masking.intersect_mask()
                 if "threshold" in t_mask:
                     intersect_params["threshold"] = t_mask["threshold"]
                     continue
-                elif "connected" in t_mask:
+                if "connected" in t_mask:
                     intersect_params["connected"] = t_mask["connected"]
                     continue
-            # All the other elements are masks
-            true_masks.append(t_mask)
+            # Add mask spec
+            mask_specs.append(t_mask)
 
-        if len(true_masks) == 0:
+        if not mask_specs:
             raise_error("No mask was passed. At least one mask is required.")
 
         # Get the nested mask data type for the input data type
@@ -437,7 +442,7 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
 
         # Get all the masks
         all_masks = []
-        for t_mask in true_masks:
+        for t_mask in mask_specs:
             if isinstance(t_mask, dict):
                 mask_name = next(iter(t_mask.keys()))
                 mask_params = t_mask[mask_name]
@@ -461,20 +466,26 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                 mask_object, _, mask_space = self.load(
                     mask_name, path_only=False, resolution=resolution
                 )
-                # Replace mask space with target space if mask's space is
-                # inherit
-                if mask_space == "inherit":
-                    mask_space = target_std_space
-                # If mask is callable like from nilearn
+                # If mask is callable like from nilearn; space will be inherit
+                # so no check for that
                 if callable(mask_object):
                     if mask_params is None:
                         mask_params = {}
                     # From nilearn
-                    if mask_name != "compute_brain_mask":
+                    if mask_name in [
+                        "compute_epi_mask",
+                        "compute_background_mask",
+                    ]:
                         mask_img = mask_object(target_img, **mask_params)
-                    # Not from nilearn
+                    # custom compute_brain_mask
+                    elif mask_name == "compute_brain_mask":
+                        mask_img = mask_object(
+                            target_data, warper_spec, **mask_params
+                        )
+                    # custom registered; arm kept for clarity
                     else:
-                        mask_img = mask_object(target_data, **mask_params)
+                        mask_img = mask_object(target_img, **mask_params)
+
                 # Mask is a Nifti1Image
                 else:
                     # Mask params provided
@@ -484,23 +495,43 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                             "Cannot pass callable params to a non-callable "
                             "mask."
                         )
-                    # Resample mask to target image
-                    mask_img = resample_to_img(
-                        source_img=mask_object,
-                        target_img=target_img,
-                        interpolation="nearest",
-                        copy=True,
-                    )
-                # Convert mask space if required
-                if mask_space != target_std_space:
-                    mask_img = ANTsMaskWarper().warp(
-                        mask_name=mask_name,
-                        mask_img=mask_img,
-                        src=mask_space,
-                        dst=target_std_space,
-                        target_data=target_data,
-                        warp_data=None,
-                    )
+                    # Resample and warp mask if target data is native
+                    if target_space == "native":
+                        mask_name = f"{mask_name}_to_native"
+                        # extra_input check done earlier and warper_spec exists
+                        if warper_spec["warper"] == "fsl":
+                            mask_img = FSLMaskWarper().warp(
+                                mask_name=mask_name,
+                                mask_img=mask_object,
+                                target_data=target_data,
+                                warp_data=warper_spec,
+                            )
+                        elif warper_spec["warper"] == "ants":
+                            mask_img = ANTsMaskWarper().warp(
+                                mask_name=mask_name,
+                                mask_img=mask_object,
+                                src="",
+                                dst="native",
+                                target_data=target_data,
+                                warp_data=warper_spec,
+                            )
+                    else:
+                        # Resample and warp mask
+                        if mask_space != target_std_space:
+                            mask_img = ANTsMaskWarper().warp(
+                                mask_name=mask_name,
+                                mask_img=mask_object,
+                                src=mask_space,
+                                dst=target_std_space,
+                                target_data=target_data,
+                                warp_data=warper_spec,
+                            )
+                        # Resample mask to target image
+                        else:
+                            mask_img = resample_to_img(
+                                source_img=mask_object,
+                                target_img=target_data["data"],
+                            )
 
             all_masks.append(mask_img)
 
@@ -510,7 +541,7 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             mask_img = intersect_masks(all_masks, **intersect_params)
         # Single mask
         else:
-            if len(intersect_params) > 0:
+            if intersect_params:
                 # Yes, I'm this strict!
                 raise_error(
                     "Cannot pass parameters to the intersection function "
@@ -518,26 +549,6 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                 )
             mask_img = all_masks[0]
 
-        # Warp mask if target data is native
-        if target_space == "native":
-            # extra_input check done earlier and warper_spec exists
-            if warper_spec["warper"] == "fsl":
-                mask_img = FSLMaskWarper().warp(
-                    mask_name="native",
-                    mask_img=mask_img,
-                    target_data=target_data,
-                    warp_data=warper_spec,
-                )
-            elif warper_spec["warper"] == "ants":
-                mask_img = ANTsMaskWarper().warp(
-                    mask_name="native",
-                    mask_img=mask_img,
-                    src="",
-                    dst="native",
-                    target_data=target_data,
-                    warp_data=warper_spec,
-                )
-
         return mask_img
 
 
-- 
2.52.0


From 99e0d6ebfe6b020adb3e1250160832958dffd784 Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Mon, 18 Nov 2024 14:49:18 +0100
Subject: [PATCH 3/8] refactor: fix mask logic

---
 junifer/data/masks/_masks.py | 43 ++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 19 deletions(-)

diff --git a/junifer/data/masks/_masks.py b/junifer/data/masks/_masks.py
index 71da67f8f..3848b6a3a 100644
--- a/junifer/data/masks/_masks.py
+++ b/junifer/data/masks/_masks.py
@@ -495,6 +495,28 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                             "Cannot pass callable params to a non-callable "
                             "mask."
                         )
+
+                    # Resample and warp mask to standard space
+                    if mask_space != target_std_space:
+                        mask_img = ANTsMaskWarper().warp(
+                            mask_name=mask_name,
+                            mask_img=mask_object,
+                            src=mask_space,
+                            dst=target_std_space,
+                            target_data=target_data,
+                            warp_data=warper_spec,
+                        )
+
+                    else:
+                        # Resample mask to target image; no further warping
+                        if target_space != "native":
+                            mask_img = resample_to_img(
+                                source_img=mask_object,
+                                target_img=target_data["data"],
+                            )
+                        # Set mask_img in case no warping happens before this
+                        else:
+                            mask_img = mask_object
                     # Resample and warp mask if target data is native
                     if target_space == "native":
                         mask_name = f"{mask_name}_to_native"
@@ -502,36 +524,19 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                         if warper_spec["warper"] == "fsl":
                             mask_img = FSLMaskWarper().warp(
                                 mask_name=mask_name,
-                                mask_img=mask_object,
+                                mask_img=mask_img,
                                 target_data=target_data,
                                 warp_data=warper_spec,
                             )
                         elif warper_spec["warper"] == "ants":
                             mask_img = ANTsMaskWarper().warp(
                                 mask_name=mask_name,
-                                mask_img=mask_object,
+                                mask_img=mask_img,
                                 src="",
                                 dst="native",
                                 target_data=target_data,
                                 warp_data=warper_spec,
                             )
-                    else:
-                        # Resample and warp mask
-                        if mask_space != target_std_space:
-                            mask_img = ANTsMaskWarper().warp(
-                                mask_name=mask_name,
-                                mask_img=mask_object,
-                                src=mask_space,
-                                dst=target_std_space,
-                                target_data=target_data,
-                                warp_data=warper_spec,
-                            )
-                        # Resample mask to target image
-                        else:
-                            mask_img = resample_to_img(
-                                source_img=mask_object,
-                                target_img=target_data["data"],
-                            )
 
             all_masks.append(mask_img)
 
-- 
2.52.0


From 1260cd34269ecba509819aef2cf55f006b9354a1 Mon Sep 17 00:00:00 2001
From: Fede Raimondo <f.raimondo@fz-juelich.de>
Date: Mon, 25 Nov 2024 11:23:10 +0100
Subject: [PATCH 4/8] Fix docstring + add debug logging

---
 junifer/data/masks/_ants_mask_warper.py      |  4 +--
 junifer/data/masks/_masks.py                 | 33 ++++++++++++++++++++
 junifer/data/parcellations/_parcellations.py | 19 ++++++++++-
 3 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/junifer/data/masks/_ants_mask_warper.py b/junifer/data/masks/_ants_mask_warper.py
index 5b91cd381..e68d78873 100644
--- a/junifer/data/masks/_ants_mask_warper.py
+++ b/junifer/data/masks/_ants_mask_warper.py
@@ -46,7 +46,7 @@ class ANTsMaskWarper:
             The mask image to transform.
         src : str
             The data type or template space to warp from.
-            It should be empty string if ``dst="T1w"``.
+            It should be empty string if ``dst="native"``.
         dst : str
             The data type or template space to warp to.
             `"native"` is the only allowed data type and it uses the resampled
@@ -58,7 +58,7 @@ class ANTsMaskWarper:
             will be applied.
         warp_data : dict or None
             The warp data item of the data object. The value is unused if
-            ``dst!="T1w"``.
+            ``dst!="native"``.
 
         Returns
         -------
diff --git a/junifer/data/masks/_masks.py b/junifer/data/masks/_masks.py
index 3848b6a3a..437f7cca9 100644
--- a/junifer/data/masks/_masks.py
+++ b/junifer/data/masks/_masks.py
@@ -376,6 +376,8 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
         """
         # Check pre-requirements for space manipulation
         target_space = target_data["space"]
+        logger.debug(f"Getting masks: {masks} in {target_space} space")
+
         # Extra data type requirement check if target space is native
         if target_space == "native":
             # Check for extra inputs
@@ -392,6 +394,9 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             )
             # Set target standard space to warp file space source
             target_std_space = warper_spec["src"]
+            logger.debug(
+                f"Target space is native. Will warp from {target_std_space}"
+            )
         else:
             # Set warper_spec so that compute_brain_mask does not fail when
             # target space is non-native
@@ -453,22 +458,40 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             # If mask is being inherited from the datagrabber or a
             # preprocessor, check that it's accessible
             if mask_name == "inherit":
+                logger.debug("Using inherited mask.")
                 if inherited_mask_item is None:
                     raise_error(
                         "Cannot inherit mask from the target data. Either the "
                         "DataGrabber or a Preprocessor does not provide "
                         "`mask` for the target data type."
                     )
+                logger.debug(
+                    f"Inherited mask is in {inherited_mask_item['space']} "
+                    "space."
+                )
                 mask_img = inherited_mask_item["data"]
+
+                if inherited_mask_item["space"] != target_space:
+                    raise_error(
+                        "Inherited mask space does not match target space."
+                    )
+                logger.debug("Resampling inherited mask to target image.")
+                # Resample inherited mask to target image
+                mask_img = resample_to_img(
+                    source_img=mask_img,
+                    target_img=target_data["data"],
+                )
             # Starting with new mask
             else:
                 # Load mask
+                logger.debug(f"Loading parcellation {t_mask}.")
                 mask_object, _, mask_space = self.load(
                     mask_name, path_only=False, resolution=resolution
                 )
                 # If mask is callable like from nilearn; space will be inherit
                 # so no check for that
                 if callable(mask_object):
+                    logger.debug("Computing mask (callable).")
                     if mask_params is None:
                         mask_params = {}
                     # From nilearn
@@ -498,6 +521,10 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
 
                     # Resample and warp mask to standard space
                     if mask_space != target_std_space:
+                        logger.debug(
+                            f"Warping {t_mask} to {target_std_space} space "
+                            "using ants."
+                        )
                         mask_img = ANTsMaskWarper().warp(
                             mask_name=mask_name,
                             mask_img=mask_object,
@@ -509,6 +536,7 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
 
                     else:
                         # Resample mask to target image; no further warping
+                        logger.debug(f"Resampling {t_mask} to target image.")
                         if target_space != "native":
                             mask_img = resample_to_img(
                                 source_img=mask_object,
@@ -519,6 +547,10 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                             mask_img = mask_object
                     # Resample and warp mask if target data is native
                     if target_space == "native":
+                        logger.debug(
+                            "Warping mask to native space using "
+                            f"{warper_spec['warper']}."
+                        )
                         mask_name = f"{mask_name}_to_native"
                         # extra_input check done earlier and warper_spec exists
                         if warper_spec["warper"] == "fsl":
@@ -543,6 +575,7 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
         # Multiple masks, need intersection / union
         if len(all_masks) > 1:
             # Intersect / union of masks
+            logger.debug("Intersecting masks.")
             mask_img = intersect_masks(all_masks, **intersect_params)
         # Single mask
         else:
diff --git a/junifer/data/parcellations/_parcellations.py b/junifer/data/parcellations/_parcellations.py
index dfe79add3..c339dac12 100644
--- a/junifer/data/parcellations/_parcellations.py
+++ b/junifer/data/parcellations/_parcellations.py
@@ -400,6 +400,7 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
         """
         # Check pre-requirements for space manipulation
         target_space = target_data["space"]
+        logger.debug(f"Getting {parcellations} in{target_space} space.")
         # Extra data type requirement check if target space is native
         if target_space == "native":
             # Check for extra inputs
@@ -416,6 +417,9 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             )
             # Set target standard space to warp file space source
             target_std_space = warper_spec["src"]
+            logger.debug(
+                f"Target space is native. Will warp from {target_std_space}"
+            )
         else:
             # Set target standard space to target space
             target_std_space = target_space
@@ -433,6 +437,7 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
         all_labels = []
         for name in parcellations:
             # Load parcellation
+            logger.debug(f"Loading parcellation {name}")
             img, labels, _, space = self.load(
                 name=name,
                 resolution=resolution,
@@ -441,6 +446,9 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             # Convert parcellation spaces if required;
             # cannot be "native" due to earlier check
             if space != target_std_space:
+                logger.debug(
+                    f"Warping {name} to {target_std_space} space using ants."
+                )
                 raw_img = ANTsParcellationWarper().warp(
                     parcellation_name=name,
                     parcellation_img=img,
@@ -452,6 +460,7 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
                 # Remove extra dimension added by ANTs
                 img = image.math_img("np.squeeze(img)", img=raw_img)
 
+            logger.debug(f"Resampling {name} to target image.")
             # Resample parcellation to target image
             img_to_merge = image.resample_to_img(
                 source_img=img,
@@ -469,6 +478,7 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             labels = all_labels[0]
         # Parcellations are already transformed to target standard space
         else:
+            logger.debug("Merging parcellations.")
             resampled_parcellation_img, labels = merge_parcellations(
                 parcellations_list=all_parcellations,
                 parcellations_names=parcellations,
@@ -477,6 +487,10 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
 
         # Warp parcellation if target space is native
         if target_space == "native":
+            logger.debug(
+                "Warping parcellation to native space using "
+                f"{warper_spec['warper']}."
+            )
             # extra_input check done earlier and warper_spec exists
             if warper_spec["warper"] == "fsl":
                 resampled_parcellation_img = FSLParcellationWarper().warp(
@@ -1194,7 +1208,10 @@ def _retrieve_aicha(
 
     # Load labels
     labels = pd.read_csv(
-        parcellation_lname, sep="\t", header=None, skiprows=[0]  # type: ignore
+        parcellation_lname,
+        sep="\t",
+        header=None,
+        skiprows=[0],  # type: ignore
     )[0].to_list()
 
     return parcellation_fname, labels
-- 
2.52.0


From 001462f14b9cfebab312dbc625ea8893b35452f2 Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Tue, 26 Nov 2024 17:48:36 +0100
Subject: [PATCH 5/8] chore: format + lint

---
 docs/changes/newsfragments/367.enh  |  2 +-
 docs/changes/newsfragments/389.enh  |  2 +-
 docs/changes/newsfragments/391.misc |  2 +-
 docs/extending/extension.rst        | 17 ++++++++---------
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/docs/changes/newsfragments/367.enh b/docs/changes/newsfragments/367.enh
index fc0ca2502..c492e0969 100644
--- a/docs/changes/newsfragments/367.enh
+++ b/docs/changes/newsfragments/367.enh
@@ -1 +1 @@
-Allow for external python files in the ``with`` section of the yaml to import other files by `Fede Raimondo`_
\ No newline at end of file
+Allow for external python files in the ``with`` section of the yaml to import other files by `Fede Raimondo`_
diff --git a/docs/changes/newsfragments/389.enh b/docs/changes/newsfragments/389.enh
index ecc98e392..16e09ad40 100644
--- a/docs/changes/newsfragments/389.enh
+++ b/docs/changes/newsfragments/389.enh
@@ -1 +1 @@
-Add a ``cleanup`` parameter to the :class:`junifer.pipeline.WorkDirManager`, which allows to disable cleaning up temporary directories by `Fede Raimondo`_
\ No newline at end of file
+Add a ``cleanup`` parameter to the :class:`junifer.pipeline.WorkDirManager`, which allows to disable cleaning up temporary directories by `Fede Raimondo`_
diff --git a/docs/changes/newsfragments/391.misc b/docs/changes/newsfragments/391.misc
index 7a5a1ee42..84a69622f 100644
--- a/docs/changes/newsfragments/391.misc
+++ b/docs/changes/newsfragments/391.misc
@@ -1 +1 @@
-Refactor ``Singleton`` class to use a metaclass instead of a class decorator by `Fede Raimondo`_
\ No newline at end of file
+Refactor ``Singleton`` class to use a metaclass instead of a class decorator by `Fede Raimondo`_
diff --git a/docs/extending/extension.rst b/docs/extending/extension.rst
index 7ceee4bc5..2e2b52d3f 100644
--- a/docs/extending/extension.rst
+++ b/docs/extending/extension.rst
@@ -31,7 +31,7 @@ This is the ideal place to include ``junifer`` extensions.
    Some ``junifer`` commands will not consider files imported from files
    included in the ``with`` statement, unless this is known to junifer. If
    ``my_file.py`` imports ``my_other_file.py``, the ``run`` command will work,
-   but ``queue`` will not create a proper job. This is because we need to 
+   but ``queue`` will not create a proper job. This is because we need to
    let junifer know that ``my_other_file.py`` is also part of the code. To do
    so, we need to include a special function in ``my_file.py`` which tells
    ``junifer`` about the dependencies of the module:
@@ -39,14 +39,13 @@ This is the ideal place to include ``junifer`` extensions.
    .. code-block:: python
 
       def junifer_module_deps() -> List[str]:
-        """Return the dependencies of the module.
+          """Return the dependencies of the module.
 
-        Returns
-        -------
-        list of str
-            The list of dependencies.
+          Returns
+          -------
+          list of str
+              The list of dependencies.
 
-        """
-
-        return ["my_other_file.py"]
+          """
 
+          return ["my_other_file.py"]
-- 
2.52.0


From 8cd321800c43a15c6e2915f532d366ebab281d18 Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Tue, 3 Dec 2024 12:58:39 +0100
Subject: [PATCH 6/8] chore: improve log message in MaskRegistry

---
 junifer/data/masks/_masks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/junifer/data/masks/_masks.py b/junifer/data/masks/_masks.py
index 437f7cca9..a603f75e3 100644
--- a/junifer/data/masks/_masks.py
+++ b/junifer/data/masks/_masks.py
@@ -484,7 +484,7 @@ class MaskRegistry(BasePipelineDataRegistry, metaclass=Singleton):
             # Starting with new mask
             else:
                 # Load mask
-                logger.debug(f"Loading parcellation {t_mask}.")
+                logger.debug(f"Loading mask {t_mask}.")
                 mask_object, _, mask_space = self.load(
                     mask_name, path_only=False, resolution=resolution
                 )
-- 
2.52.0


From f5f7cbd9490452e7e0726e4be4494e96b49d25db Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Tue, 3 Dec 2024 12:59:01 +0100
Subject: [PATCH 7/8] chore: improve log message in ParcellationRegistry

---
 junifer/data/parcellations/_parcellations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/junifer/data/parcellations/_parcellations.py b/junifer/data/parcellations/_parcellations.py
index c339dac12..9035c8332 100644
--- a/junifer/data/parcellations/_parcellations.py
+++ b/junifer/data/parcellations/_parcellations.py
@@ -400,7 +400,7 @@ class ParcellationRegistry(BasePipelineDataRegistry, metaclass=Singleton):
         """
         # Check pre-requirements for space manipulation
         target_space = target_data["space"]
-        logger.debug(f"Getting {parcellations} in{target_space} space.")
+        logger.debug(f"Getting {parcellations} in {target_space} space.")
         # Extra data type requirement check if target space is native
         if target_space == "native":
             # Check for extra inputs
-- 
2.52.0


From c160e88ae431ed483f6f83b2126b52bf0d23b566 Mon Sep 17 00:00:00 2001
From: Synchon Mandal <synchon@protonmail.com>
Date: Tue, 3 Dec 2024 13:05:32 +0100
Subject: [PATCH 8/8] chore: add changelog 395.enh

---
 docs/changes/newsfragments/395.enh | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 docs/changes/newsfragments/395.enh

diff --git a/docs/changes/newsfragments/395.enh b/docs/changes/newsfragments/395.enh
new file mode 100644
index 000000000..e4cd97722
--- /dev/null
+++ b/docs/changes/newsfragments/395.enh
@@ -0,0 +1 @@
+Refactor the masking logic in the pipeline to account for and optimise space transformations and merges by `Synchon Mandal`_ and `Fede Raimondo`_
-- 
2.52.0