ENH: Add type hints across miscellaneous methods

Add type hints across miscellaneous methods.

Author: jhlegarreta

src/nifreeze/data/dmri.py

@@ -369,10 +369,10 @@ 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 : int = DEFAULT_HIGHB_THRESHOLD,
 ):
     """
     Find the shelling scheme on the given b-values.

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

@@ -21,6 +21,7 @@
 #     https://www.nipreps.org/community/licensing/
 #
 
+import numbers
 from importlib import import_module
 
 import numpy as np
@@ -29,6 +30,7 @@
 from nifreeze.data.dmri import (
     DEFAULT_CLIP_PERCENTILE,
     DTI_MIN_ORIENTATIONS,
+    DWI,
 )
 from nifreeze.model.base import BaseModel, ExpectationModel
 
@@ -51,7 +53,7 @@ class BaseDWIModel(BaseModel):
         "_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 +119,7 @@ def _fit(self, index, n_jobs=None, **kwargs):
         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.
 
@@ -173,7 +175,7 @@ class AverageDWIModel(ExpectationModel):
 
     __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 : numbers.Number = 100, th_high : numbers.Number = 100, detrend : bool = False, **kwargs):
         r"""
         Implement object initialization.