Merge pull request #189 from jhlegarreta/sty/add-pet-type-hints

STY: Add type hints to PET-related classes and functions

Author: jhlegarreta

src/nifreeze/data/pet.py

@@ -27,14 +27,15 @@
 import json
 from collections import namedtuple
 from pathlib import Path
-from typing import Callable
+from typing import Any, Callable
 
 import attrs
 import h5py
 import nibabel as nb
 import numpy as np
 from nibabel.spatialimages import SpatialImage
 from nitransforms.linear import Affine
+from typing_extensions import Self
 
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
 from nifreeze.utils.ndimage import load_api
@@ -123,11 +124,10 @@ def lofo_split(self, index):
 
         return (train_data, train_timings), (test_data, test_timing)
 
-    def set_transform(self, index, affine, order=3):
+    def set_transform(self, index: int, affine: np.ndarray, order: int = 3) -> None:
         """Set an affine, and update data object and gradients."""
-        reference = namedtuple("ImageGrid", ("shape", "affine"))(
-            shape=self.dataobj.shape[:3], affine=self.affine
-        )
+        ImageGrid = namedtuple("ImageGrid", ("shape", "affine"))
+        reference = ImageGrid(self.dataobj.shape[:3], self.affine)
         xform = Affine(matrix=affine, reference=reference)
 
         if not Path(self._filepath).exists():
@@ -152,7 +152,9 @@ def set_transform(self, index, affine, order=3):
 
         self.motion_affines[index] = xform
 
-    def to_filename(self, filename, compression=None, compression_opts=None):
+    def to_filename(
+        self, filename: Path | str, compression: str | None = None, compression_opts: Any = None
+    ) -> None:
         """Write an HDF5 file to disk."""
         filename = Path(filename)
         if not filename.name.endswith(".h5"):
@@ -183,21 +185,23 @@ def to_nifti(self, filename, *_):
         nii.to_filename(filename)
 
     @classmethod
-    def from_filename(cls, filename):
+    def from_filename(cls, filename: Path | str) -> Self:
         """Read an HDF5 file from disk."""
         with h5py.File(filename, "r") as in_file:
             root = in_file["/0"]
             data = {k: np.asanyarray(v) for k, v in root.items() if not k.startswith("_")}
         return cls(**data)
 
     @classmethod
-    def load(cls, filename, json_file, brainmask_file=None):
+    def load(
+        cls, filename: Path | str, json_file: Path | str, brainmask_file: Path | str | None = None
+    ) -> Self:
         """Load PET data."""
         filename = Path(filename)
         if filename.name.endswith(".h5"):
             return cls.from_filename(filename)
 
-        img = nb.load(filename)
+        img = load_api(filename, SpatialImage)
         retval = cls(
             dataobj=img.get_fdata(dtype="float32"),
             affine=img.affine,
@@ -217,7 +221,7 @@ def load(cls, filename, json_file, brainmask_file=None):
         assert len(retval.midframe) == retval.dataobj.shape[-1]
 
         if brainmask_file:
-            mask = nb.load(brainmask_file)
+            mask = load_api(brainmask_file, SpatialImage)
             retval.brainmask = np.asanyarray(mask.dataobj)
 
         return retval

src/nifreeze/estimator.py

@@ -222,11 +222,11 @@ def run(self, dataset: DatasetT, **kwargs) -> Self:
 class PETMotionEstimator:
     """Estimates motion within PET imaging data aligned with generic Estimator workflow."""
 
-    def __init__(self, align_kwargs=None, strategy="lofo"):
+    def __init__(self, align_kwargs: dict | None = None, strategy: str = "lofo"):
         self.align_kwargs = align_kwargs or {}
         self.strategy = strategy
 
-    def run(self, pet_dataset, omp_nthreads=None):
+    def run(self, pet_dataset: PET, omp_nthreads: int | None = None) -> list:
         n_frames = len(pet_dataset)
         frame_indices = np.arange(n_frames)
 

src/nifreeze/model/pet.py

@@ -23,6 +23,7 @@
 """Models for nuclear imaging."""
 
 from os import cpu_count
+from typing import Union
 
 import nibabel as nb
 import numpy as np
@@ -31,6 +32,7 @@
 from scipy.interpolate import BSpline
 from scipy.sparse.linalg import cg
 
+from nifreeze.data.pet import PET
 from nifreeze.model.base import BaseModel
 
 DEFAULT_TIMEFRAME_MIDPOINT_TOL = 1e-2
@@ -54,13 +56,13 @@ class PETModel(BaseModel):
 
     def __init__(
         self,
-        dataset,
-        timepoints=None,
-        xlim=None,
-        n_ctrl=None,
-        order=3,
-        smooth_fwhm=10,
-        thresh_pct=20,
+        dataset: PET,
+        timepoints: list | np.ndarray | None = None,
+        xlim: float | None = None,
+        n_ctrl: int | None = None,
+        order: int = 3,
+        smooth_fwhm: float = 10.0,
+        thresh_pct: float = 20.0,
         **kwargs,
     ):
         """
@@ -105,7 +107,7 @@ def __init__(
         self._mask = None
 
     @property
-    def is_fitted(self):
+    def is_fitted(self) -> bool:
         return self._locked_fit is not None
 
     def _fit(self, index: int | None = None, n_jobs=None, **kwargs) -> int:
@@ -117,7 +119,7 @@ def _fit(self, index: int | None = None, n_jobs=None, **kwargs) -> int:
         if index is not None:
             raise NotImplementedError("Fitting with held-out data is not supported")
         timepoints = kwargs.get("timepoints", None) or self._x
-        x = (np.array(timepoints, dtype="float32") / self._xlim) * self._n_ctrl
+        x = np.asarray((np.array(timepoints, dtype="float32") / self._xlim) * self._n_ctrl)
 
         data = self._dataset.dataobj
         brainmask = self._dataset.brainmask
@@ -148,7 +150,7 @@ def _fit(self, index: int | None = None, n_jobs=None, **kwargs) -> int:
 
         return n_jobs
 
-    def fit_predict(self, index: int | None = None, **kwargs):
+    def fit_predict(self, index: int | None = None, **kwargs) -> Union[np.ndarray, None]:
         """Return the corrected volume using B-spline interpolation."""
 
         # Fit the BSpline basis on all data
@@ -159,7 +161,7 @@ def fit_predict(self, index: int | None = None, **kwargs):
             return None
 
         # Project sample timing into B-Spline coordinates
-        x = (index / self._xlim) * self._n_ctrl
+        x = np.asarray((index / self._xlim) * self._n_ctrl)
         A = BSpline.design_matrix(x, self._t, k=self._order)
 
         # A is 1 (num. timepoints) x C (num. coeff)