FIX: Fix miscellaneous oversights in the base DWI model

src/nifreeze/model/dmri.py

@@ -178,7 +178,7 @@ def fit_predict(self, index: int | None = None, **kwargs):
         kwargs.pop("omp_nthreads", None)  # Drop omp_nthreads
         n_models = self._fit(
             index,
-            n_jobs=kwargs.pop("n_jobs"),
+            n_jobs=kwargs.pop("n_jobs", None),
             **kwargs,
         )
 

test/conftest.py

@@ -29,6 +29,7 @@
 import nitransforms as nt
 import numpy as np
 import pytest
+from dipy.core.geometry import normalized_vector
 from dipy.io.gradients import read_bvals_bvecs
 
 from nifreeze.data.dmri import DWI
@@ -251,7 +252,9 @@ def setup_random_gtab_data(request):
     bvals_shells = _generate_random_choices(request, shells, n_gradients)
 
     bvals = np.hstack([b0s * [0], bvals_shells])
-    bvecs = np.hstack([np.zeros((3, b0s)), rng.random((3, n_gradients))])
+    bvecs = np.hstack(
+        [np.zeros((3, b0s)), normalized_vector(rng.random((3, n_gradients)), axis=0)]
+    )
 
     return bvals, bvecs
 

test/test_model.py

@@ -157,6 +157,31 @@ def test_gp_model(evals, S0, snr, hsph_dirs, bval_shell):
     assert prediction.shape == (2,)
 
 
+@pytest.mark.random_dwi_data(50, (14, 16, 8), True)
+def test_dti_model(setup_random_dwi_data):
+    (
+        dwi_dataobj,
+        affine,
+        brainmask_dataobj,
+        b0_dataobj,
+        gradients,
+        _,
+    ) = setup_random_dwi_data
+
+    dataset = DWI(
+        dataobj=dwi_dataobj,
+        affine=affine,
+        brainmask=brainmask_dataobj,
+        bzero=b0_dataobj,
+        gradients=gradients,
+    )
+
+    dtimodel = model.DTIModel(dataset)
+    predicted = dtimodel.fit_predict(4)
+
+    assert predicted.shape == dwi_dataobj.shape[:-1]
+
+
 def test_factory(datadir):
     """Check that the two different initialisations result in the same models"""