From f20fe5f572521f3cca78fe5c8f9e45b104e472d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Fri, 20 Dec 2024 10:19:11 -0500
Subject: [PATCH 01/29] ENH: Add static type checking

Add static type checking: use `mypy` to ensure that variable and
function calls are using the appropriate types.

Configure and run `mypy` checks in CI.

Documentation:
https://mypy.readthedocs.io/en/stable/index.html
---
 .github/workflows/test.yml |  3 +--
 pyproject.toml             | 19 +++++++++++++++++++
 tox.ini                    |  9 +++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index a1acb0d7..43d95a5b 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -95,8 +95,7 @@ jobs:
     continue-on-error: true
     strategy:
       matrix:
-        check: ['spellcheck']
-
+        check: ['spellcheck', 'typecheck']
     steps:
       - uses: actions/checkout@v4
       - name: Install the latest version of uv
diff --git a/pyproject.toml b/pyproject.toml
index 205ac696..accd832c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -75,6 +75,15 @@ test = [
     "pytest-env",
     "pytest-xdist >= 1.28"
 ]
+types = [
+  "pandas-stubs",
+  "types-setuptools",
+  "scipy-stubs",
+  "types-PyYAML",
+  "types-tqdm",
+  "pytest",
+  "microsoft-python-type-stubs @ git+https://github.com/microsoft/python-type-stubs.git",
+]
 
 notebooks = [
     "jupyter",
@@ -138,6 +147,16 @@ version-file = "src/nifreeze/_version.py"
 # Developer tool configurations
 #
 
+[[tool.mypy.overrides]]
+module = [
+  "nipype.*",
+  "nilearn.*",
+  "nireports.*",
+  "nitransforms.*",
+  "seaborn",
+]
+ignore_missing_imports = true
+
 [tool.ruff]
 line-length = 99
 target-version = "py310"
diff --git a/tox.ini b/tox.ini
index e3672d1f..089d30f8 100644
--- a/tox.ini
+++ b/tox.ini
@@ -76,6 +76,15 @@ extras = doc
 commands =
   make -C docs/ SPHINXOPTS="-W -v" BUILDDIR="$HOME/docs" OUTDIR="${CURBRANCH:-html}" html
 
+[testenv:typecheck]
+description = Run mypy type checking
+labels = check
+deps =
+  mypy
+extras = types
+commands =
+  mypy .
+
 [testenv:spellcheck]
 description = Check spelling
 labels = check

From 734862c862588373b65ecc7b01e3cf7d3f0001d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:14:35 -0500
Subject: [PATCH 02/29] ENH: Comment empty `latex_elements` section in
 documentation config file

Comment empty `latex_elements` section in documentation config file.

Fixes:
```
docs/conf.py:157: error:
 Need type annotation for "latex_elements" (hint: "latex_elements: dict[<type>, <type>] = ...")  [var-annotated]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:31
---
 docs/conf.py | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/docs/conf.py b/docs/conf.py
index 02488e10..af5a5212 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -154,20 +154,20 @@
 
 # -- Options for LaTeX output ------------------------------------------------
 
-latex_elements = {
-    # The paper size ('letterpaper' or 'a4paper').
-    #
-    # 'papersize': 'letterpaper',
-    # The font size ('10pt', '11pt' or '12pt').
-    #
-    # 'pointsize': '10pt',
-    # Additional stuff for the LaTeX preamble.
-    #
-    # 'preamble': '',
-    # Latex figure (float) alignment
-    #
-    # 'figure_align': 'htbp',
-}
+# latex_elements = {
+#     # The paper size ('letterpaper' or 'a4paper').
+#     #
+#     # 'papersize': 'letterpaper',
+#     # The font size ('10pt', '11pt' or '12pt').
+#     #
+#     # 'pointsize': '10pt',
+#     # Additional stuff for the LaTeX preamble.
+#     #
+#     # 'preamble': '',
+#     # Latex figure (float) alignment
+#     #
+#     # 'figure_align': 'htbp',
+# }
 
 # Grouping the document tree into LaTeX files. List of tuples
 # (source start file, target name, title,

From 12de4f96c23230e682c2add860b03cbcd98c4fad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:17:26 -0500
Subject: [PATCH 03/29] ENH: Fix arguments to `unique` in GP error analysis
 plot script

Fix arguments to the `unique` function in GP error analysis plot
script: get the unique values directly from the `pandas.Series`
instance and check the length of the returned values prior to getting
the first instance.

Fixes:
```
scripts/dwi_gp_estimation_error_analysis_plot.py:92: error:
 Argument 1 to "unique" has incompatible type
 "ExtensionArray | ndarray[Any, Any]"; expected "_SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]]"  [arg-type]
scripts/dwi_gp_estimation_error_analysis_plot.py:93: error:
 Argument 1 to "unique" has incompatible type
 "ExtensionArray | ndarray[Any, Any]"; expected "_SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]]"  [arg-type]
scripts/dwi_gp_estimation_error_analysis_plot.py:94: error:
 Argument 1 to "unique" has incompatible type
```

raised for example in:
---
 scripts/dwi_gp_estimation_error_analysis_plot.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/scripts/dwi_gp_estimation_error_analysis_plot.py b/scripts/dwi_gp_estimation_error_analysis_plot.py
index 1656f6b6..252dc942 100644
--- a/scripts/dwi_gp_estimation_error_analysis_plot.py
+++ b/scripts/dwi_gp_estimation_error_analysis_plot.py
@@ -89,9 +89,17 @@ def main() -> None:
     df = pd.read_csv(args.error_data_fname, sep="\t", keep_default_na=False, na_values="n/a")
 
     # Plot the prediction error
-    kfolds = sorted(np.unique(df["n_folds"].values))
-    snr = np.unique(df["snr"].values).item()
-    bval = np.unique(df["bval"].values).item()
+    kfolds = sorted(pd.unique(df["n_folds"]))
+    snr = pd.unique(df["snr"])
+    if len(snr) == 1:
+        snr = snr[0]
+    else:
+        raise ValueError(f"More than one unique SNR value: {snr}")
+    bval = pd.unique(df["bval"])
+    if len(bval) == 1:
+        bval = bval[0]
+    else:
+        raise ValueError(f"More than one unique bval value: {bval}")
     rmse_data = [df.groupby("n_folds").get_group(k)["rmse"].values for k in kfolds]
     axis = 1
     mean = np.mean(rmse_data, axis=axis)

From b06eff8474af864459b08e8b7b161460b59ffbe0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:26:25 -0500
Subject: [PATCH 04/29] ENH: Convert list to array prior to computing mean and
 std dev

Convert list to array prior to computing mean and std dev using `NumPy`.

Fixes:
```
scripts/dwi_gp_estimation_error_analysis_plot.py:97: error:
 Argument 1 to "mean" has incompatible type
 "list[ExtensionArray | ndarray[Any, Any]]"; expected "_SupportsArray[dtype[bool_ | integer[Any] | floating[Any] | complexfloating[Any, Any]]] | _NestedSequence[_SupportsArray[dtype[bool_ | integer[Any] | floating[Any] | complexfloating[Any, Any]]]] | bool | int | float | complex | _NestedSequence[bool | int | float | complex] | _SupportsArray[dtype[object_]] | _NestedSequence[_SupportsArray[dtype[object_]]]"  [arg-type]
scripts/dwi_gp_estimation_error_analysis_plot.py:98: error:
 Argument 1 to "std" has incompatible type
 "list[ExtensionArray | ndarray[Any, Any]]"; expected "_SupportsArray[dtype[bool_ | integer[Any] | floating[Any] | complexfloating[Any, Any]]] | _NestedSequence[_SupportsArray[dtype[bool_ | integer[Any] | floating[Any] | complexfloating[Any, Any]]]] | bool | int | float | complex | _NestedSequence[bool | int | float | complex] | _SupportsArray[dtype[object_]] | _NestedSequence[_SupportsArray[dtype[object_]]]"  [arg-type]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:68
---
 scripts/dwi_gp_estimation_error_analysis_plot.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/dwi_gp_estimation_error_analysis_plot.py b/scripts/dwi_gp_estimation_error_analysis_plot.py
index 252dc942..5b0d5729 100644
--- a/scripts/dwi_gp_estimation_error_analysis_plot.py
+++ b/scripts/dwi_gp_estimation_error_analysis_plot.py
@@ -100,7 +100,7 @@ def main() -> None:
         bval = bval[0]
     else:
         raise ValueError(f"More than one unique bval value: {bval}")
-    rmse_data = [df.groupby("n_folds").get_group(k)["rmse"].values for k in kfolds]
+    rmse_data = np.asarray([df.groupby("n_folds").get_group(k)["rmse"].values for k in kfolds])
     axis = 1
     mean = np.mean(rmse_data, axis=axis)
     std_dev = np.std(rmse_data, axis=axis)

From 3eb17abb975eb81ea8e9d6dde08d4c29481e8672 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:29:15 -0500
Subject: [PATCH 05/29] BUG: Use appropriate keyword arg names to instantiate
 `SphericalKriging`

Use appropriate keyword argument names to instantiate
`SphericalKriging`: left behind in commit
https://github.com/nipreps/eddymotion/pull/244/commits/9e65c34f803cc158e036ddd7f0dfd39eb5059cff

Fixes:
```
scripts/dwi_gp_estimation_simulated_signal.py:139: error:
 Unexpected keyword argument "a" for "SphericalKriging"  [call-arg]
src/nifreeze/model/gpr.py:365: note: "SphericalKriging" defined here
scripts/dwi_gp_estimation_simulated_signal.py:139: error:
 Unexpected keyword argument "lambda_s" for "SphericalKriging"  [call-arg]
src/nifreeze/model/gpr.py:365: note: "SphericalKriging" defined here
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:105
---
 scripts/dwi_gp_estimation_simulated_signal.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/scripts/dwi_gp_estimation_simulated_signal.py b/scripts/dwi_gp_estimation_simulated_signal.py
index fd22a81a..c6803595 100644
--- a/scripts/dwi_gp_estimation_simulated_signal.py
+++ b/scripts/dwi_gp_estimation_simulated_signal.py
@@ -132,11 +132,11 @@ def main() -> None:
 
     # Fit the Gaussian Process regressor and predict on an arbitrary number of
     # directions
-    a = 1.15
-    lambda_s = 120
+    beta_a = 1.15
+    beta_l = 120
     alpha = 100
     gpr = DiffusionGPR(
-        kernel=SphericalKriging(a=a, lambda_s=lambda_s),
+        kernel=SphericalKriging(beta_a=beta_a, beta_l=beta_l),
         alpha=alpha,
         optimizer=None,
     )

From 781104aec04a37a68975b9068edd1fc16d97c0bb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:33:34 -0500
Subject: [PATCH 06/29] ENH: Select appropriate element in case `predict`
 returns a tuple

Select appropriate element in case the GPR prediction method returns a
tuple.

Fixes:
```
scripts/dwi_gp_estimation_simulated_signal.py:159: error:
 Item "tuple[ndarray[Any, Any], ...]" of
 "ndarray[Any, Any] | tuple[ndarray[Any, Any], ndarray[Any, Any], ndarray[Any, Any]] | tuple[ndarray[Any, Any], ndarray[Any, Any]]"
 has no attribute "T"  [union-attr]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:109
---
 scripts/dwi_gp_estimation_simulated_signal.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/scripts/dwi_gp_estimation_simulated_signal.py b/scripts/dwi_gp_estimation_simulated_signal.py
index c6803595..3b2081e5 100644
--- a/scripts/dwi_gp_estimation_simulated_signal.py
+++ b/scripts/dwi_gp_estimation_simulated_signal.py
@@ -154,6 +154,8 @@ def main() -> None:
     X_test = np.vstack([gtab[~gtab.b0s_mask].bvecs, sph.vertices])
 
     predictions = gpr_fit.predict(X_test)
+    if isinstance(predictions, tuple):
+        predictions = predictions[0]
 
     # Save the predicted data
     testsims.serialize_dwi(predictions.T, args.dwi_pred_data_fname)

From d42675758ff1d90656dfccac8a0f0ef824f6d23b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:38:14 -0500
Subject: [PATCH 07/29] ENH: Group keyword arguments into a single dictionary

Group keyword arguments into a single dictionary.

Fixes:
```
scripts/optimize_registration.py:133: error:
 Argument "output_transform_prefix" to "generate_command" has incompatible type "str"; expected "dict[Any, Any]"  [arg-type]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:141
---
 scripts/optimize_registration.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/scripts/optimize_registration.py b/scripts/optimize_registration.py
index d9f732ad..abc05d6f 100644
--- a/scripts/optimize_registration.py
+++ b/scripts/optimize_registration.py
@@ -127,12 +127,13 @@ async def train_coro(
             moving_path = tmp_folder / f"test-{index:04d}.nii.gz"
             (~xfm).apply(refnii, reference=refnii).to_filename(moving_path)
 
+            _kwargs = {"output_transform_prefix": f"conversion-{index:04d}", **align_kwargs}
+
             cmdline = erants.generate_command(
                 fixed_path,
                 moving_path,
                 fixedmask_path=brainmask_path,
-                output_transform_prefix=f"conversion-{index:04d}",
-                **align_kwargs,
+                **_kwargs,
             )
 
             tasks.append(

From bc48dcf336b75e06a49524bbc05b454631ec540c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:41:10 -0500
Subject: [PATCH 08/29] ENH: Use `str` instead of `Path` for
 `_parse_yaml_config` parameter

Use `str` instead of `Path` for `_parse_yaml_config` parameter since
`argparse` passes the argument as a string, not as a `Path`.

Fixes:
```
src/nifreeze/cli/parser.py:74: error:
 Argument "type" to "add_argument" of "_ActionsContainer" has incompatible type
 "Callable[[Path], dict[Any, Any]]"; expected "Callable[[str], Any] | FileType | str"  [arg-type]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:47
---
 src/nifreeze/cli/parser.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/nifreeze/cli/parser.py b/src/nifreeze/cli/parser.py
index 98655f9b..117bcde7 100644
--- a/src/nifreeze/cli/parser.py
+++ b/src/nifreeze/cli/parser.py
@@ -29,13 +29,13 @@
 import yaml
 
 
-def _parse_yaml_config(file_path: Path) -> dict:
+def _parse_yaml_config(file_path: str) -> dict:
     """
     Parse YAML configuration file.
 
     Parameters
     ----------
-    file_path : Path
+    file_path : str
         Path to the YAML configuration file.
 
     Returns

From d3591b4d2c8cca32992cb4a60872eadbaabe8046 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:44:20 -0500
Subject: [PATCH 09/29] ENH: Use arrays in NumPy's `percentile` arguments

Use arrays in NumPy's `percentile` function arguments.

Fixes:
```
src/nifreeze/data/filtering.py:80: error:
 Argument 1 to "percentile" has incompatible type "ndarray[Any, dtype[number[Any] | bool_]] | ndarray[tuple[int, ...], dtype[number[Any] | bool_]]";
 expected "_SupportsArray[dtype[bool_ | integer[Any] | floating[Any]]] | _NestedSequence[_SupportsArray[dtype[bool_ | integer[Any] | floating[Any]]]] | bool | int | float | _NestedSequence[bool | int | float]"  [arg-type]
src/nifreeze/data/filtering.py:81: error:
 Argument 1 to "percentile" has incompatible type
 "ndarray[Any, dtype[number[Any] | bool_]] | ndarray[tuple[int, ...], dtype[number[Any] | bool_]]";
 expected "_SupportsArray[dtype[bool_ | integer[Any] | floating[Any]]] | _NestedSequence[_SupportsArray[dtype[bool_ | integer[Any] | floating[Any]]]] | bool | int | float | _NestedSequence[bool | int | float]"  [arg-type]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:130
---
 src/nifreeze/data/filtering.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/nifreeze/data/filtering.py b/src/nifreeze/data/filtering.py
index 0f14604a..6f5b725e 100644
--- a/src/nifreeze/data/filtering.py
+++ b/src/nifreeze/data/filtering.py
@@ -77,8 +77,12 @@ def advanced_clip(
     # Calculate stats on denoised version to avoid outlier bias
     denoised = median_filter(data, footprint=ball(3))
 
-    a_min = np.percentile(denoised[denoised >= 0] if nonnegative else denoised, p_min)
-    a_max = np.percentile(denoised[denoised >= 0] if nonnegative else denoised, p_max)
+    a_min = np.percentile(
+        np.asarray([denoised[denoised >= 0] if nonnegative else denoised]), p_min
+    )
+    a_max = np.percentile(
+        np.asarray([denoised[denoised >= 0] if nonnegative else denoised]), p_max
+    )
 
     # Clip and scale data
     data = np.clip(data, a_min=a_min, a_max=a_max)

From 3b214c316d8c82e910971cc7f2673c084852c663 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:47:45 -0500
Subject: [PATCH 10/29] ENH: List `sigma_sq` in the GP model slots

List `sigma_sq` in the GP model slots.

Fixes:
```
src/nifreeze/model/_dipy.py:128: error:
 Trying to assign name "sigma_sq" that is not in "__slots__" of type "nifreeze.model._dipy.GaussianProcessModel"  [misc]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:102
---
 src/nifreeze/model/_dipy.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/nifreeze/model/_dipy.py b/src/nifreeze/model/_dipy.py
index b501c8b6..c5fbadae 100644
--- a/src/nifreeze/model/_dipy.py
+++ b/src/nifreeze/model/_dipy.py
@@ -87,6 +87,7 @@ class GaussianProcessModel(ReconstModel):
     __slots__ = (
         "kernel",
         "_modelfit",
+        "sigma_sq",
     )
 
     def __init__(

From 590d23594b27f4a7c61c57e836424dcef9a031c1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:50:13 -0500
Subject: [PATCH 11/29] ENH: Add the dimensionality to the `mask` ndarray
 parameter annotation

Add the dimensionality to the `mask` ndarray parameter annotation.

Fixes:
```
src/nifreeze/model/_dipy.py:140: error:
 "ndarray" expects 2 type arguments, but 1 given  [type-arg]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:103
---
 src/nifreeze/model/_dipy.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/nifreeze/model/_dipy.py b/src/nifreeze/model/_dipy.py
index c5fbadae..b37d5e5f 100644
--- a/src/nifreeze/model/_dipy.py
+++ b/src/nifreeze/model/_dipy.py
@@ -25,6 +25,7 @@
 from __future__ import annotations
 
 import warnings
+from typing import Any
 
 import numpy as np
 from dipy.core.gradients import GradientTable
@@ -138,7 +139,7 @@ def fit(
         self,
         data: np.ndarray,
         gtab: GradientTable | np.ndarray,
-        mask: np.ndarray[bool] | None = None,
+        mask: np.ndarray[bool, Any] | None = None,
         random_state: int = 0,
     ) -> GPFit:
         """Fit method of the DTI model class

From ecbd57434f3d37cf518d2693b13080e2a7be43e7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:52:07 -0500
Subject: [PATCH 12/29] ENH: Fix type hint for `figsize` parameter

Fix type hint for `figsize` parameter: use `float` as values can be
floating point numbers.

Fixes:
```
src/nifreeze/viz/signals.py:40: error:
 Incompatible default for argument "figsize" (default has type "tuple[float, float]", argument has type "tuple[int, int]")  [assignment]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:49
---
 src/nifreeze/viz/signals.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/nifreeze/viz/signals.py b/src/nifreeze/viz/signals.py
index 37798e06..2322177d 100644
--- a/src/nifreeze/viz/signals.py
+++ b/src/nifreeze/viz/signals.py
@@ -37,7 +37,7 @@ def plot_error(
     ylabel: str,
     title: str,
     color: str = "orange",
-    figsize: tuple[int, int] = (19.2, 10.8),
+    figsize: tuple[float, float] = (19.2, 10.8),
 ) -> plt.Figure:
     """
     Plot the error and standard deviation.

From 73afbf733dd1cf1e2e7f8cce7e98d739692732d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 17:57:27 -0500
Subject: [PATCH 13/29] ENH: Provide appropriate type hints to
 `reg_target_type`

Provide appropriate type hints to the `reg_target_type` parameter of the
`ants._run_registration` function. The expression is set to a tuple in
https://github.com/nipreps/nifreeze/blob/6b6ba707c47b2763c040f3d3a2ceabdcf450eefe/src/nifreeze/estimator.py#L103

Fixes:
```
src/nifreeze/registration/ants.py:459: error:
 Incompatible types in assignment (expression has type "tuple[str, str]", variable has type "str")  [assignment]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:133
---
 src/nifreeze/registration/ants.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/nifreeze/registration/ants.py b/src/nifreeze/registration/ants.py
index f6bba237..e296b808 100644
--- a/src/nifreeze/registration/ants.py
+++ b/src/nifreeze/registration/ants.py
@@ -412,7 +412,7 @@ def _run_registration(
     i_iter: int,
     vol_idx: int,
     dirname: Path,
-    reg_target_type: str,
+    reg_target_type: str | tuple[str, str],
     align_kwargs: dict,
 ) -> nt.base.BaseTransform:
     """
@@ -440,7 +440,7 @@ def _run_registration(
         DWI frame index.
     dirname : :obj:`Path`
         Directory name where the transformation is saved.
-    reg_target_type : :obj:`str`
+    reg_target_type : :obj:`str` or tuple of :obj:`str`
         Target registration type.
     align_kwargs : :obj:`dict`
         Parameters to configure the image registration process.

From db1a1e7c5796968ccb15762045cf9c3d71676331 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 18:24:08 -0500
Subject: [PATCH 14/29] ENH: Import `Bounds` from `scipy.optimize`

Import `Bounds` from `scipy.optimize`.

Fixes:
```
src/nifreeze/model/gpr.py:32: error:
 Module "scipy.optimize._minimize" does not explicitly export attribute "Bounds"  [attr-defined]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:70
---
 src/nifreeze/model/gpr.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/nifreeze/model/gpr.py b/src/nifreeze/model/gpr.py
index 19f88bf4..e4a99427 100644
--- a/src/nifreeze/model/gpr.py
+++ b/src/nifreeze/model/gpr.py
@@ -29,7 +29,7 @@
 
 import numpy as np
 from scipy import optimize
-from scipy.optimize._minimize import Bounds
+from scipy.optimize import Bounds
 from sklearn.gaussian_process import GaussianProcessRegressor
 from sklearn.gaussian_process.kernels import (
     Hyperparameter,

From 87f7e90781d32b7d43f40077338bc79e1c5adbd8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 18:26:40 -0500
Subject: [PATCH 15/29] ENH: Avoid type checking for private function import
 statement

Avoid type checking for private function import statement.

Fixes:
```
src/nifreeze/model/gpr.py:215: error:
 Module "sklearn.utils.optimize" has no attribute "_check_optimize_result"  [attr-defined]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:73
---
 src/nifreeze/model/gpr.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/nifreeze/model/gpr.py b/src/nifreeze/model/gpr.py
index e4a99427..e0b88edc 100644
--- a/src/nifreeze/model/gpr.py
+++ b/src/nifreeze/model/gpr.py
@@ -212,7 +212,7 @@ def _constrained_optimization(
     ) -> tuple[float, float]:
         options = {}
         if self.optimizer == "fmin_l_bfgs_b":
-            from sklearn.utils.optimize import _check_optimize_result
+            from sklearn.utils.optimize import _check_optimize_result  # type: ignore
 
             for name in LBFGS_CONFIGURABLE_OPTIONS:
                 if (value := getattr(self, name, None)) is not None:

From b8601d9ac491b05936e9113632bfbdc9db04d138 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 18:42:17 -0500
Subject: [PATCH 16/29] ENH: Remove unused `namedtuple` definition in test

Remove unused `namedtuple` definition in test.

Fixes:
```
test/test_gpr.py:31: error:
 First argument to namedtuple() should be "GradientTablePatch", not "gtab"  [name-match]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:129
---
 test/test_gpr.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/test/test_gpr.py b/test/test_gpr.py
index 8d199748..b714678a 100644
--- a/test/test_gpr.py
+++ b/test/test_gpr.py
@@ -20,7 +20,6 @@
 #
 #     https://www.nipreps.org/community/licensing/
 #
-from collections import namedtuple
 
 import numpy as np
 import pytest
@@ -28,9 +27,6 @@
 
 from nifreeze.model import gpr
 
-GradientTablePatch = namedtuple("gtab", ["bvals", "bvecs"])
-
-
 THETAS = np.linspace(0, np.pi / 2, num=50)
 EXPECTED_EXPONENTIAL = [
     1.0,

From 24e4ec1391fc1b19280d1e2c14547cf0927da893 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 18:59:48 -0500
Subject: [PATCH 17/29] ENH: Use `ClassVar` for class variable type hinting

Use `ClassVar` for type hinting the `GaussianProcessRegressor`
`_parameter_constraints` class variable in child class.

Fixes:
```
src/nifreeze/model/gpr.py:156: error:
 Cannot override class variable (previously declared on base class "GaussianProcessRegressor") with instance variable  [misc]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:71
---
 src/nifreeze/model/gpr.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/nifreeze/model/gpr.py b/src/nifreeze/model/gpr.py
index e0b88edc..72c33741 100644
--- a/src/nifreeze/model/gpr.py
+++ b/src/nifreeze/model/gpr.py
@@ -25,7 +25,7 @@
 from __future__ import annotations
 
 from numbers import Integral, Real
-from typing import Callable, Mapping, Sequence
+from typing import Callable, ClassVar, Mapping, Sequence
 
 import numpy as np
 from scipy import optimize
@@ -153,7 +153,7 @@ class DiffusionGPR(GaussianProcessRegressor):
 
     """
 
-    _parameter_constraints: dict = {
+    _parameter_constraints: ClassVar[dict] = {
         "kernel": [None, Kernel],
         "alpha": [Interval(Real, 0, None, closed="left"), np.ndarray],
         "optimizer": [StrOptions(SUPPORTED_OPTIMIZERS), callable, None],

From cc735d1d41ef13068ee97e0d37decc4d81fe8337 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sat, 21 Dec 2024 19:03:59 -0500
Subject: [PATCH 18/29] ENH: Annotate `optimizer` attribute type DiffusionGPR

Annotate the `optimizer` attribute type in the DiffusionGPR GPR child
class.

Fixes:
```
src/nifreeze/model/gpr.py:234: error:
 "DiffusionGPR" has no attribute "optimizer"  [attr-defined]
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:82
---
 src/nifreeze/model/gpr.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/nifreeze/model/gpr.py b/src/nifreeze/model/gpr.py
index 72c33741..20821f38 100644
--- a/src/nifreeze/model/gpr.py
+++ b/src/nifreeze/model/gpr.py
@@ -25,7 +25,7 @@
 from __future__ import annotations
 
 from numbers import Integral, Real
-from typing import Callable, ClassVar, Mapping, Sequence
+from typing import Callable, ClassVar, Mapping, Optional, Sequence, Union
 
 import numpy as np
 from scipy import optimize
@@ -153,6 +153,8 @@ class DiffusionGPR(GaussianProcessRegressor):
 
     """
 
+    optimizer: Optional[Union[StrOptions, Callable, None]] = None
+
     _parameter_constraints: ClassVar[dict] = {
         "kernel": [None, Kernel],
         "alpha": [Interval(Real, 0, None, closed="left"), np.ndarray],

From 052cf36bc29e687938bf89d6c023eb009f43cc88 Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 09:48:24 -0500
Subject: [PATCH 19/29] chore: Ignore more untyped dependencies

---
 pyproject.toml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index accd832c..6252ce98 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -154,6 +154,11 @@ module = [
   "nireports.*",
   "nitransforms.*",
   "seaborn",
+  "dipy.*",
+  "smac.*",
+  "joblib",
+  "h5py",
+  "ConfigSpace",
 ]
 ignore_missing_imports = true
 

From 374fe48067b5a1431764f45a7672a4d9c945d1c9 Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 09:49:02 -0500
Subject: [PATCH 20/29] type: Fix annotations in data.base and utils.iterators

---
 src/nifreeze/data/base.py       | 25 ++++++++++++++-----------
 src/nifreeze/utils/iterators.py | 10 +++++-----
 test/test_data_base.py          |  3 ++-
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/src/nifreeze/data/base.py b/src/nifreeze/data/base.py
index 9eae6edd..147d06ae 100644
--- a/src/nifreeze/data/base.py
+++ b/src/nifreeze/data/base.py
@@ -33,8 +33,12 @@
 import h5py
 import nibabel as nb
 import numpy as np
+from nibabel.arrayproxy import ArrayLike
+from nibabel.spatialimages import SpatialHeader, SpatialImage
 from nitransforms.linear import Affine
 
+from ..utils.ndimage import load_api
+
 NFDH5_EXT = ".h5"
 
 
@@ -68,15 +72,15 @@ class BaseDataset:
 
     """
 
-    dataobj = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
+    dataobj: np.ndarray = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
     """A :obj:`~numpy.ndarray` object for the data array."""
-    affine = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
+    affine: np.ndarray = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
     """Best affine for RAS-to-voxel conversion of coordinates (NIfTI header)."""
-    brainmask = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
+    brainmask: np.ndarray = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
     """A boolean ndarray object containing a corresponding brainmask."""
-    motion_affines = attr.ib(default=None, eq=attr.cmp_using(eq=_cmp))
+    motion_affines: np.ndarray = attr.ib(default=None, eq=attr.cmp_using(eq=_cmp))
     """List of :obj:`~nitransforms.linear.Affine` realigning the dataset."""
-    datahdr = attr.ib(default=None)
+    datahdr: SpatialHeader = attr.ib(default=None)
     """A :obj:`~nibabel.spatialimages.SpatialHeader` header corresponding to the data."""
 
     _filepath = attr.ib(
@@ -159,9 +163,8 @@ def set_transform(self, index: int, affine: np.ndarray, order: int = 3) -> None:
             The order of the spline interpolation.
 
         """
-        reference = namedtuple("ImageGrid", ("shape", "affine"))(
-            shape=self.dataobj.shape[:3], affine=self.affine
-        )
+        ImageGrid = namedtuple("ImageGrid", ("shape", "affine"))
+        reference = ImageGrid(shape=self.dataobj.shape[:3], affine=self.affine)
 
         xform = Affine(matrix=affine, reference=reference)
 
@@ -279,11 +282,11 @@ def load(
     if filename.name.endswith(NFDH5_EXT):
         return BaseDataset.from_filename(filename)
 
-    img = nb.load(filename)
-    retval = BaseDataset(dataobj=img.dataobj, affine=img.affine)
+    img = load_api(filename, SpatialImage)
+    retval = BaseDataset(dataobj=np.asanyarray(img.dataobj), affine=img.affine)
 
     if brainmask_file:
-        mask = nb.load(brainmask_file)
+        mask = load_api(brainmask_file, SpatialImage)
         retval.brainmask = np.asanyarray(mask.dataobj)
 
     return retval
diff --git a/src/nifreeze/utils/iterators.py b/src/nifreeze/utils/iterators.py
index 83886e5d..0cf52d21 100644
--- a/src/nifreeze/utils/iterators.py
+++ b/src/nifreeze/utils/iterators.py
@@ -27,7 +27,7 @@
 from typing import Iterator
 
 
-def linear_iterator(size: int = None, **kwargs) -> Iterator[int]:
+def linear_iterator(size: int | None = None, **kwargs) -> Iterator[int]:
     """
     Traverse the dataset volumes in ascending order.
 
@@ -53,10 +53,10 @@ def linear_iterator(size: int = None, **kwargs) -> Iterator[int]:
     if size is None:
         raise TypeError("Cannot build iterator without size")
 
-    return range(size)
+    return iter(range(size))
 
 
-def random_iterator(size: int = None, **kwargs) -> Iterator[int]:
+def random_iterator(size: int | None = None, **kwargs) -> Iterator[int]:
     """
     Traverse the dataset volumes randomly.
 
@@ -105,7 +105,7 @@ def random_iterator(size: int = None, **kwargs) -> Iterator[int]:
     return (x for x in index_order)
 
 
-def bvalue_iterator(size: int = None, **kwargs) -> Iterator[int]:
+def bvalue_iterator(size: int | None = None, **kwargs) -> Iterator[int]:
     """
     Traverse the volumes in a DWI dataset by growing b-value.
 
@@ -132,7 +132,7 @@ def bvalue_iterator(size: int = None, **kwargs) -> Iterator[int]:
     return (index[1] for index in indexed_bvals)
 
 
-def centralsym_iterator(size: int = None, **kwargs) -> Iterator[int]:
+def centralsym_iterator(size: int | None = None, **kwargs) -> Iterator[int]:
     """
     Traverse the dataset starting from the center and alternatingly progressing to the sides.
 
diff --git a/test/test_data_base.py b/test/test_data_base.py
index 7f2c7956..abf6c400 100644
--- a/test/test_data_base.py
+++ b/test/test_data_base.py
@@ -90,6 +90,7 @@ def test_set_transform(random_dataset: BaseDataset):
     assert random_dataset.motion_affines is not None
     np.testing.assert_array_equal(random_dataset.motion_affines[idx], affine)
     # The returned affine from __getitem__ should be the same.
+    assert aff0 is not None
     np.testing.assert_array_equal(aff0, affine)
 
 
@@ -121,7 +122,7 @@ def test_to_nifti(random_dataset: BaseDataset):
         assert nifti_file.is_file()
 
         # Load the saved file with nibabel
-        img = nb.load(nifti_file)
+        img = nb.Nifti1Image.from_filename(nifti_file)
         data = img.get_fdata(dtype=np.float32)
         assert data.shape == (32, 32, 32, 5)
         assert np.allclose(data, random_dataset.dataobj)

From 478c8602705e3d3921e5bea7ae3c5ee7e481a2fc Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 09:50:30 -0500
Subject: [PATCH 21/29] type: Fix expected type of set_xticklabels arg

---
 src/nifreeze/viz/signals.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/nifreeze/viz/signals.py b/src/nifreeze/viz/signals.py
index 2322177d..5c2db0b0 100644
--- a/src/nifreeze/viz/signals.py
+++ b/src/nifreeze/viz/signals.py
@@ -74,7 +74,7 @@ def plot_error(
     ax.set_xlabel(xlabel)
     ax.set_ylabel(ylabel)
     ax.set_xticks(kfolds)
-    ax.set_xticklabels(kfolds)
+    ax.set_xticklabels(map(str, kfolds))
     ax.set_title(title)
     fig.tight_layout()
     return fig

From 0037efdd9518cd437bd2a6fc86d2eb441e643e1d Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 10:10:37 -0500
Subject: [PATCH 22/29] chore: Fix DiffusionGPR annotations

---
 src/nifreeze/model/gpr.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/nifreeze/model/gpr.py b/src/nifreeze/model/gpr.py
index 20821f38..08805373 100644
--- a/src/nifreeze/model/gpr.py
+++ b/src/nifreeze/model/gpr.py
@@ -25,7 +25,7 @@
 from __future__ import annotations
 
 from numbers import Integral, Real
-from typing import Callable, ClassVar, Mapping, Optional, Sequence, Union
+from typing import Callable, ClassVar, Literal, Mapping, Optional, Sequence, Union
 
 import numpy as np
 from scipy import optimize
@@ -171,7 +171,7 @@ def __init__(
         kernel: Kernel | None = None,
         *,
         alpha: float = 0.5,
-        optimizer: str | Callable | None = "fmin_l_bfgs_b",
+        optimizer: Literal["fmin_l_bfgs_b"] | Callable | None = "fmin_l_bfgs_b",
         n_restarts_optimizer: int = 0,
         copy_X_train: bool = True,
         normalize_y: bool = True,
@@ -229,7 +229,7 @@ def _constrained_optimization(
                 options=options,
                 args=(self.eval_gradient,),
                 tol=self.tol,
-            )
+            )  # type: ignore[call-overload]
             _check_optimize_result("lbfgs", opt_res)
             return opt_res.x, opt_res.fun
 

From 9f3fdd54e356aa1630433ee28cb864ad98bda9bc Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 10:36:30 -0500
Subject: [PATCH 23/29] type: Annotate PET Dataset

Requires making BaseDataset generic on indexable fields
---
 src/nifreeze/data/base.py | 20 +++++++++++++-------
 src/nifreeze/data/pet.py  | 28 ++++++++++++++++------------
 test/test_data_base.py    |  4 ++--
 3 files changed, 31 insertions(+), 21 deletions(-)

diff --git a/src/nifreeze/data/base.py b/src/nifreeze/data/base.py
index 147d06ae..3eea6561 100644
--- a/src/nifreeze/data/base.py
+++ b/src/nifreeze/data/base.py
@@ -27,13 +27,12 @@
 from collections import namedtuple
 from pathlib import Path
 from tempfile import mkdtemp
-from typing import Any
+from typing import Any, Generic, TypeVarTuple
 
 import attr
 import h5py
 import nibabel as nb
 import numpy as np
-from nibabel.arrayproxy import ArrayLike
 from nibabel.spatialimages import SpatialHeader, SpatialImage
 from nitransforms.linear import Affine
 
@@ -42,6 +41,9 @@
 NFDH5_EXT = ".h5"
 
 
+Ts = TypeVarTuple("Ts")
+
+
 def _data_repr(value: np.ndarray | None) -> str:
     if value is None:
         return "None"
@@ -56,7 +58,7 @@ def _cmp(lh: Any, rh: Any) -> bool:
 
 
 @attr.s(slots=True)
-class BaseDataset:
+class BaseDataset(Generic[*Ts]):
     """
     Base dataset representation structure.
 
@@ -97,9 +99,13 @@ def __len__(self) -> int:
 
         return self.dataobj.shape[-1]
 
+    def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[*Ts]:
+        # PY312: Default values for TypeVarTuples are not yet supported
+        return ()  # type: ignore[return-value]
+
     def __getitem__(
         self, idx: int | slice | tuple | np.ndarray
-    ) -> tuple[np.ndarray, np.ndarray | None]:
+    ) -> tuple[np.ndarray, np.ndarray | None, *Ts]:
         """
         Returns volume(s) and corresponding affine(s) through fancy indexing.
 
@@ -122,7 +128,7 @@ def __getitem__(
             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
+        return self.dataobj[..., idx], affine, *self._getextra(idx)
 
     @classmethod
     def from_filename(cls, filename: Path | str) -> BaseDataset:
@@ -250,7 +256,7 @@ def load(
     filename: Path | str,
     brainmask_file: Path | str | None = None,
     motion_file: Path | str | None = None,
-) -> BaseDataset:
+) -> BaseDataset[()]:
     """
     Load 4D data from a filename or an HDF5 file.
 
@@ -283,7 +289,7 @@ def load(
         return BaseDataset.from_filename(filename)
 
     img = load_api(filename, SpatialImage)
-    retval = BaseDataset(dataobj=np.asanyarray(img.dataobj), affine=img.affine)
+    retval: BaseDataset[()] = BaseDataset(dataobj=np.asanyarray(img.dataobj), affine=img.affine)
 
     if brainmask_file:
         mask = load_api(brainmask_file, SpatialImage)
diff --git a/src/nifreeze/data/pet.py b/src/nifreeze/data/pet.py
index dd0e4543..74753255 100644
--- a/src/nifreeze/data/pet.py
+++ b/src/nifreeze/data/pet.py
@@ -29,14 +29,15 @@
 
 import attr
 import h5py
-import nibabel as nb
 import numpy as np
+from nibabel.spatialimages import SpatialImage
 
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
+from nifreeze.utils.ndimage import load_api
 
 
 @attr.s(slots=True)
-class PET(BaseDataset):
+class PET(BaseDataset[np.ndarray | None]):
     """Data representation structure for PET data."""
 
     frame_time: np.ndarray | None = attr.ib(
@@ -46,6 +47,10 @@ class PET(BaseDataset):
     total_duration: float | None = attr.ib(default=None, repr=True)
     """A float representing the total duration of the dataset."""
 
+    def _getextra(self, idx: int | slice | tuple | np.ndarray) -> tuple[np.ndarray | None]:
+        return (self.frame_time[idx] if self.frame_time is not None else None,)
+
+    # For the sake of the docstring
     def __getitem__(
         self, idx: int | slice | tuple | np.ndarray
     ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None]:
@@ -69,9 +74,7 @@ def __getitem__(
             The frame time corresponding to the index(es).
 
         """
-
-        data, affine = super().__getitem__(idx)
-        return data, affine, self.frame_time[idx]
+        return super().__getitem__(idx)
 
     @classmethod
     def from_filename(cls, filename: Union[str, Path]) -> PET:
@@ -181,7 +184,7 @@ def load(
         pet_obj = PET.from_filename(filename)
     else:
         # Load from NIfTI
-        img = nb.load(str(filename))
+        img = load_api(filename, SpatialImage)
         data = img.get_fdata(dtype=np.float32)
         pet_obj = PET(
             dataobj=data,
@@ -204,11 +207,12 @@ def load(
 
     # If the user doesn't provide frame_duration, we derive it:
     if frame_duration is None:
-        frame_time_arr = pet_obj.frame_time
-        # If shape is e.g. (N,), then we can do
-        durations = np.diff(frame_time_arr)
-        if len(durations) == (len(frame_time_arr) - 1):
-            durations = np.append(durations, durations[-1])  # last frame same as second-last
+        if pet_obj.frame_time is not None:
+            frame_time_arr = pet_obj.frame_time
+            # If shape is e.g. (N,), then we can do
+            durations = np.diff(frame_time_arr)
+            if len(durations) == (len(frame_time_arr) - 1):
+                durations = np.append(durations, durations[-1])  # last frame same as second-last
     else:
         durations = np.array(frame_duration, dtype=np.float32)
 
@@ -218,7 +222,7 @@ def load(
 
     # If a brain mask is provided, load and attach
     if brainmask_file is not None:
-        mask_img = nb.load(str(brainmask_file))
+        mask_img = load_api(brainmask_file, SpatialImage)
         pet_obj.brainmask = np.asanyarray(mask_img.dataobj, dtype=bool)
 
     return pet_obj
diff --git a/test/test_data_base.py b/test/test_data_base.py
index abf6c400..4afa6446 100644
--- a/test/test_data_base.py
+++ b/test/test_data_base.py
@@ -53,7 +53,7 @@ def test_len(random_dataset: BaseDataset):
     assert len(random_dataset) == 5  # last dimension is 5 volumes
 
 
-def test_getitem_volume_index(random_dataset: BaseDataset):
+def test_getitem_volume_index(random_dataset: BaseDataset[()]):
     """
     Test that __getitem__ returns the correct (volume, affine) tuple.
 
@@ -71,7 +71,7 @@ def test_getitem_volume_index(random_dataset: BaseDataset):
     assert aff_slice is None
 
 
-def test_set_transform(random_dataset: BaseDataset):
+def test_set_transform(random_dataset: BaseDataset[()]):
     """
     Test that calling set_transform changes the data and motion_affines.
     For simplicity, we'll apply an identity transform and check that motion_affines is updated.

From 05237a6892223dc729973512efd641a389fd4170 Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 10:46:03 -0500
Subject: [PATCH 24/29] type: Fix annotations for dMRI model

---
 src/nifreeze/data/base.py |  4 ++--
 src/nifreeze/data/dmri.py | 17 +++++++++++------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/nifreeze/data/base.py b/src/nifreeze/data/base.py
index 3eea6561..4fee5a01 100644
--- a/src/nifreeze/data/base.py
+++ b/src/nifreeze/data/base.py
@@ -36,7 +36,7 @@
 from nibabel.spatialimages import SpatialHeader, SpatialImage
 from nitransforms.linear import Affine
 
-from ..utils.ndimage import load_api
+from nifreeze.utils.ndimage import load_api
 
 NFDH5_EXT = ".h5"
 
@@ -236,7 +236,7 @@ def to_filename(
                         compression_opts=compression_opts,
                     )
 
-    def to_nifti(self, filename: Path) -> None:
+    def to_nifti(self, filename: Path | str) -> None:
         """
         Write a NIfTI file to disk.
 
diff --git a/src/nifreeze/data/dmri.py b/src/nifreeze/data/dmri.py
index 6a32496c..69423c8c 100644
--- a/src/nifreeze/data/dmri.py
+++ b/src/nifreeze/data/dmri.py
@@ -33,13 +33,15 @@
 import h5py
 import nibabel as nb
 import numpy as np
+from nibabel.spatialimages import SpatialImage
 from nitransforms.linear import Affine
 
 from nifreeze.data.base import BaseDataset, _cmp, _data_repr
+from nifreeze.utils.ndimage import load_api
 
 
 @attr.s(slots=True)
-class DWI(BaseDataset):
+class DWI(BaseDataset[np.ndarray | None]):
     """Data representation structure for dMRI data."""
 
     bzero = attr.ib(default=None, repr=_data_repr, eq=attr.cmp_using(eq=_cmp))
@@ -49,6 +51,10 @@ class DWI(BaseDataset):
     eddy_xfms = attr.ib(default=None)
     """List of transforms to correct for estimatted 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,)
+
+    # For the sake of the docstring
     def __getitem__(
         self, idx: int | slice | tuple | np.ndarray
     ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray | None]:
@@ -74,8 +80,7 @@ def __getitem__(
 
         """
 
-        data, affine = super().__getitem__(idx)
-        return data, affine, self.gradients[..., idx]
+        return super().__getitem__(idx)
 
     @classmethod
     def from_filename(cls, filename: Path | str) -> DWI:
@@ -276,7 +281,7 @@ def load(
         return DWI.from_filename(filename)
 
     # 2) Otherwise, load a NIfTI
-    img = nb.load(str(filename))
+    img = load_api(filename, SpatialImage)
     fulldata = img.get_fdata(dtype=np.float32)
     affine = img.affine
 
@@ -319,7 +324,7 @@ def load(
     # 6) b=0 volume (bzero)
     #    If the user provided a b0_file, load it
     if b0_file:
-        b0img = nb.load(str(b0_file))
+        b0img = load_api(b0_file, SpatialImage)
         b0vol = np.asanyarray(b0img.dataobj)
         # We'll assume your DWI class has a bzero: np.ndarray | None attribute
         dwi_obj.bzero = b0vol
@@ -333,7 +338,7 @@ def load(
 
     # 7) If a brainmask_file was provided, load it
     if brainmask_file:
-        mask_img = nb.load(str(brainmask_file))
+        mask_img = load_api(brainmask_file, SpatialImage)
         dwi_obj.brainmask = np.asanyarray(mask_img.dataobj, dtype=bool)
 
     return dwi_obj

From 4eb80dd670c7c0949e495250de54c34530eba640 Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 10:46:59 -0500
Subject: [PATCH 25/29] type: Fix complaint about mismatched return type

model.predict can return 1, 2 or 3 arrays
---
 src/nifreeze/model/_dipy.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/nifreeze/model/_dipy.py b/src/nifreeze/model/_dipy.py
index b37d5e5f..4f4b3020 100644
--- a/src/nifreeze/model/_dipy.py
+++ b/src/nifreeze/model/_dipy.py
@@ -79,7 +79,9 @@ def gp_prediction(
         raise RuntimeError("Model is not yet fitted.")
 
     # Extract orientations from bvecs, and highly likely, the b-value too.
-    return model.predict(X, return_std=return_std)
+    orientations = model.predict(X, return_std=return_std)
+    assert isinstance(orientations, np.ndarray)
+    return orientations
 
 
 class GaussianProcessModel(ReconstModel):

From e240ac6e2727b40ef05528b02b0e7052edd4cefb Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 13:23:17 -0500
Subject: [PATCH 26/29] type: Ignore warning related to bad upstream stubs

---
 scripts/dwi_gp_estimation_error_analysis.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/scripts/dwi_gp_estimation_error_analysis.py b/scripts/dwi_gp_estimation_error_analysis.py
index eb393b97..4ae7c5bf 100644
--- a/scripts/dwi_gp_estimation_error_analysis.py
+++ b/scripts/dwi_gp_estimation_error_analysis.py
@@ -74,7 +74,14 @@ def cross_validate(
     """
 
     rkf = RepeatedKFold(n_splits=cv, n_repeats=n_repeats)
-    scores = cross_val_score(gpr, X, y, scoring="neg_root_mean_squared_error", cv=rkf)
+    # scikit-learn stubs do not recognize rkf as a BaseCrossValidator
+    scores = cross_val_score(
+        gpr,
+        X,
+        y,
+        scoring="neg_root_mean_squared_error",
+        cv=rkf,  # type: ignore[arg-type]
+    )
     return scores
 
 

From b450b8d57e7f4e27b459686d0da1de72a5dbe0a1 Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 13:23:57 -0500
Subject: [PATCH 27/29] type: Fixes

---
 scripts/dwi_gp_estimation_error_analysis.py | 8 ++++----
 src/nifreeze/registration/ants.py           | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/scripts/dwi_gp_estimation_error_analysis.py b/scripts/dwi_gp_estimation_error_analysis.py
index 4ae7c5bf..3feef2c1 100644
--- a/scripts/dwi_gp_estimation_error_analysis.py
+++ b/scripts/dwi_gp_estimation_error_analysis.py
@@ -49,7 +49,7 @@ def cross_validate(
     cv: int,
     n_repeats: int,
     gpr: DiffusionGPR,
-) -> dict[int, list[tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]]:
+) -> np.ndarray:
     """
     Perform the experiment by estimating the dMRI signal using a Gaussian process model.
 
@@ -211,10 +211,10 @@ def main() -> None:
 
     if args.kfold:
         # Use Scikit-learn cross validation
-        scores = defaultdict(list, {})
+        scores: dict[str, list] = defaultdict(list, {})
         for n in args.kfold:
             for i in range(args.repeats):
-                cv_scores = -1.0 * cross_validate(X, y.T, n, gpr)
+                cv_scores = -1.0 * cross_validate(X, y.T, n, i, gpr)
                 scores["rmse"] += cv_scores.tolist()
                 scores["repeat"] += [i] * len(cv_scores)
                 scores["n_folds"] += [n] * len(cv_scores)
@@ -224,7 +224,7 @@ def main() -> None:
             print(f"Finished {n}-fold cross-validation")
 
         scores_df = pd.DataFrame(scores)
-        scores_df.to_csv(args.output_scores, sep="\t", index=None, na_rep="n/a")
+        scores_df.to_csv(args.output_scores, sep="\t", index=False, na_rep="n/a")
 
         grouped = scores_df.groupby(["n_folds"])
         print(grouped[["rmse"]].mean())
diff --git a/src/nifreeze/registration/ants.py b/src/nifreeze/registration/ants.py
index e296b808..712a0874 100644
--- a/src/nifreeze/registration/ants.py
+++ b/src/nifreeze/registration/ants.py
@@ -213,7 +213,7 @@ def generate_command(
     movingmask_path: str | Path | list[str] | None = None,
     init_affine: str | Path | None = None,
     default: str = "b0-to-b0_level0",
-    **kwargs: dict,
+    **kwargs,
 ) -> str:
     """
     Generate an ANTs' command line.

From 231094579137801c02ac1b7c06b5c952698f19b6 Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 15:03:45 -0500
Subject: [PATCH 28/29] doc: Unmock numpy to allow np.ndarray to be found as a
 type

---
 docs/conf.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/docs/conf.py b/docs/conf.py
index af5a5212..93397d99 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -54,7 +54,6 @@
     "nipype",
     "nitime",
     "nitransforms",
-    "numpy",
     "pandas",
     "scipy",
     "seaborn",

From 02ff9c4120a0e21887f669dd67e78c8b31b0838a Mon Sep 17 00:00:00 2001
From: "Christopher J. Markiewicz" <markiewicz@stanford.edu>
Date: Wed, 22 Jan 2025 15:16:15 -0500
Subject: [PATCH 29/29] type: Annotate `Kernel.diag()` argument `X`

Annotate `Kernel.diag()` argument `X`: use `npt.ArrayLike`.

Fixes:
```
src/nifreeze/model/gpr.py:335: error:
 Argument 1 of "diag" is incompatible with supertype "Kernel";
 supertype defines the argument type as
 "Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | bool | int | float | complex | str | bytes | _NestedSequence[bool | int | float | complex | str | bytes]"
  [override]
src/nifreeze/model/gpr.py:335: note: This violates the Liskov substitution principle
src/nifreeze/model/gpr.py:335: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
src/nifreeze/model/gpr.py:445: error:
 Argument 1 of "diag" is incompatible with supertype "Kernel";
 supertype defines the argument type as
 "Buffer | _SupportsArray[dtype[Any]] | _NestedSequence[_SupportsArray[dtype[Any]]] | bool | int | float | complex | str | bytes | _NestedSequence[bool | int | float | complex | str | bytes]"
  [override]
src/nifreeze/model/gpr.py:445: note: This violates the Liskov substitution principle
src/nifreeze/model/gpr.py:335: note: See https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
```

raised for example in:
https://github.com/nipreps/nifreeze/actions/runs/12437972140/job/34728973936#step:8:93

Documentation:
https://mypy.readthedocs.io/en/stable/common_issues.html#incompatible-overrides
---
 src/nifreeze/model/gpr.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/nifreeze/model/gpr.py b/src/nifreeze/model/gpr.py
index 08805373..a3c612bb 100644
--- a/src/nifreeze/model/gpr.py
+++ b/src/nifreeze/model/gpr.py
@@ -28,6 +28,7 @@
 from typing import Callable, ClassVar, Literal, Mapping, Optional, Sequence, Union
 
 import numpy as np
+import numpy.typing as npt
 from scipy import optimize
 from scipy.optimize import Bounds
 from sklearn.gaussian_process import GaussianProcessRegressor
@@ -334,7 +335,7 @@ def __call__(
 
         return self.beta_l * C_theta, K_gradient
 
-    def diag(self, X: np.ndarray) -> np.ndarray:
+    def diag(self, X: npt.ArrayLike) -> np.ndarray:
         """Returns the diagonal of the kernel k(X, X).
 
         The result of this method is identical to np.diag(self(X)); however,
@@ -351,7 +352,7 @@ def diag(self, X: np.ndarray) -> np.ndarray:
         K_diag : :obj:`~numpy.ndarray` of shape (n_samples_X,)
             Diagonal of kernel k(X, X)
         """
-        return self.beta_l * np.ones(X.shape[0])
+        return self.beta_l * np.ones(np.asanyarray(X).shape[0])
 
     def is_stationary(self) -> bool:
         """Returns whether the kernel is stationary."""
@@ -444,7 +445,7 @@ def __call__(
 
         return self.beta_l * C_theta, K_gradient
 
-    def diag(self, X: np.ndarray) -> np.ndarray:
+    def diag(self, X: npt.ArrayLike) -> np.ndarray:
         """Returns the diagonal of the kernel k(X, X).
 
         The result of this method is identical to np.diag(self(X)); however,
@@ -461,7 +462,7 @@ def diag(self, X: np.ndarray) -> np.ndarray:
         K_diag : :obj:`~numpy.ndarray` of shape (n_samples_X,)
             Diagonal of kernel k(X, X)
         """
-        return self.beta_l * np.ones(X.shape[0])
+        return self.beta_l * np.ones(np.asanyarray(X).shape[0])
 
     def is_stationary(self) -> bool:
         """Returns whether the kernel is stationary."""