enh: add tests to docstring and integrate to overhaul

Closes: #332.

Author: oesteban

src/nifreeze/data/dmri.py

@@ -37,164 +37,21 @@
 from typing_extensions import Self
 
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
-from nifreeze.utils.ndimage import get_data, load_api
-
-DEFAULT_CLIP_PERCENTILE = 75
-"""Upper percentile threshold for intensity clipping."""
-
-DEFAULT_MIN_S0 = 1e-5
-"""Minimum value when considering the :math:`S_{0}` DWI signal."""
-
-DEFAULT_MAX_S0 = 1.0
-"""Maximum value when considering the :math:`S_{0}` DWI signal."""
-
-DEFAULT_LOWB_THRESHOLD = 50
-"""The lower bound for the b-value so that the orientation is considered a DW volume."""
-
-DEFAULT_GRADIENT_EPS = 1e-8
-"""Epsilon value for b-vector normalization."""
-
-DEFAULT_HIGHB_THRESHOLD = 8000
-"""A b-value cap for DWI data."""
-
-DEFAULT_NUM_BINS = 15
-"""Number of bins to classify b-values."""
-
-DEFAULT_MULTISHELL_BIN_COUNT_THR = 7
-"""Default bin count to consider a multishell scheme."""
-
-DTI_MIN_ORIENTATIONS = 6
-"""Minimum number of nonzero b-values in a DWI dataset."""
-
-GRADIENT_ABSENCE_ERROR_MSG = "Gradient table may not be None."
-"""Gradient absence error message."""
-
-GRADIENT_OBJECT_ERROR_MSG = "Gradient table must be a numeric homogeneous array-like object"
-"""Gradient object error message."""
-
-GRADIENT_VOLUME_DIMENSIONALITY_MISMATCH_ERROR = """\
-Gradient table shape does not match the number of diffusion volumes: \
-expected {n_volumes} rows, found {n_gradients}."""
-"""dMRI volume count vs. gradient count mismatch error message."""
-
-GRADIENT_BVAL_BVEC_PRIORITY_WARN_MSG = """\
-Both a gradients table file and b-vec/val files are defined; \
-ignoring b-vec/val files in favor of the gradients_file."""
-""""dMRI gradient file priority warning message."""
-
-GRADIENT_NDIM_ERROR_MSG = "Gradient table must be a 2D array"
-"""dMRI gradient dimensionality error message."""
-
-GRADIENT_DATA_MISSING_ERROR = "No gradient data provided."
-"""dMRI missing gradient data error message."""
-
-GRADIENT_EXPECTED_COLUMNS_ERROR_MSG = (
-    "Gradient table must have four columns (3 direction components and one b-value)."
+from nifreeze.data.dmriutils import (
+    DEFAULT_HIGHB_THRESHOLD,
+    DEFAULT_LOWB_THRESHOLD,
+    DEFAULT_MULTISHELL_BIN_COUNT_THR,
+    DEFAULT_NUM_BINS,
+    DTI_MIN_ORIENTATIONS,
+    GRADIENT_BVAL_BVEC_PRIORITY_WARN_MSG,
+    GRADIENT_DATA_MISSING_ERROR,
+    GRADIENT_EXPECTED_COLUMNS_ERROR_MSG,
+    GRADIENT_VOLUME_DIMENSIONALITY_MISMATCH_ERROR,
+    find_shelling_scheme,
+    format_gradients,
+    transform_fsl_bvec,
 )
-"""dMRI gradient expected columns error message."""
-
-
-def format_gradients(
-    value: npt.ArrayLike | None,
-) -> np.ndarray | None:
-    """
-    Validate and orient gradient tables to row-major convention.
-
-    Parameters
-    ----------
-    value : :obj:`ArrayLike`
-        The value to format.
-
-    Returns
-    -------
-    :obj:`~numpy.ndarray`
-        Row-major convention gradient table.
-
-    Raises
-    ------
-    exc:`ValueError`
-        If ``value`` is not a 2D :obj:`~numpy.ndarray` (``value.ndim != 2``).
-
-    Examples
-    --------
-    Passing an already well-formed table returns the data unchanged::
-
-        >>> format_gradients(
-        ...     [
-        ...         [1, 0, 0, 0],
-        ...         [0, 1, 0, 1000],
-        ...         [0, 0, 1, 2000],
-        ...         [0, 0, 0, 0],
-        ...         [0, 0, 0, 1000],
-        ...     ]
-        ... )
-        array([[   1,    0,    0,    0],
-               [   0,    1,    0, 1000],
-               [   0,    0,    1, 2000],
-               [   0,    0,    0,    0],
-               [   0,    0,    0, 1000]])
-
-    Column-major inputs are automatically transposed when an expected
-    number of diffusion volumes is provided::
-
-        >>> format_gradients(
-        ...     [[1, 0], [0, 1], [0, 0], [1000, 2000]],
-        ... )
-        array([[   1,    0,    0, 1000],
-               [   0,    1,    0, 2000]])
-
-    Gradient tables must always have two dimensions::
-
-        >>> format_gradients([0, 1, 0, 1000])
-        Traceback (most recent call last):
-        ...
-        ValueError: Gradient table must be a 2D array
-
-    Gradient tables must have a regular shape::
-
-        >>> format_gradients([[1, 2], [3, 4, 5]])
-        Traceback (most recent call last):
-        ...
-        TypeError: Gradient table must be a numeric homogeneous array-like object
-
-    Gradient tables must always have two dimensions::
-
-        >>> format_gradients([0, 1, 0, 1000])
-        Traceback (most recent call last):
-        ...
-        ValueError: Gradient table must be a 2D array
-
-    """
-
-    if value is None:
-        raise ValueError(GRADIENT_ABSENCE_ERROR_MSG)
-
-    try:
-        formatted = np.asarray(value, dtype=float)
-    except (TypeError, ValueError) as exc:
-        # Conversion failed (e.g. nested ragged objects, non-numeric)
-        raise TypeError(GRADIENT_OBJECT_ERROR_MSG) from exc
-
-    if formatted.ndim != 2:
-        raise ValueError(GRADIENT_NDIM_ERROR_MSG)
-
-    # If the numeric values are all integers, preserve integer dtype
-    if np.all(np.isfinite(formatted)) and np.allclose(formatted, np.round(formatted)):
-        formatted = formatted.astype(int)
-
-    # Transpose if column-major
-    formatted = formatted.T if formatted.shape[0] == 4 and formatted.shape[1] != 4 else formatted
-
-    # Normalize b-vectors in-place
-    bvecs = formatted[:, :3]
-    norms = np.linalg.norm(bvecs, axis=1)
-    mask = norms > DEFAULT_GRADIENT_EPS
-    if np.any(mask):
-        formatted[mask, :3] = bvecs[mask] / norms[mask, None]  # Norm b-vectors
-        formatted[mask, 3] *= norms[mask]  # Scale b-values by norm
-        formatted[~mask, :] = 0.0  # Zero-out small b-vectors
-
-    return formatted
+from nifreeze.utils.ndimage import get_data, load_api
 
 
 def validate_gradients(
@@ -216,10 +73,33 @@ def validate_gradients(
         The attribute being validated; attr.name is used in the error message.
     value : :obj:`~npt.NDArray`
         The value to validate.
+
+    Raises
+    ------
+    :exc:`ValueError`
+        If the gradient table is invalid.
+
+
+    Examples
+    --------
+    Non-finite inputs are rejected::
+
+        >>> validate_gradients(None, None, [[np.inf, 0.0, 0.0, 1000]])
+        Traceback (most recent call last):
+        ...
+        ValueError: Gradient table contains NaN or infinite values.
+        >>> validate_gradients(None, None, [[np.nan, 0.0, 0.0, 1000]])
+        Traceback (most recent call last):
+        ...
+        ValueError: Gradient table contains NaN or infinite values.
+
     """
-    if value.shape[1] != 4:
+    if np.shape(value)[1] != 4:
         raise ValueError(GRADIENT_EXPECTED_COLUMNS_ERROR_MSG)
 
+    if not np.all(np.isfinite(value)):
+        raise ValueError("Gradient table contains NaN or infinite values.")
+
 
 @attrs.define(slots=True)
 class DWI(BaseDataset[np.ndarray]):
@@ -567,180 +447,3 @@ def from_nii(
         bzero=b0_data,
         brainmask=brainmask_data,
     )
-
-
-def find_shelling_scheme(
-    bvals: np.ndarray,
-    num_bins: int = DEFAULT_NUM_BINS,
-    multishell_nonempty_bin_count_thr: int = DEFAULT_MULTISHELL_BIN_COUNT_THR,
-    bval_cap: float = DEFAULT_HIGHB_THRESHOLD,
-) -> tuple[str, list[npt.NDArray[np.floating]], list[np.floating]]:
-    """
-    Find the shelling scheme on the given b-values.
-
-    Computes the histogram of the b-values according to ``num_bins``
-    and depending on the nonempty bin count, classify the shelling scheme
-    as single-shell if they are 2 (low-b and a shell); multi-shell if they are
-    below the ``multishell_nonempty_bin_count_thr`` value; and DSI otherwise.
-
-    Parameters
-    ----------
-    bvals : :obj:`list` or :obj:`~numpy.ndarray`
-         List or array of b-values.
-    num_bins : :obj:`int`, optional
-        Number of bins.
-    multishell_nonempty_bin_count_thr : :obj:`int`, optional
-        Bin count to consider a multi-shell scheme.
-    bval_cap : :obj:`float`, optional
-        Maximum b-value to be considered in a multi-shell scheme.
-
-    Returns
-    -------
-    scheme : :obj:`str`
-        Shelling scheme.
-    bval_groups : :obj:`list`
-        List of grouped b-values.
-    bval_estimated : :obj:`list`
-        List of 'estimated' b-values as the median value of each b-value group.
-
-    """
-
-    # Bin the b-values: use -1 as the lower bound to be able to appropriately
-    # include b0 values
-    hist, bin_edges = np.histogram(bvals, bins=num_bins, range=(-1, min(max(bvals), bval_cap)))
-
-    # Collect values in each bin
-    bval_groups = []
-    bval_estimated = []
-    for lower, upper in zip(bin_edges[:-1], bin_edges[1:], strict=False):
-        # Add only if a nonempty b-values mask
-        if (mask := (bvals > lower) & (bvals <= upper)).sum():
-            bval_groups.append(bvals[mask])
-            bval_estimated.append(np.median(bvals[mask]))
-
-    nonempty_bins = len(bval_groups)
-
-    if nonempty_bins < 2:
-        raise ValueError("DWI must have at least one high-b shell")
-
-    if nonempty_bins == 2:
-        scheme = "single-shell"
-    elif nonempty_bins < multishell_nonempty_bin_count_thr:
-        scheme = "multi-shell"
-    else:
-        scheme = "DSI"
-
-    return scheme, bval_groups, bval_estimated
-
-
-def transform_fsl_bvec(
-    b_ijk: np.ndarray, xfm: np.ndarray, imaffine: np.ndarray, invert: bool = False
-) -> np.ndarray:
-    """
-    Transform a b-vector from the original space to the new space defined by the affine.
-
-    Parameters
-    ----------
-    b_ijk : :obj:`~numpy.ndarray`
-        The b-vector in FSL/DIPY conventions (i.e., voxel coordinates).
-    xfm : :obj:`~numpy.ndarray`
-        The affine transformation to apply.
-        Please note that this is the inverse of the head-motion-correction affine,
-        which maps coordinates from the realigned space to the moved (scan) space.
-        In this case, we want to move the b-vector from the moved (scan) space into
-        the realigned space.
-    imaffine : :obj:`~numpy.ndarray`
-        The image's affine, to convert.
-    invert : :obj:`bool`, optional
-        If ``True``, the transformation will be inverted.
-
-    Returns
-    -------
-    :obj:`~numpy.ndarray`
-        The transformed b-vector in voxel coordinates (FSL/DIPY).
-
-    """
-    xfm = np.linalg.inv(xfm) if invert else xfm.copy()
-
-    # Go from world coordinates (xfm) to voxel coordinates
-    ijk2ijk_xfm = np.linalg.inv(imaffine) @ xfm @ imaffine
-
-    return ijk2ijk_xfm[:3, :3] @ b_ijk[:3]
-
-
-def normalize_gradients(value: np.ndarray, eps: float = 1e-8, copy: bool = True) -> np.ndarray:
-    """Normalize b-vectors in arrays of common shapes.
-
-    Parameters
-    ----------
-    value : :obj:`~numpy.ndarray`
-        Input array with shape one of:
-        - (N, 3) : rows are b-vector components (e.g., [gx gy gz])
-        - (N, 4) : first 3 columns are b-vector components (e.g., [gx gy gz b])
-        - (3, N) : columns are b-vector components (e.g., [gx gy gz].T)
-        - (4, N) : first 3 rows are b-vector components (e.g., [gx gy gz b].T)
-        - (3,) or (1,3) or (3,1) : single b-vector
-        Columns are checked first to disambiguate Nx3/Nx4 cases.
-    eps : float, optional
-        Threshold below which a vector is considered zero and left unchanged.
-    copy : bool, optional
-        If ``True``, returns a new array; modify in-place otherwise.
-
-    Returns
-    -------
-    out : :obj:`~numpy.ndarray`
-        Array with the same shape as ``value`` with each 3-component b-vector
-        normalized.
-    """
-    arr = np.asarray(value, dtype=float)
-
-    # 1D single vector
-    if arr.ndim == 1:
-        if arr.size != 3:
-            raise ValueError(GRADIENT_NORMALIZATION_LENGTH_ERROR_MSG)
-        norm = np.linalg.norm(arr)
-        if norm > eps:
-            if copy:
-                return arr / norm
-            else:
-                # Perform in-place normalization on the array view
-                arr[:] = arr / norm
-                return arr
-        else:
-            return arr.copy() if copy else arr
-
-    if arr.ndim != 2:
-        raise ValueError(GRADIENT_NORMALIZATION_SHAPE_ERROR_MSG)
-
-    rows, cols = arr.shape
-
-    # Prepare output (copy or in-place)
-    normalized_arr = arr.copy() if copy else arr
-
-    # Determine where the 3-component vectors live and create a (N, 3) view
-    # Check columns first to make Nx3/Nx4 deterministic
-    if cols == 4:
-        # Nx4: first 3 columns are b-vectors components, last are b-values
-        vecs = normalized_arr[:, :3]  # shape (N, 3)
-    elif cols == 3:
-        # Nx3: rows are vectors
-        vecs = normalized_arr  # shape (N, 3)
-    elif rows == 4:
-        # 4xN: first 3 rows are b-vector components, last row are b-values
-        # Create a (N, 3) view by transposing first 3 rows
-        vecs = normalized_arr[:3, :].T  # shape (N, 3)
-    elif rows == 3:
-        # 3xN: columns are vectors: normalize per-column
-        vecs = normalized_arr.T  # shape (N, 3)
-    else:
-        raise ValueError(
-            GRADIENT_NORMALIZATION_UNRECOGNIZED_SHAPE_ERROR_MSG.format(shape=arr.shape)
-        )
-
-    # Normalize in-place on vecs (which is a view into output)
-    norms = np.linalg.norm(vecs, axis=1)
-    mask = norms > eps
-    if np.any(mask):
-        vecs[mask] = vecs[mask] / norms[mask, None]
-
-    return normalized_arr