REF: Refactor filters into filtering

Refactor code blocks that perform filtering operations within models
into the `filtering` module so that the model and filter concepts are
separated, supporting more cleanly pipelining models with filters within
the new `Estimator` philosophy.

Add the percentile-based detrending feature that was temporarily removed
in commit 7322e93.

Take advantage of the commit to rename the `th_low` and `th_high`
attributes of the `AverageDWIModel` to `atol_low` and `atol_high`.

Fix the mask treatment in the robust minmax normalization computation
function: when the original mask value is not `None`, the `data[mask]`
indexing results in an error since the mask is a 1D array whose
dimension does not match the indices along dimension 0 of the data
array.

Fixes:
```
>       reshaped_data = data.reshape((-1, data.shape[-1])) if mask is None else data[mask]
E       IndexError: boolean index did not match indexed array along dimension 0; dimension is 32 but corresponding boolean dimension is 5
```

Co-authored-by: Oscar Esteban <[email protected]>

Author: jhlegarreta

src/nifreeze/data/filtering.py

@@ -28,8 +28,12 @@
 from scipy.ndimage import median_filter
 from skimage.morphology import ball
 
+from nifreeze.data.dmri import DEFAULT_CLIP_PERCENTILE
+
 DEFAULT_DTYPE = "int16"
 """The default image's data type."""
+BVAL_ATOL = 100.0
+"""b-value tolerance value."""
 
 
 def advanced_clip(
@@ -96,3 +100,161 @@ def advanced_clip(
         data = np.round(255 * data).astype(dtype)
 
     return data
+
+
+def robust_minmax_normalization(
+    data: np.ndarray,
+    mask: np.ndarray | None = None,
+    p_min: float = 5.0,
+    p_max: float = 95.0,
+    inplace: bool = False,
+) -> np.ndarray | None:
+    r"""Normalize min-max percentiles of each volume to the grand min-max
+    percentiles.
+
+    Robust min/max normalization of the volumes in the dataset following:
+
+    .. math::
+        \text{data}_{\text{normalized}} = \frac{(\text{data} - p_{min}) \cdot p_{\text{mean}}}{p_{\text{range}}} + p_{min}^{\text{mean}}
+
+    where
+
+    .. math::
+        p_{\text{range}} = p_{max} - p_{min}, \quad p_{\text{mean}} = \frac{1}{N} \sum_{i=1}^N p_{\text{range}_i}, \quad p_{min}^{\text{mean}} = \frac{1}{N} \sum_{i=1}^N p_{5_i}
+
+    If a mask is provided, only the data within the mask are considered.
+
+    Parameters
+    ----------
+    data : :obj:`~numpy.ndarray`
+        Data to be normalized.
+    mask : :obj:`~numpy.ndarray`, optional
+        Mask. If provided, only the data within the mask are considered.
+    p_min : :obj:`float`, optional
+        The lower percentile value for normalization.
+    p_max : :obj:`float`, optional
+        The upper percentile value for normalization.
+    inplace : :obj:`bool`, optional
+        If ``False``, the normalization is performed on the original data.
+
+    Returns
+    -------
+    data : :obj:`~numpy.ndarray` or None
+        Normalized data or ``None`` if ``inplace`` is ``True``.
+    """
+
+    normalized = data if inplace else data.copy()
+
+    mask = mask if mask is not None else np.ones(data.shape[-1], dtype=bool)
+    volumes = data[..., mask]
+    reshape_shape = (-1, volumes.shape[-1]) if mask is None else (-1, sum(mask))
+    reshaped_data = volumes.reshape(reshape_shape)
+    p5 = np.percentile(reshaped_data, p_min, axis=0)
+    p95 = np.percentile(reshaped_data, p_max, axis=0) - p5
+    normalized[..., mask] = (volumes - p5) * p95.mean() / p95 + p5.mean()
+
+    if inplace:
+        return None
+
+    return normalized
+
+
+def grand_mean_normalization(
+    data: np.ndarray,
+    mask: np.ndarray | None = None,
+    center: float = DEFAULT_CLIP_PERCENTILE,
+    inplace: bool = False,
+) -> np.ndarray | None:
+    """Robust grand mean normalization.
+
+    Regresses out global signal differences so that data are normalized and
+    centered around a given value.
+
+    If a mask is provided, only the data within the mask are considered.
+
+    Parameters
+    ----------
+    data : :obj:`~numpy.ndarray`
+        Data to be normalized.
+    mask : :obj:`~numpy.ndarray`, optional
+        Mask. If provided, only the data within the mask are considered.
+    center : float, optional
+        Central value around which to normalize the data.
+    inplace : :obj:`bool`, optional
+        If ``False``, the normalization is performed on the original data.
+
+    Returns
+    -------
+    data : :obj:`~numpy.ndarray` or None
+        Normalized data or ``None`` if ``inplace`` is ``True``.
+    """
+
+    normalized = data if inplace else data.copy()
+
+    mask = mask if mask is not None else np.ones(data.shape[-1], dtype=bool)
+    volumes = data[..., mask]
+
+    centers = np.median(volumes, axis=(0, 1, 2))
+    reference = np.percentile(centers[centers >= 1.0], center)
+    centers[centers < 1.0] = reference
+    drift = reference / centers
+    normalized[..., mask] = volumes * drift
+
+    if inplace:
+        return None
+
+    return normalized
+
+
+def dwi_select_shells(
+    gradients: np.ndarray,
+    index: int,
+    atol_low: float | None = None,
+    atol_high: float | None = None,
+) -> np.ndarray:
+    """Select DWI shells around the given index and lower and upper b-value
+    bounds.
+
+    Computes a boolean mask of the DWI shells around the given index with the
+    provided lower and upper bound b-values.
+
+    If ``atol_low`` and ``atol_high`` are both ``None``, the returned shell mask
+    corresponds to the lengths of the diffusion-sensitizing gradients.
+
+    Parameters
+    ----------
+    gradients : :obj:`~numpy.ndarray`
+        Gradients.
+    index : :obj:`int`
+        Index of the shell data.
+    atol_low : :obj:`float`, optional
+        A lower bound for the b-value.
+    atol_high : :obj:`float`, optional
+        An upper bound for the b-value.
+
+    Returns
+    -------
+    shellmask : :obj:`~numpy.ndarray`
+        Shell mask.
+    """
+
+    bvalues = gradients[:, -1]
+    bcenter = bvalues[index]
+
+    shellmask = np.ones(len(bvalues), dtype=bool)
+    shellmask[index] = False  # Drop the held-out index
+
+    if atol_low is None and atol_high is None:
+        return shellmask
+
+    atol_low = 0 if atol_low is None else atol_low
+    atol_high = gradients[:, -1].max() if atol_high is None else atol_high
+
+    # Keep only b-values within the range defined by atol_high and atol_low
+    shellmask[bvalues > (bcenter + atol_high)] = False
+    shellmask[bvalues < (bcenter - atol_low)] = False
+
+    if not shellmask.sum():
+        raise RuntimeError(f"Shell corresponding to index {index} (b={bcenter}) is empty.")
+
+    return shellmask

src/nifreeze/model/dmri.py

@@ -27,11 +27,8 @@
 from dipy.core.gradients import gradient_table_from_bvals_bvecs
 from joblib import Parallel, delayed
 
-from nifreeze.data.dmri import (
-    DEFAULT_CLIP_PERCENTILE,
-    DTI_MIN_ORIENTATIONS,
-    DWI,
-)
+from nifreeze.data.dmri import DTI_MIN_ORIENTATIONS, DWI
+from nifreeze.data.filtering import BVAL_ATOL, dwi_select_shells, grand_mean_normalization
 from nifreeze.model.base import BaseModel, ExpectationModel
 
 S0_EPSILON = 1e-6
@@ -215,14 +212,14 @@ def fit_predict(self, index: int | None = None, **kwargs):
 class AverageDWIModel(ExpectationModel):
     """A trivial model that returns an average DWI volume."""
 
-    __slots__ = ("_th_low", "_th_high", "_detrend")
+    __slots__ = ("_atol_low", "_atol_high", "_detrend")
 
     def __init__(
         self,
         dataset: DWI,
         stat: str = "median",
-        th_low: float = 100.0,
-        th_high: float = 100.0,
+        atol_low: float = BVAL_ATOL,
+        atol_high: float = BVAL_ATOL,
         detrend: bool = False,
         **kwargs,
     ):
@@ -235,10 +232,10 @@ def __init__(
             Reference to a DWI object.
         stat : :obj:`str`, optional
             Whether the summary statistic to apply is ``"mean"`` or ``"median"``.
-        th_low : :obj:`float`, optional
+        atol_low : :obj:`float`, optional
             A lower bound for the b-value corresponding to the diffusion weighted images
             that will be averaged.
-        th_high : :obj:`float`, optional
+        atol_low : :obj:`float`, optional
             An upper bound for the b-value corresponding to the diffusion weighted images
             that will be averaged.
         detrend : :obj:`bool`, optional
@@ -249,8 +246,8 @@ def __init__(
         """
         super().__init__(dataset, stat=stat, **kwargs)
 
-        self._th_low = th_low
-        self._th_high = th_high
+        self._atol_low = atol_low
+        self._atol_high = atol_high
         self._detrend = detrend
 
     def fit_predict(self, index: int | None = None, *_, **kwargs):
@@ -259,31 +256,22 @@ def fit_predict(self, index: int | None = None, *_, **kwargs):
         if index is None:
             raise RuntimeError(f"Model {self.__class__.__name__} does not allow locking.")
 
-        bvalues = self._dataset.gradients[:, -1]
-        bcenter = bvalues[index]
-
-        shellmask = np.ones(len(self._dataset), dtype=bool)
-
-        # Keep only bvalues within the range defined by th_high and th_low
-        shellmask[index] = False
-        shellmask[bvalues > (bcenter + self._th_high)] = False
-        shellmask[bvalues < (bcenter - self._th_low)] = False
-
-        if not shellmask.sum():
-            raise RuntimeError(f"Shell corresponding to index {index} (b={bcenter}) is empty.")
+        shellmask = dwi_select_shells(
+            self._dataset.gradients,
+            index,
+            atol_low=self._atol_low,
+            atol_high=self._atol_high,
+        )
 
         shelldata = self._dataset.dataobj[..., shellmask]
 
         # Regress out global signal differences
         if self._detrend:
-            centers = np.median(shelldata, axis=(0, 1, 2))
-            reference = np.percentile(centers[centers >= 1.0], DEFAULT_CLIP_PERCENTILE)
-            centers[centers < 1.0] = reference
-            drift = reference / centers
-            shelldata = shelldata * drift
+            shelldata = grand_mean_normalization(shelldata, mask=None)
 
         # Select the summary statistic
         avg_func = np.median if self._stat == "median" else np.mean
+
         # Calculate the average
         return avg_func(shelldata, axis=-1)