ENH: Add type hints across miscellaneous methods

src/nifreeze/data/dmri.py

@@ -33,6 +33,7 @@
 import h5py
 import nibabel as nb
 import numpy as np
+import numpy.typing as npt
 from nibabel.spatialimages import SpatialImage
 from nitransforms.linear import Affine
 
@@ -369,11 +370,11 @@ def load(
 
 
 def find_shelling_scheme(
-    bvals,
-    num_bins=DEFAULT_NUM_BINS,
-    multishell_nonempty_bin_count_thr=DEFAULT_MULTISHELL_BIN_COUNT_THR,
-    bval_cap=DEFAULT_HIGHB_THRESHOLD,
-):
+    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.
 
@@ -390,7 +391,7 @@ def find_shelling_scheme(
         Number of bins.
     multishell_nonempty_bin_count_thr : :obj:`int`, optional
         Bin count to consider a multi-shell scheme.
-    bval_cap : :obj:`int`, optional
+    bval_cap : :obj:`float`, optional
         Maximum b-value to be considered in a multi-shell scheme.
 
     Returns

src/nifreeze/model/base.py

@@ -36,7 +36,7 @@ class ModelFactory:
     """A factory for instantiating data models."""
 
     @staticmethod
-    def init(model=None, **kwargs):
+    def init(model: str | None = None, **kwargs):
         """
         Instantiate a diffusion model.
 
@@ -136,7 +136,7 @@ def __init__(self, dataset, stat="median", **kwargs):
         super().__init__(dataset, **kwargs)
         self._stat = stat
 
-    def fit_predict(self, index, **kwargs):
+    def fit_predict(self, index: int, **kwargs):
         """
         Return the expectation map.
 

src/nifreeze/model/dmri.py

@@ -29,6 +29,7 @@
 from nifreeze.data.dmri import (
     DEFAULT_CLIP_PERCENTILE,
     DTI_MIN_ORIENTATIONS,
+    DWI,
 )
 from nifreeze.model.base import BaseModel, ExpectationModel
 
@@ -51,7 +52,7 @@
         "_modelargs": "Arguments acceptable by the underlying DIPY-like model.",
     }
 
-    def __init__(self, dataset, **kwargs):
+    def __init__(self, dataset: DWI, **kwargs):
         r"""Initialization.
 
         Parameters
@@ -117,7 +118,7 @@
         self._model = None  # Preempt further actions on the model
         return n_jobs
 
-    def fit_predict(self, index, **kwargs):
+    def fit_predict(self, index: int, **kwargs):
         """
         Predict asynchronously chunk-by-chunk the diffusion signal.
 
@@ -140,7 +141,7 @@
         if n_models == 1:
             predicted, _ = _exec_predict(self._model, **(kwargs | {"gtab": gradient, "S0": S0}))
         else:
-            S0 = np.array_split(S0, n_models) if S0 is not None else [None] * n_models
+            S0 = np.array_split(S0, n_models) if S0 is not None else np.full(n_models, None)
 
             predicted = [None] * n_models
 
@@ -173,7 +174,15 @@
 
     __slots__ = ("_th_low", "_th_high", "_detrend")
 
-    def __init__(self, dataset, stat="median", th_low=100, th_high=100, detrend=False, **kwargs):
+    def __init__(
+        self,
+        dataset: DWI,
+        stat: str = "median",
+        th_low: float = 100.0,
+        th_high: float = 100.0,
+        detrend: bool = False,
+        **kwargs,
+    ):
         r"""
         Implement object initialization.
 
@@ -183,10 +192,10 @@
             Reference to a DWI object.
         stat : :obj:`str`, optional
             Whether the summary statistic to apply is ``"mean"`` or ``"median"``.
-        th_low : :obj:`numbers.Number`, optional
+        th_low : :obj:`float`, optional
             A lower bound for the b-value corresponding to the diffusion weighted images
             that will be averaged.
-        th_high : :obj:`numbers.Number`, optional
+        th_high : :obj:`float`, optional
             An upper bound for the b-value corresponding to the diffusion weighted images
             that will be averaged.
         detrend : :obj:`bool`, optional