enh: force parallelization even for dti and dki

Author: oesteban

src/nifreeze/model/_dipy.py

@@ -29,6 +29,7 @@
 import numpy as np
 from dipy.core.gradients import GradientTable
 from dipy.reconst.base import ReconstModel
+from dipy.reconst.multi_voxel import multi_voxel_fit
 from sklearn.gaussian_process import GaussianProcessRegressor
 
 from nifreeze.model.gpr import (
@@ -38,6 +39,11 @@
 )
 
 
+@multi_voxel_fit
+def multi_fit(obj, data, *, mask=None, **kwargs):
+    return obj.fit(data, *obj.args, mask=mask)
+
+
 def gp_prediction(
     model: GaussianProcessRegressor,
     gtab: GradientTable | np.ndarray,

src/nifreeze/model/dmri.py

@@ -123,11 +123,21 @@ def _fit(self, index: int | None = None, n_jobs=None, **kwargs):
                 class_name,
             )(gtab, **kwargs)
 
-        self._model_fit = model.fit(
-            data,
-            engine="serial" if n_jobs == 1 else "joblib",
-            n_jobs=n_jobs,
-        )
+        try:
+            self._model_fit = model.fit(
+                data,
+                engine="serial" if n_jobs == 1 else "joblib",
+                n_jobs=n_jobs,
+            )
+        except TypeError:
+            from nifreeze.model._dipy import multi_fit
+
+            self._model_fit = multi_fit(
+                model,
+                data,
+                engine="serial" if n_jobs == 1 else "ray",
+                n_jobs=n_jobs,
+            )
         return n_jobs
 
     def fit_predict(self, index: int | None = None, **kwargs):