ENH: Validate data objects' attributes at instantiation

jhlegarreta · jhlegarreta · commit 702e7f29a33f · 2025-11-17T20:09:49.000-05:00
Validate data objects' attributes at instantiation: ensures that the
attributes are present and match the expected dimensionalities.
diff --git a/src/nifreeze/data/base.py b/src/nifreeze/data/base.py
@@ -44,6 +44,21 @@
 
 ImageGrid = namedtuple("ImageGrid", ("shape", "affine"))
 
+DATAOBJ_ABSENCE_ERROR_MSG = "BaseDataset 'dataobj' may not be None"
+"""BaseDataset initialization dataobj absence error message."""
+
+DATAOBJ_NDIM_ERROR_MSG = "BaseDataset 'dataobj' must be a 4-D numpy array"
+"""BaseDataset initialization dataobj dimensionality error message."""
+
+AFFINE_ABSENCE_ERROR_MSG = "BaseDataset 'affine' may not be None"
+"""BaseDataset initialization affine absence error message."""
+
+AFFINE_SHAPE_ERROR_MSG = "BaseDataset 'affine' must be a 2-D numpy array (4 x 4)"
+"""BaseDataset initialization affine shape error message."""
+
+BRAINMASK_SHAPE_MISMATCH_ERROR_MSG = "BaseDataset brainmask shape ({brainmask_shape}) does not match dataset volumes ({data_shape})."
+"""BaseDataset brainmask shape mismatch error message."""
+
 
 def _data_repr(value: np.ndarray | None) -> str:
     if value is None:
@@ -58,6 +73,20 @@ def _cmp(lh: Any, rh: Any) -> bool:
     return lh == rh
 
 
+def _dataobj_validator(inst, attr, value) -> None:
+    if value is None:
+        raise ValueError(DATAOBJ_ABSENCE_ERROR_MSG)
+    if not isinstance(value, np.ndarray) or value.ndim != 4:
+        raise ValueError(DATAOBJ_NDIM_ERROR_MSG)
+
+
+def _affine_validator(inst, attr, value) -> None:
+    if value is None:
+        raise ValueError(AFFINE_ABSENCE_ERROR_MSG)
+    if not isinstance(value, np.ndarray) or value.shape != (4, 4):
+        raise ValueError(AFFINE_SHAPE_ERROR_MSG)
+
+
 @attrs.define(slots=True)
 class BaseDataset(Generic[Unpack[Ts]]):
     """
@@ -75,9 +104,13 @@ class BaseDataset(Generic[Unpack[Ts]]):
 
     """
 
-    dataobj: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
+    dataobj: np.ndarray = attrs.field(
+        default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp), validator=_dataobj_validator
+    )
     """A :obj:`~numpy.ndarray` object for the data array."""
-    affine: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
+    affine: np.ndarray = attrs.field(
+        default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp), validator=_affine_validator
+    )
     """Best affine for RAS-to-voxel conversion of coordinates (NIfTI header)."""
     brainmask: np.ndarray | None = attrs.field(
         default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp)
@@ -95,6 +128,22 @@ class BaseDataset(Generic[Unpack[Ts]]):
     )
     """A path to an HDF5 file to store the whole dataset."""
 
+    def __attrs_post_init__(self) -> None:
+        """Enforce basic consistency of base dataset fields at instantiation
+        time.
+
+        Specifically, the brainmask (if present) must match spatial shape of
+        dataobj.
+        """
+
+        if self.brainmask is not None:
+            if self.brainmask.shape != tuple(self.dataobj.shape[:3]):
+                raise ValueError(
+                    BRAINMASK_SHAPE_MISMATCH_ERROR_MSG.format(
+                        brainmask_shape=self.brainmask.shape, data_shape=self.dataobj.shape[:3]
+                    )
+                )
+
     def __len__(self) -> int:
         """Obtain the number of volumes/frames in the dataset."""
         return self.dataobj.shape[-1]
diff --git a/src/nifreeze/data/dmri.py b/src/nifreeze/data/dmri.py
@@ -39,6 +39,17 @@
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
 from nifreeze.utils.ndimage import get_data, load_api
 
+GRADIENT_ABSENCE_ERROR_MSG = "DWI 'gradients' may not be None"
+"""DWI initialization gradient absence error message."""
+
+GRADIENT_SHAPE_ERROR_MSG = "DWI 'gradients' must be a 2-D numpy array (N x 4)"
+"""DWI initialization gradient shape error message."""
+
+GRADIENT_COUNT_MISMATCH_ERROR_MSG = (
+    "DWI gradients count ({n_gradients}) does not match dataset volumes ({data_vols})."
+)
+"""DWI initialization gradient count mismatch error message."""
+
 DEFAULT_CLIP_PERCENTILE = 75
 """Upper percentile threshold for intensity clipping."""
 
@@ -64,6 +75,13 @@
 """Minimum number of nonzero b-values in a DWI dataset."""
 
 
+def _gradients_validator(inst, attr, value) -> None:
+    if value is None:
+        raise ValueError(GRADIENT_ABSENCE_ERROR_MSG)
+    if not isinstance(value, np.ndarray) or value.shape[1] != 4:
+        raise ValueError(GRADIENT_SHAPE_ERROR_MSG)
+
+
 @attrs.define(slots=True)
 class DWI(BaseDataset[np.ndarray]):
     """Data representation structure for dMRI data."""
@@ -72,14 +90,43 @@ class DWI(BaseDataset[np.ndarray]):
         default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp)
     )
     """A *b=0* reference map, preferably obtained by some smart averaging."""
-    gradients: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
+    gradients: np.ndarray = attrs.field(
+        default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp), validator=_gradients_validator
+    )
     """A 2D numpy array of the gradient table (``N`` orientations x ``C`` components)."""
     eddy_xfms: list = attrs.field(default=None)
     """List of transforms to correct for estimated eddy current distortions."""
 
     def __attrs_post_init__(self) -> None:
+        """Enforce presence and basic consistency of required dMRI fields at
+        instantiation time.
+
+        Specifically, the number of gradient directions must match the last
+        dimension of the data (number of volumes).
+        """
+
         self._normalize_gradients()
 
+        # If the data object exists and has a time/volume axis, ensure sizes
+        # match.
+        data_vols = None
+        if getattr(self, "dataobj", None) is not None:
+            shape = getattr(self.dataobj, "shape", None)
+            if isinstance(shape, (tuple, list)) and len(shape) >= 1:
+                try:
+                    data_vols = int(shape[-1])
+                except (TypeError, ValueError):
+                    data_vols = None
+
+        if data_vols is not None:
+            n_gradients = self.gradients.shape[1]
+            if n_gradients != data_vols:
+                raise ValueError(
+                    GRADIENT_COUNT_MISMATCH_ERROR_MSG.format(
+                        n_gradients=n_gradients, data_vols=data_vols
+                    )
+                )
+
     def _normalize_gradients(self) -> None:
         if self.gradients is None:
             return
diff --git a/src/nifreeze/data/pet.py b/src/nifreeze/data/pet.py
@@ -41,18 +41,70 @@
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
 from nifreeze.utils.ndimage import load_api
 
+ARRAY_ATTRIBUTE_SHAPE_ERROR_MSG = "PET {attribute} must be a 1-D numpy array."
+"""PET array attribute shape error message."""
+
+SCALAR_ATTRIBUTE_ERROR_MSG = "PET {attribute} must be a scalar."
+"""PET scalar attribute shape error message."""
+
+ATTRIBUTE_SHAPE_MISMATCH_ERROR_MSG = (
+    "PET {attribute} length ({attr_len}) does not match number of frames ({data_frames})"
+)
+"""PET attribute shape mismatch error message."""
+
+
+def _1d_array_validator(inst, attr, value) -> None:
+    if not isinstance(value, np.ndarray) or value.ndim != 1:
+        raise ValueError(ARRAY_ATTRIBUTE_SHAPE_ERROR_MSG.format(attribute=attr.name))
+
+
+def _scalar_validator(inst, attr, value) -> None:
+    if not isinstance(value, (int, float, np.integer, np.floating)):
+        raise ValueError(SCALAR_ATTRIBUTE_ERROR_MSG.format(attribute=attr.name))
+
 
 @attrs.define(slots=True)
 class PET(BaseDataset[np.ndarray]):
     """Data representation structure for PET data."""
 
-    midframe: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
+    midframe: np.ndarray = attrs.field(
+        default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp), validator=_1d_array_validator
+    )
     """A (N,) numpy array specifying the midpoint timing of each sample or frame."""
-    total_duration: float = attrs.field(default=None, repr=True)
+    total_duration: float = attrs.field(default=None, repr=True, validator=_scalar_validator)
     """A float representing the total duration of the dataset."""
-    uptake: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
+    uptake: np.ndarray = attrs.field(
+        default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp), validator=_1d_array_validator
+    )
     """A (N,) numpy array specifying the uptake value of each sample or frame."""
 
+    def __attrs_post_init__(self) -> None:
+        """Enforce presence and basic consistency of required PET fields at
+        instantiation time.
+
+        Specifically, the length of the midframe and uptake attributes must
+        match the last  dimension of the data (number of frames).
+        """
+        data_frames = int(self.dataobj.shape[-1])
+
+        if len(self.midframe) != data_frames:
+            raise ValueError(
+                ATTRIBUTE_SHAPE_MISMATCH_ERROR_MSG.format(
+                    attribute=attrs.fields_dict(self.__class__)["midframe"].name,
+                    attr_len=len(self.midframe),
+                    data_frames=data_frames,
+                )
+            )
+
+        if len(self.uptake) != data_frames:
+            raise ValueError(
+                ATTRIBUTE_SHAPE_MISMATCH_ERROR_MSG.format(
+                    attribute=attrs.fields_dict(self.__class__)["uptake"].name,
+                    attr_len=len(self.uptake),
+                    data_frames=data_frames,
+                )
+            )
+
     def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[np.ndarray]:
         return (self.midframe[idx],)
 
diff --git a/test/test_data_base.py b/test/test_data_base.py
@@ -22,6 +22,7 @@
 #
 """Test dataset base class."""
 
+import re
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Any
@@ -31,6 +32,13 @@
 import pytest
 
 from nifreeze.data import NFDH5_EXT, BaseDataset, load
+from nifreeze.data.base import (
+    AFFINE_ABSENCE_ERROR_MSG,
+    AFFINE_SHAPE_ERROR_MSG,
+    BRAINMASK_SHAPE_MISMATCH_ERROR_MSG,
+    DATAOBJ_ABSENCE_ERROR_MSG,
+    DATAOBJ_NDIM_ERROR_MSG,
+)
 from nifreeze.utils.ndimage import get_data
 
 DEFAULT_RANDOM_DATASET_SHAPE = (32, 32, 32, 5)
@@ -51,6 +59,51 @@ def random_dataset(setup_random_uniform_spatial_data) -> BaseDataset:
     return BaseDataset(dataobj=data, affine=affine)
 
 
+def test_missing_dataobj_error():
+    with pytest.raises(ValueError, match=DATAOBJ_ABSENCE_ERROR_MSG):
+        BaseDataset()
+
+
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4, 6), 0.0, 1.0)
+def test_dataobj_ndim_error(setup_random_uniform_spatial_data):
+    data, _ = setup_random_uniform_spatial_data
+    with pytest.raises(ValueError, match=DATAOBJ_NDIM_ERROR_MSG):
+        BaseDataset(dataobj=data)
+
+
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4), 0.0, 1.0)
+def test_missing_affine_error(setup_random_uniform_spatial_data):
+    data, _ = setup_random_uniform_spatial_data
+    with pytest.raises(ValueError, match=AFFINE_ABSENCE_ERROR_MSG):
+        BaseDataset(dataobj=data)
+
+
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4), 0.0, 1.0)
+@pytest.mark.parametrize("size", ((2, 2), (3, 4), (4, 3), (5, 5)))
+def test_affine_shape_error(setup_random_uniform_ndim_data, size):
+    data = setup_random_uniform_ndim_data
+    affine = np.ones(size)
+    with pytest.raises(ValueError, match=re.escape(AFFINE_SHAPE_ERROR_MSG)):
+        BaseDataset(dataobj=data, affine=affine)
+
+
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4), 0.0, 1.0)
+def test_brainmask_volume_mismatch_error(request, setup_random_uniform_spatial_data):
+    data, affine = setup_random_uniform_spatial_data
+    data_shape = data.shape[:3]
+    brainmask_size = tuple(map(lambda x: x + 1, data_shape))
+    brainmask = request.node.rng.choice([True, False], size=brainmask_size)
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            BRAINMASK_SHAPE_MISMATCH_ERROR_MSG.format(
+                brainmask_shape=brainmask.shape, data_shape=data_shape
+            )
+        ),
+    ):
+        BaseDataset(dataobj=data, affine=affine, brainmask=brainmask)
+
+
 def test_base_dataset_init(random_dataset: BaseDataset):
     """Test that the BaseDataset can be initialized with random data."""
     assert random_dataset.dataobj is not None
diff --git a/test/test_data_dmri.py b/test/test_data_dmri.py
@@ -22,14 +22,23 @@
 #
 """Unit tests exercising the dMRI data structure."""
 
+import re
 from pathlib import Path
 
 import nibabel as nb
 import numpy as np
 import pytest
 
 from nifreeze.data import load
-from nifreeze.data.dmri import DWI, find_shelling_scheme, from_nii, transform_fsl_bvec
+from nifreeze.data.dmri import (
+    DWI,
+    GRADIENT_ABSENCE_ERROR_MSG,
+    GRADIENT_COUNT_MISMATCH_ERROR_MSG,
+    GRADIENT_SHAPE_ERROR_MSG,
+    find_shelling_scheme,
+    from_nii,
+    transform_fsl_bvec,
+)
 from nifreeze.utils.ndimage import load_api
 
 
@@ -82,6 +91,36 @@ def test_motion_file_not_implemented():
         from_nii("dmri.nii.gz", motion_file="motion.x5")
 
 
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4), 0.0, 1.0)
+def test_missing_gradients_error(setup_random_uniform_spatial_data):
+    data, affine = setup_random_uniform_spatial_data
+    with pytest.raises(ValueError, match=GRADIENT_ABSENCE_ERROR_MSG):
+        DWI(dataobj=data, affine=affine)
+
+
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4), 0.0, 1.0)
+def test_gradients_shape_error(setup_random_uniform_spatial_data):
+    data, affine = setup_random_uniform_spatial_data
+    gradients = np.zeros((3, data.shape[-1]))
+    with pytest.raises(ValueError, match=re.escape(GRADIENT_SHAPE_ERROR_MSG)):
+        DWI(dataobj=data, affine=affine, gradients=gradients)
+
+
+@pytest.mark.random_uniform_spatial_data((2, 2, 2, 4), 0.0, 1.0)
+def test_gradients_volume_mismatch_error(setup_random_uniform_spatial_data):
+    data, affine = setup_random_uniform_spatial_data
+    data_vols = data.shape[-1]
+    n_gradients = data_vols + 1
+    gradients = np.zeros((4, n_gradients))
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            GRADIENT_COUNT_MISMATCH_ERROR_MSG.format(n_gradients=n_gradients, data_vols=data_vols)
+        ),
+    ):
+        DWI(dataobj=data, affine=affine, gradients=gradients)
+
+
 @pytest.mark.parametrize("insert_b0", (False, True))
 @pytest.mark.parametrize("rotate_bvecs", (False, True))
 def test_load(datadir, tmp_path, insert_b0, rotate_bvecs):  # noqa: C901
diff --git a/test/test_data_pet.py b/test/test_data_pet.py
