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

junifer/markers/reho/_junifer_reho.py

@@ -4,7 +4,6 @@
 # License: AGPL
 
 from functools import lru_cache
-from itertools import product
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -15,12 +14,12 @@ import nibabel as nib
 import numpy as np
 import scipy as sp
 from nilearn import image as nimg
-from nilearn import masking as nmask
 
 from ...pipeline import WorkDirManager
 from ...typing import Dependencies
-from ...utils import logger, raise_error
+from ...utils import raise_error
 from ...utils.singleton import Singleton
+from ..base import logger
 
 
 if TYPE_CHECKING:
@@ -50,6 +49,7 @@ class JuniferReHo(metaclass=Singleton):
         self,
         input_path: Path,
         nneigh: int = 27,
+        mask_path: Path | None = None,
     ) -> tuple["Nifti1Image", Path]:
         """Compute ReHo map.
 
@@ -66,6 +66,11 @@ class JuniferReHo(metaclass=Singleton):
             * 125 : for 5x5 cuboidal volume
 
             (default 27).
+        mask_path : pathlib.Path or None, optional
+            Path to the mask to restrict the computation to. Voxels outside
+            the mask are excluded from every neighbourhood and set to 0 in
+            the output. If None, voxels with non-zero time series are used,
+            same as AFNI's ``3dReHo`` without ``-mask`` (default None).
 
         Returns
         -------
@@ -90,44 +95,21 @@ class JuniferReHo(metaclass=Singleton):
 
         # Get scan data
         niimg = nib.load(input_path)
-        niimg_data = niimg.get_fdata().copy()
+        niimg_data = niimg.get_fdata(caching="unchanged")
         # Get scan dimensions
-        n_x, n_y, n_z, _ = niimg_data.shape
-
-        # Get rank of every voxel across time series
-        ranks_niimg_data = sp.stats.rankdata(niimg_data, axis=-1)
-
-        # Initialize 3D array to store tied rank correction for every voxel
-        tied_rank_corrections = np.zeros((n_x, n_y, n_z), dtype=np.float64)
-        # Calculate tied rank correction for every voxel
-        for i_x, i_y, i_z in product(range(n_x), range(n_y), range(n_z)):
-            # Calculate tied rank count for every voxel across time series
-            _, tie_count = np.unique(
-                ranks_niimg_data[i_x, i_y, i_z, :],
-                return_counts=True,
-            )
-            # Calculate and store tied rank correction for every voxel across
-            # timeseries
-            tied_rank_corrections[i_x, i_y, i_z] = np.sum(
-                tie_count**3 - tie_count
-            )
+        n_x, n_y, n_z, n_t = niimg_data.shape
 
-        # Initialize 3D array to store reho map
-        reho_map = np.ones((n_x, n_y, n_z), dtype=np.float32)
-
-        # TODO(synchon): this will give incorrect results if
-        # template doesn't match, hence needs to be changed
-        # after #299 is merged
-        # Calculate whole brain mask
-        mni152_whole_brain_mask = nmask.compute_brain_mask(
-            target_img=niimg,
-            threshold=0.5,
-            mask_type="whole-brain",
-        )
-        # Convert 0 / 1 array to bool
-        logical_mni152_whole_brain_mask = (
-            mni152_whole_brain_mask.get_fdata().astype(bool)
-        )
+        # Get mask as bool array
+        if mask_path is not None:
+            logical_mask = np.squeeze(nib.load(mask_path).get_fdata()) != 0
+        else:
+            # Same as AFNI's 3dReHo without -mask
+            logical_mask = np.sum(np.abs(niimg_data), axis=-1) > 1e-6
+
+        # Get rank of every voxel across time series and tied rank correction
+        # for every voxel
+        ranks_niimg_data, tied_rank_corrections = _rank_with_ties(niimg_data)
+        del niimg_data
 
         # Create mask cluster and set start and end indices
         if nneigh in (7, 19, 27):
@@ -165,66 +147,49 @@ class JuniferReHo(metaclass=Singleton):
                 mask_cluster[2, 0, 2] = 0
                 mask_cluster[2, 2, 2] = 0
 
-            start_idx = 1
-            end_idx = 2
-
         elif nneigh == 125:
             mask_cluster = np.ones((5, 5, 5))
-            start_idx = 2
-            end_idx = 3
-
-        # Convert 0 / 1 array to bool
-        logical_mask_cluster = mask_cluster.astype(bool)
-
-        for i, j, k in product(
-            range(start_idx, n_x - (end_idx - 1)),
-            range(start_idx, n_y - (end_idx - 1)),
-            range(start_idx, n_z - (end_idx - 1)),
-        ):
-            # Get mask only for neighbourhood
-            logical_neighbourhood_mni152_whole_brain_mask = (
-                logical_mni152_whole_brain_mask[
-                    i - start_idx : i + end_idx,
-                    j - start_idx : j + end_idx,
-                    k - start_idx : k + end_idx,
-                ]
-            )
-            # Perform logical AND to get neighbourhood mask;
-            # done to take care of brain boundaries
-            neighbourhood_mask = (
-                logical_mask_cluster
-                & logical_neighbourhood_mni152_whole_brain_mask
-            )
-            # Continue if voxel is restricted by mask
-            if neighbourhood_mask[1, 1, 1] == 0:
-                continue
-
-            # Get ranks for the neighbourhood
-            neighbourhood_ranks = ranks_niimg_data[
-                i - start_idx : i + end_idx,
-                j - start_idx : j + end_idx,
-                k - start_idx : k + end_idx,
-                :,
-            ]
-            # Get tied ranks corrections for the neighbourhood
-            neighbourhood_tied_ranks_corrections = tied_rank_corrections[
-                i - start_idx : i + end_idx,
-                j - start_idx : j + end_idx,
-                k - start_idx : k + end_idx,
-            ]
-            # Mask neighbourhood ranks
-            masked_neighbourhood_ranks = neighbourhood_ranks[
-                logical_mask_cluster, :
-            ]
-            # Mask tied ranks corrections for the neighbourhood
-            masked_tied_rank_corrections = (
-                neighbourhood_tied_ranks_corrections[logical_mask_cluster]
-            )
-            # Calculate KCC
-            reho_map[i, j, k] = _kendall_w_reho(
-                timeseries_ranks=masked_neighbourhood_ranks,
-                tied_rank_corrections=masked_tied_rank_corrections,
-            )
+
+        # Exclude voxels outside the mask from every neighbourhood
+        ranks_niimg_data[~logical_mask] = 0
+        tied_rank_corrections[~logical_mask] = 0
+        # Sum over the neighbourhood of every voxel: number of voxels in the
+        # mask, ranks per timepoint and tied rank corrections; at the edges
+        # of the volume, only the neighbours within it are used
+        n_neighbours = sp.ndimage.correlate(
+            logical_mask.astype(np.float64),
+            mask_cluster,
+            mode="constant",
+        )
+        neighbourhood_rank_sums = sp.ndimage.correlate(
+            ranks_niimg_data,
+            mask_cluster[..., np.newaxis],
+            mode="constant",
+        )
+        del ranks_niimg_data
+        neighbourhood_tied_rank_corrections = sp.ndimage.correlate(
+            tied_rank_corrections,
+            mask_cluster,
+            mode="constant",
+        )
+        # Calculate Kendall's coefficient of concordance (KCC)
+        np.square(neighbourhood_rank_sums, out=neighbourhood_rank_sums)
+        numerator = (12 * np.sum(neighbourhood_rank_sums, axis=-1)) - (
+            3 * n_neighbours**2 * n_t * (n_t + 1) ** 2
+        )
+        del neighbourhood_rank_sums
+        denominator = (n_neighbours**2 * n_t * (n_t**2 - 1)) - (
+            n_neighbours * neighbourhood_tied_rank_corrections
+        )
+        # If all the time series in the neighbourhood are fully tied, the
+        # denominator is 0 and KCC is set to 0, same as AFNI's 3dReHo
+        with np.errstate(divide="ignore", invalid="ignore"):
+            kcc = np.where(denominator == 0, 0.0, numerator / denominator)
+
+        # Initialize 3D array to store reho map; voxels outside the mask stay
+        # at 0, same as AFNI's 3dReHo
+        reho_map = np.zeros((n_x, n_y, n_z), dtype=np.float32)
+        reho_map[logical_mask] = kcc[logical_mask]
 
         # Create new image like target image
         output_data = nimg.new_img_like(
@@ -239,45 +204,69 @@ class JuniferReHo(metaclass=Singleton):
         element_tempdir = WorkDirManager().get_element_tempdir(
             prefix="junifer_reho"
         )
-        output_path = element_tempdir / "output.nii.gz"
+        output_path = element_tempdir / "output.nii"
         # Save computed data to file
         nib.save(output_data, output_path)
 
         return output_data, output_path  # type: ignore
 
 
-def _kendall_w_reho(
-    timeseries_ranks: np.ndarray, tied_rank_corrections: np.ndarray
-) -> float:
-    """Calculate Kendall's coefficient of concordance (KCC) for ReHo map.
+def _rank_with_ties(data: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Rank data along the last axis and compute tied rank corrections.
 
-    ..note:: This function should only be used to calculate KCC for a ReHo map.
-             For general use, check out ``junifer.stats.kendall_w``.
+    Equivalent to ``scipy.stats.rankdata(data, axis=-1)`` (tied values get
+    their average rank and NaN propagates to the whole time series), but
+    sorting only once.
 
     Parameters
     ----------
-    timeseries_ranks : 2D numpy.ndarray
-        A matrix of ranks of a subset subject's brain voxels.
-    tied_rank_corrections : 3D numpy.ndarray
-        A 3D array consisting of the tied rank corrections for the ranks
-        of a subset subject's brain voxels.
+    data : numpy.ndarray
+        The data to rank along the last axis.
 
     Returns
     -------
-    float
-        Kendall's W (KCC) of the given timeseries matrix.
+    numpy.ndarray
+        The ranks of ``data`` along the last axis.
+    numpy.ndarray
+        The tied rank correction, the sum of ``t^3 - t`` over the groups of
+        ``t`` tied values, for every element along the last axis.
 
     """
-    m, n = timeseries_ranks.shape  # annotators X items
-
-    numerator = (12 * np.sum(np.square(np.sum(timeseries_ranks, axis=0)))) - (
-        3 * m**2 * n * (n + 1) ** 2
+    n = data.shape[-1]
+    sort_idx = np.argsort(data, axis=-1, kind="stable")
+    sorted_data = np.take_along_axis(data, sort_idx, axis=-1)
+    positions = np.arange(1, n + 1)
+    # Mark the first and last sample of every tie group in sorted order
+    is_first = np.ones(data.shape, dtype=bool)
+    is_first[..., 1:] = sorted_data[..., 1:] != sorted_data[..., :-1]
+    del sorted_data
+    is_last = np.ones(data.shape, dtype=bool)
+    is_last[..., :-1] = is_first[..., 1:]
+    # Min and max rank of the tie group of every sample in sorted order
+    min_ranks = np.maximum.accumulate(
+        np.where(is_first, positions, 0), axis=-1
     )
-    denominator = (m**2 * n * (n**2 - 1)) - (m * np.sum(tied_rank_corrections))
-
-    if denominator == 0:
-        kcc = 1.0
-    else:
-        kcc = numerator / denominator
-
-    return kcc
+    del is_first
+    max_ranks = np.flip(
+        np.minimum.accumulate(
+            np.flip(np.where(is_last, positions, n + 1), axis=-1), axis=-1
+        ),
+        axis=-1,
+    )
+    del is_last
+    # Every sample belongs to a tie group of size t = max - min + 1, and
+    # summing t^2 - 1 over the samples is the same as summing t^3 - t over
+    # the tie groups
+    tie_sizes = max_ranks - min_ranks + 1
+    tied_rank_corrections = np.sum(tie_sizes**2 - 1, axis=-1).astype(
+        np.float64
+    )
+    del tie_sizes
+    # Tied values get their average rank; put back in the original order
+    ranks = np.empty(data.shape, dtype=np.float64)
+    np.put_along_axis(ranks, sort_idx, (min_ranks + max_ranks) / 2, axis=-1)
+    # Propagate NaN to the whole time series, same as scipy.stats.rankdata
+    has_nan = np.isnan(data).any(axis=-1)
+    ranks[has_nan] = np.nan
+    tied_rank_corrections[has_nan] = np.nan
+    return ranks, tied_rank_corrections

junifer/markers/reho/reho_base.py

@@ -3,17 +3,24 @@
 # Authors: Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
+from enum import StrEnum
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
+    Annotated,
     Any,
     ClassVar,
-    Optional,
 )
 
+from pydantic import BeforeValidator
+
+from ...data import get_data
+from ...datagrabber import DataType
+from ...storage import StorageType
 from ...typing import ConditionalDependencies, MarkerInOutMappings
-from ...utils import logger, raise_error
-from ..base import BaseMarker
+from ...utils import ensure_list_or_none
+from ..base import BaseMarker, logger
+from ..utils import _save_mask
 from ._afni_reho import AFNIReHo
 from ._junifer_reho import JuniferReHo
 
@@ -21,7 +28,20 @@ from ._junifer_reho import JuniferReHo
 if TYPE_CHECKING:
     from nibabel.nifti1 import Nifti1Image
 
-__all__ = ["ReHoBase"]
+
+__all__ = ["ReHoBase", "ReHoImpl"]
+
+
+class ReHoImpl(StrEnum):
+    """Accepted ReHo implementations.
+
+    * ``junifer`` : ``junifer``'s ReHo
+    * ``afni`` : AFNI's ``3dReHo``
+
+    """
+
+    junifer = "junifer"
+    afni = "afni"
 
 
 class ReHoBase(BaseMarker):
@@ -29,57 +49,98 @@ class ReHoBase(BaseMarker):
 
     Parameters
     ----------
-    using : {"junifer", "afni"}
-        Implementation to use for computing ReHo:
-
-        * "junifer" : Use ``junifer``'s own ReHo implementation
-        * "afni" : Use AFNI's ``3dReHo``
-
-    name : str, optional
-        The name of the marker. If None, it will use the class name
-        (default None).
-
-    Raises
-    ------
-    ValueError
-        If ``using`` is invalid.
+    using : :enum:`.ReHoImpl`
+    reho_params : dict or None, optional
+        Extra parameters for computing ReHo map as a dictionary (default None).
+        If ``using=ReHoImpl.afni``, then the valid keys are:
+
+        * ``nneigh`` : {7, 19, 27}, optional (default 27)
+            Number of voxels in the neighbourhood, inclusive. Can be:
+
+            - 7 : for facewise neighbours only
+            - 19 : for face- and edge-wise neighbours
+            - 27 : for face-, edge-, and node-wise neighbors
+
+        * ``neigh_rad`` : positive float, optional
+            The radius of a desired neighbourhood (default None).
+        * ``neigh_x`` : positive float, optional
+            The semi-radius for x-axis of ellipsoidal volumes (default None).
+        * ``neigh_y`` : positive float, optional
+            The semi-radius for y-axis of ellipsoidal volumes (default None).
+        * ``neigh_z`` : positive float, optional
+            The semi-radius for z-axis of ellipsoidal volumes (default None).
+        * ``box_rad`` : positive int, optional
+            The number of voxels outward in a given cardinal direction for a
+            cubic box centered on a given voxel (default None).
+        * ``box_x`` : positive int, optional
+            The number of voxels for +/- x-axis of cuboidal volumes
+            (default None).
+        * ``box_y`` : positive int, optional
+            The number of voxels for +/- y-axis of cuboidal volumes
+            (default None).
+        * ``box_z`` : positive int, optional
+            The number of voxels for +/- z-axis of cuboidal volumes
+            (default None).
+
+        else if ``using=ReHoImpl.junifer``, then the valid keys are:
+
+        * ``nneigh`` : {7, 19, 27, 125}, optional (default 27)
+            Number of voxels in the neighbourhood, inclusive. Can be:
+
+            * 7 : for facewise neighbours only
+            * 19 : for face- and edge-wise neighbours
+            * 27 : for face-, edge-, and node-wise neighbors
+            * 125 : for 5x5 cuboidal volume
+
+    agg_method : str, optional
+        The aggregation function to use.
+        See :func:`.get_aggfunc_by_name` for options
+        (default "mean").
+    agg_method_params : dict or None, optional
+        The parameters to pass to the aggregation function.
+        See :func:`.get_aggfunc_by_name` for options (default None).
+    masks : str, dict, list of them or None, optional
+        The specification of the masks to apply. The ReHo computation is
+        restricted to the voxels in the mask, and the mask is also applied
+        to regions before extracting signals. Check
+        :ref:`Using Masks <using_masks>` for more details.
+        If None, will not apply any mask (default None).
+    name : str or None, optional
+        The name of the marker.
+        If None, it will use the class name (default None).
 
     """
 
     _CONDITIONAL_DEPENDENCIES: ClassVar[ConditionalDependencies] = [
         {
-            "using": "afni",
-            "depends_on": AFNIReHo,
+            "using": ReHoImpl.afni,
+            "depends_on": [AFNIReHo],
         },
         {
-            "using": "junifer",
-            "depends_on": JuniferReHo,
+            "using": ReHoImpl.junifer,
+            "depends_on": [JuniferReHo],
         },
     ]
 
     _MARKER_INOUT_MAPPINGS: ClassVar[MarkerInOutMappings] = {
-        "BOLD": {
-            "reho": "vector",
+        DataType.BOLD: {
+            "reho": StorageType.Vector,
         },
     }
 
-    def __init__(
-        self,
-        using: str,
-        name: Optional[str] = None,
-    ) -> None:
-        # 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
-        super().__init__(on="BOLD", name=name)
+    using: ReHoImpl
+    reho_params: dict | None = None
+    agg_method: str = "mean"
+    agg_method_params: dict | None = None
+    masks: Annotated[
+        dict | str | list[dict | str] | None,
+        BeforeValidator(ensure_list_or_none),
+    ] = None
 
     def _compute(
         self,
         input_data: dict[str, Any],
+        extra_input: dict[str, Any] | None = None,
         **reho_params: Any,
     ) -> tuple["Nifti1Image", Path]:
         """Compute voxel-wise ReHo.
@@ -95,6 +156,9 @@ class ReHoBase(BaseMarker):
         ----------
         input_data : dict
             The BOLD data as dictionary.
+        extra_input : dict, optional
+            The other fields in the pipeline data object. Used to compute
+            the mask in the space of ``input_data`` (default None).
         **reho_params : dict
             Extra keyword arguments for ReHo.
 
@@ -122,9 +186,23 @@ class ReHoBase(BaseMarker):
             estimator = AFNIReHo()
         elif self.using == "junifer":
             estimator = JuniferReHo()
+        # Get mask in the space of the input data
+        mask_path = None
+        if self.masks is not None:
+            logger.debug(f"Masking with {self.masks}")
+            mask_path = _save_mask(
+                get_data(
+                    kind="mask",
+                    names=self.masks,
+                    target_data=input_data,
+                    extra_input=extra_input,
+                ),
+                prefix="reho_mask",
+            )
         # Compute reho
         reho_map, reho_map_path = estimator.compute(  # type: ignore
             input_path=input_data["path"],
+            mask_path=mask_path,
             **reho_params,
         )
 

junifer/markers/reho/reho_maps.py

@@ -3,12 +3,13 @@
 # Authors: Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
-from typing import Any, Optional, Union
+from typing import Any
 
 import numpy as np
 
 from ...api.decorators import register_marker
-from ...utils import logger
+from ...datagrabber import DataType
+from ..base import logger
 from ..maps_aggregation import MapsAggregation
 from .reho_base import ReHoBase
 
@@ -25,15 +26,10 @@ class ReHoMaps(ReHoBase):
     maps : str
         The name of the map(s) to use.
         See :func:`.list_data` for options.
-    using : {"junifer", "afni"}
-        Implementation to use for computing ReHo:
-
-        * "junifer" : Use ``junifer``'s own ReHo implementation
-        * "afni" : Use AFNI's ``3dReHo``
-
-    reho_params : dict, optional
+    using : :enum:`.ReHoImpl`
+    reho_params : dict or None, optional
         Extra parameters for computing ReHo map as a dictionary (default None).
-        If ``using="afni"``, then the valid keys are:
+        If ``using=ReHoImpl.afni``, then the valid keys are:
 
         * ``nneigh`` : {7, 19, 27}, optional (default 27)
             Number of voxels in the neighbourhood, inclusive. Can be:
@@ -63,7 +59,7 @@ class ReHoMaps(ReHoBase):
             The number of voxels for +/- z-axis of cuboidal volumes
             (default None).
 
-        else if ``using="junifer"``, then the valid keys are:
+        else if ``using=ReHoImpl.junifer``, then the valid keys are:
 
         * ``nneigh`` : {7, 19, 27, 125}, optional (default 27)
             Number of voxels in the neighbourhood, inclusive. Can be:
@@ -73,34 +69,22 @@ class ReHoMaps(ReHoBase):
             * 27 : for face-, edge-, and node-wise neighbors
             * 125 : for 5x5 cuboidal volume
 
-    masks : str, dict or list of dict or str, optional
+    masks : str, dict, list of them or None, optional
         The specification of the masks to apply to regions before extracting
         signals. Check :ref:`Using Masks <using_masks>` for more details.
         If None, will not apply any mask (default None).
-    name : str, optional
-        The name of the marker. If None, it will use the class name
-        (default None).
+    name : str or None, optional
+        The name of the marker.
+        If None, it will use the class name (default None).
 
     """
 
-    def __init__(
-        self,
-        maps: str,
-        using: str,
-        reho_params: Optional[dict] = None,
-        masks: Union[str, dict, list[Union[dict, str]], None] = None,
-        name: Optional[str] = None,
-    ) -> None:
-        # Superclass init first to validate `using` parameter
-        super().__init__(using=using, name=name)
-        self.maps = maps
-        self.reho_params = reho_params
-        self.masks = masks
+    maps: str
 
     def compute(
         self,
         input: dict[str, Any],
-        extra_input: Optional[dict[str, Any]] = None,
+        extra_input: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """Compute.
 
@@ -132,10 +116,12 @@ class ReHoMaps(ReHoBase):
         # transformation to native space.
         if self.reho_params is not None:
             reho_map, reho_file_path = self._compute(
-                input_data=input, **self.reho_params
+                input_data=input, extra_input=extra_input, **self.reho_params
             )
         else:
-            reho_map, reho_file_path = self._compute(input_data=input)
+            reho_map, reho_file_path = self._compute(
+                input_data=input, extra_input=extra_input
+            )
 
         # Perform aggregation on reho map
         aggregation_input = dict(input.items())
@@ -144,7 +130,7 @@ class ReHoMaps(ReHoBase):
         maps_aggregation = MapsAggregation(
             maps=self.maps,
             masks=self.masks,
-            on="BOLD",
+            on=DataType.BOLD,
         ).compute(
             input=aggregation_input,
             extra_input=extra_input,