Merge branch 'main' into tst/test-pet-init-load

Author: mnoergaard

.github/workflows/benchmark.yml

@@ -55,5 +55,6 @@ jobs:
         export OMP_NUM_THREADS=1
     - name: Run benchmarks
       run: |
-        asv machine --yes --config benchmarks/asv.conf.json
-        asv run --config benchmarks/asv.conf.json --show-stderr
+        CONFIG_FILE="$(pwd)/benchmarks/asv.conf.json"
+        asv machine --yes --config "$CONFIG_FILE"
+        asv run --config "$CONFIG_FILE" --show-stderr

benchmarks/benchmarks/bench_model.py

@@ -34,6 +34,7 @@
 from skimage.morphology import ball
 
 from nifreeze.model.gpr import DiffusionGPR, SphericalKriging
+from nifreeze.utils.ndimage import load_api
 
 
 class DiffusionGPRBenchmark(ABC):
@@ -67,7 +68,7 @@ def make_data(self):
         name = "sherbrooke_3shell"
 
         dwi_fname, bval_fname, bvec_fname = dpd.get_fnames(name=name)
-        dwi_data = nb.load(dwi_fname).get_fdata()
+        dwi_data = load_api(dwi_fname, nb.Nifti1Image).get_fdata()
         bvals, bvecs = read_bvals_bvecs(bval_fname, bvec_fname)
 
         _, brain_mask = median_otsu(dwi_data, vol_idx=[0])
@@ -96,7 +97,10 @@ def make_data(self):
         self._y_test = y[:, train_test_mask]
 
     def time_fit(self, *args):
+        assert self._estimator is not None
+        assert self._y_train is not None
         self._estimator = self._estimator.fit(self._X_train, self._y_train.T)
 
     def time_predict(self):
+        assert self._estimator is not None
         self._estimator.predict(self._X_test)

docs/_static/nifreeze-flowchart.svg

@@ -90,7 +90,7 @@
     "    if not dilmask_path.exists():\n",
     "        niimsk = nb.load(bmask_path)\n",
     "        niimsk.__class__(\n",
-    "            binary_dilation(niimsk.get_fdata() > 0.0, ball(4)).astype(\"uint8\"),\n",
+    "            binary_dilation(niimsk.get_fdata() > 0.0, ball(4)).astype(np.uint8),\n",
     "            niimsk.affine,\n",
     "            niimsk.header,\n",
     "        ).to_filename(dilmask_path)\n",

docs/notebooks/bold_realignment.ipynb

@@ -141,6 +141,9 @@ version-file = "src/nifreeze/_version.py"
 # Developer tool configurations
 #
 
+[tool.mypy]
+check_untyped_defs = true
+
 [[tool.mypy.overrides]]
 module = [
   "nipype.*",
@@ -233,9 +236,7 @@ branch = true
 parallel = true
 concurrency = ["multiprocessing"]
 omit = [
-  "*/tests/*",
-  "*/testing/*",
-  "*/viz/*",
+  "^test/*",
   "*/__init__.py",
   "*/conftest.py",
   "src/nifreeze/_version.py"

pyproject.toml

@@ -34,12 +34,12 @@ def _parse_yaml_config(file_path: str) -> dict:
 
     Parameters
     ----------
-    file_path : str
+    file_path : :obj:`str`
         Path to the YAML configuration file.
 
     Returns
     -------
-    dict
+    :obj:`dict`
         A dictionary containing the parsed YAML configuration.
     """
     with open(file_path, "r") as file:
@@ -201,20 +201,20 @@ def _normalize_model_name(model_name: str) -> str:
     return model_name.lower().replace("single", "")
 
 
-def parse_args(argv: list) -> tuple[Namespace, dict, dict, dict]:
+def parse_args(argv: list[str] | None = None) -> tuple[Namespace, dict, dict, dict]:
     """Parse the command line arguments and return a curated arguments.
 
     Performs further checks to ensure that necessary data is provided for the
     estimation process.
 
     Parameters
     ----------
-    argv : list
+    argv : :obj:`list`, optional
         Arguments.
 
     Returns
     -------
-    args : :obj:`Namespace`
+    args : :obj:`~argparse.Namespace`
         Populated namespace.
     extra_kwargs : :obj:`dict`
         Extra keyword arguments passed to the dataset.

src/nifreeze/cli/parser.py

@@ -29,10 +29,15 @@
 from nifreeze.estimator import Estimator
 
 
-def main(argv=None) -> None:
+def main(argv: list[str] | None = None) -> None:
     """
     Entry point.
 
+    Parameters
+    ----------
+    argv : obj:`list`, optional
+        Arguments.
+
     Returns
     -------
     None

src/nifreeze/cli/run.py

@@ -34,7 +34,6 @@
 import h5py
 import nibabel as nb
 import numpy as np
-from nibabel.spatialimages import SpatialHeader
 from nitransforms.linear import LinearTransformsMapping
 from typing_extensions import Self, TypeVarTuple, Unpack
 
@@ -79,12 +78,14 @@ class BaseDataset(Generic[Unpack[Ts]]):
     """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))
     """Best affine for RAS-to-voxel conversion of coordinates (NIfTI header)."""
-    brainmask: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
+    brainmask: np.ndarray | None = attrs.field(
+        default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp)
+    )
     """A boolean ndarray object containing a corresponding brainmask."""
-    motion_affines: np.ndarray = attrs.field(default=None, eq=attrs.cmp_using(eq=_cmp))
-    """List of :obj:`~nitransforms.linear.Affine` realigning the dataset."""
-    datahdr: SpatialHeader = attrs.field(default=None)
-    """A :obj:`~nibabel.spatialimages.SpatialHeader` header corresponding to the data."""
+    motion_affines: np.ndarray | None = attrs.field(default=None, eq=attrs.cmp_using(eq=_cmp))
+    """Array of :obj:`~nitransforms.linear.Affine` realigning the dataset."""
+    datahdr: nb.Nifti1Header | None = attrs.field(default=None)
+    """A :obj:`~nibabel.Nifti1Header` header corresponding to the data."""
 
     _filepath: Path = attrs.field(
         factory=lambda: Path(mkdtemp()) / "hmxfms_cache.h5",
@@ -95,12 +96,24 @@ class BaseDataset(Generic[Unpack[Ts]]):
 
     def __len__(self) -> int:
         """Obtain the number of volumes/frames in the dataset."""
-        if self.dataobj is None:
-            return 0
-
         return self.dataobj.shape[-1]
 
     def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[Unpack[Ts]]:
+        """
+        Extracts extra fields synchronized with the indexed access of the corresponding data object.
+
+        Parameters
+        ----------
+        idx : :obj:`int` or :obj:`slice` or :obj:`tuple` or :obj:`~numpy.ndarray`
+            Index (or indexing type/object) for which extra information will be extracted.
+
+        Returns
+        -------
+        :obj:`tuple`
+            A tuple with the extra fields (may be an empty tuple if no extra fields are defined).
+
+        """
+        _ = idx  # Avoid unused parameter warning
         return ()  # type: ignore[return-value]
 
     def __getitem__(
@@ -116,16 +129,24 @@ def __getitem__(
 
         Returns
         -------
-        volumes : :obj:`~numpy.ndarray`
+        :obj:`~numpy.ndarray`
             The selected data subset.
             If ``idx`` is a single integer, this will have shape ``(X, Y, Z)``,
             otherwise it may have shape ``(X, Y, Z, k)``.
-        motion_affine : :obj:`~numpy.ndarray` or ``None``
+        affine : :obj:`~numpy.ndarray` or ``None``
             The corresponding per-volume motion affine(s) or ``None`` if identity transform(s).
+        Unpack[:obj:`~nifreeze.data.base.Ts`]
+            Zero or more additional per-volume fields returned as unpacked
+            trailing elements. The exact number, order, and types of elements
+            are determined by the type variables :obj:`~nifreeze.data.base.Ts`
+            and by the values returned from :meth:`_getextra`. Subclasses
+            provide these values by implementing :meth:`_getextra`. If no extra
+            fields are defined, no element is returned.
+            Example usages:
+            - vols, aff, *extras = dataset[0:10]
+            - vol, aff, bvecs, bvals = dataset[0]  # when two extras are present
 
         """
-        if self.dataobj is None:
-            raise ValueError("No data available (dataobj is None).")
 
         affine = self.motion_affines[idx] if self.motion_affines is not None else None
         return self.dataobj[..., idx], affine, *self._getextra(idx)

src/nifreeze/data/base.py

@@ -65,7 +65,7 @@
 
 
 @attrs.define(slots=True)
-class DWI(BaseDataset[np.ndarray | None]):
+class DWI(BaseDataset[np.ndarray]):
     """Data representation structure for dMRI data."""
 
     bzero: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
@@ -75,13 +75,13 @@ class DWI(BaseDataset[np.ndarray | None]):
     eddy_xfms: list = attrs.field(default=None)
     """List of transforms to correct for estimated eddy current distortions."""
 
-    def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[np.ndarray | None]:
-        return (self.gradients[..., idx] if self.gradients is not None else None,)
+    def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[np.ndarray]:
+        return (self.gradients[..., idx],)
 
     # For the sake of the docstring
     def __getitem__(
         self, idx: int | slice | tuple | np.ndarray
-    ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None]:
+    ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]:
         """
         Returns volume(s) and corresponding affine(s) and gradient(s) through fancy indexing.
 
@@ -210,7 +210,7 @@ def to_nifti(
 
         Parameters
         ----------
-        filename : :obj:`os.pathlike`
+        filename : :obj:`os.pathlike`, optional
             The output NIfTI file path.
         write_hmxfms : :obj:`bool`, optional
             If ``True``, the head motion affines will be written out to filesystem
@@ -310,7 +310,7 @@ def from_nii(
     brainmask_file : :obj:`os.pathlike`, optional
         A brainmask NIfTI file. If provided, will be loaded and
         stored in the returned dataset.
-    motion_file : :obj:`os.pathlike`
+    motion_file : :obj:`os.pathlike`, optional
         A file containing head motion affine matrices (linear)
     gradients_file : :obj:`os.pathlike`, optional
         A text file containing the gradients table, shape (4, N) or (N, 4).

src/nifreeze/data/dmri.py

@@ -42,7 +42,7 @@
 
 
 @attrs.define(slots=True)
-class PET(BaseDataset[np.ndarray | None]):
+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))
@@ -52,13 +52,13 @@ class PET(BaseDataset[np.ndarray | None]):
     uptake: np.ndarray = attrs.field(default=None, repr=_data_repr, eq=attrs.cmp_using(eq=_cmp))
     """A (N,) numpy array specifying the uptake value of each sample or frame."""
 
-    def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[np.ndarray | None]:
-        return (self.midframe[idx] if self.midframe is not None else None,)
+    def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[np.ndarray]:
+        return (self.midframe[idx],)
 
     # For the sake of the docstring
     def __getitem__(
         self, idx: int | slice | tuple | np.ndarray
-    ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None]:
+    ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]:
         """
         Returns volume(s) and corresponding affine(s) and timing(s) through fancy indexing.
 
@@ -75,7 +75,7 @@ def __getitem__(
             otherwise it may have shape ``(X, Y, Z, k)``.
         motion_affine : :obj:`~numpy.ndarray` or ``None``
             The corresponding per-volume motion affine(s) or ``None`` if identity transform(s).
-        time : :obj:`float` or ``None``
+        time : :obj:`~numpy.ndarray`
             The frame time corresponding to the index(es).
 
         """
@@ -105,18 +105,17 @@ def lofo_split(self, index):
         with h5py.File(self._filepath, "r") as in_file:
             root = in_file["/0"]
             pet_frame = np.asanyarray(root["dataobj"][..., index])
-            if self.midframe is not None:
-                timing_frame = np.asanyarray(root["midframe"][..., index])
+            timing_frame = np.asanyarray(root["midframe"][..., index])
 
         # Mask to exclude the selected frame
         mask = np.ones(self.dataobj.shape[-1], dtype=bool)
         mask[index] = False
 
         train_data = self.dataobj[..., mask]
-        train_timings = self.midframe[mask] if self.midframe is not None else None
+        train_timings = self.midframe[mask]
 
         test_data = pet_frame
-        test_timing = timing_frame if self.midframe is not None else None
+        test_timing = timing_frame
 
         return (train_data, train_timings), (test_data, test_timing)
 
@@ -144,7 +143,7 @@ def set_transform(self, index: int, affine: np.ndarray, order: int = 3) -> None:
 
         # update transform
         if self.motion_affines is None:
-            self.motion_affines = [None] * len(self)
+            self.motion_affines = np.asarray([None] * len(self))
 
         self.motion_affines[index] = xform
 
@@ -174,12 +173,6 @@ def to_filename(
                         compression_opts=compression_opts,
                     )
 
-    def to_nifti(self, filename, *_):
-        """Write a NIfTI 1.0 file to disk."""
-        nii = nb.Nifti1Image(self.dataobj, self.affine, None)
-        nii.header.set_xyzt_units("mm")
-        nii.to_filename(filename)
-
     @classmethod
     def from_filename(cls, filename: Path | str) -> Self:
         """Read an HDF5 file from disk."""
@@ -225,9 +218,9 @@ def load(
 
 def from_nii(
     filename: Path | str,
+    frame_time: np.ndarray | list[float],
     brainmask_file: Path | str | None = None,
     motion_file: Path | str | None = None,
-    frame_time: np.ndarray | list[float] | None = None,
     frame_duration: np.ndarray | list[float] | None = None,
 ) -> PET:
     """
@@ -237,14 +230,13 @@ def from_nii(
     ----------
     filename : :obj:`os.pathlike`
         The NIfTI file.
+    frame_time : :obj:`numpy.ndarray` or :obj:`list` of :obj:`float`
+        The start times of each frame relative to the beginning of the acquisition.
     brainmask_file : :obj:`os.pathlike`, optional
         A brainmask NIfTI file. If provided, will be loaded and
         stored in the returned dataset.
     motion_file : :obj:`os.pathlike`, optional
         A file containing head motion affine matrices (linear).
-    frame_time : :obj:`numpy.ndarray` or :obj:`list` of :obj:`float`, optional
-        The start times of each frame relative to the beginning of the acquisition.
-        If ``None``, an error is raised (since BIDS requires ``FrameTimesStart``).
     frame_duration : :obj:`numpy.ndarray` or :obj:`list` of :obj:`float`, optional
         The duration of each frame.
         If ``None``, it is derived by the difference of consecutive frame times,
@@ -261,8 +253,6 @@ def from_nii(
         If ``frame_time`` is not provided (BIDS requires it).
 
     """
-    if frame_time is None:
-        raise RuntimeError("frame_time must be provided")
     if motion_file:
         raise NotImplementedError
 
@@ -277,17 +267,14 @@ def from_nii(
 
     pet_obj.uptake = _compute_uptake_statistic(data, stat_func=np.sum)
 
-    # If the user supplied new values, set them
-    if frame_time is not None:
-        # Convert to a float32 numpy array and zero out the earliest time
-        frame_time_arr = np.array(frame_time, dtype=np.float32)
-        frame_time_arr -= frame_time_arr[0]
-        pet_obj.midframe = frame_time_arr
+    # Convert to a float32 numpy array and zero out the earliest time
+    frame_time_arr = np.array(frame_time, dtype=np.float32)
+    frame_time_arr -= frame_time_arr[0]
+    pet_obj.midframe = frame_time_arr
 
     # If the user doesn't provide frame_duration, we derive it:
     if frame_duration is None:
-        if pet_obj.midframe is not None:
-            durations = _compute_frame_duration(pet_obj.midframe)
+        durations = _compute_frame_duration(pet_obj.midframe)
     else:
         durations = np.array(frame_duration, dtype=np.float32)