Fix MultiShellKernel type issues

Author: oesteban

pyproject.toml

@@ -155,6 +155,11 @@ module = [
   "joblib",
   "h5py",
   "ConfigSpace",
+  "scipy.*",
+  "sklearn.*",
+  "skimage.*",
+  "pandas",
+  "attrs",
 ]
 ignore_missing_imports = true
 

src/nifreeze/model/gpr.py

@@ -252,7 +252,7 @@ def _constrained_optimization(
                 options=options,
                 args=(self.eval_gradient,),
                 tol=self.tol,
-            )
+            )  # type: ignore[call-overload]
             return opt_res.x, opt_res.fun
 
         if callable(self.optimizer):
@@ -479,6 +479,9 @@ def __repr__(self) -> str:
 class MultiShellKernel(KernelOperator):
     """Composite kernel for multi-shell diffusion data."""
 
+    k1: Kernel
+    k2: Kernel
+
     def __init__(
         self,
         orientation_kernel: Kernel | None = None,
@@ -511,10 +514,13 @@ def __call__(
         eval_gradient: bool = False,
     ) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
         X_o, X_b = self._split(X)
+        Y_o: np.ndarray | None
+        Y_b: np.ndarray | None
         if Y is not None:
             Y_o, Y_b = self._split(Y)
         else:
-            Y_o = Y_b = None
+            Y_o = None
+            Y_b = None
 
         if eval_gradient:
             K1, g1 = self.k1(X_o, Y_o, eval_gradient=True)