FIX: Fix miscellaneous oversights in the base DWI model

Fix miscellaneous oversights in the base DWI model:
- Provide the `None` default value for the `n_jobs` kwarg if not
  provided in the dictionary.
  Fixes:
  ```
          n_models = self._fit(
              index,
  >           n_jobs=kwargs.pop("n_jobs"),
              **kwargs,
          )
  E       KeyError: 'n_jobs'
  ```

- Make the diffusion gradients have a unit norm in the gtab generation
  fixture.
  Fixes:
  ```
          bvecs = np.where(np.isnan(bvecs), 0, bvecs)
          bvecs_close_to_1 = abs(vector_norm(bvecs) - 1) <= atol
          if bvecs.shape[1] != 3:
              raise ValueError("bvecs should be (N, 3)")
          if not np.all(bvecs_close_to_1[dwi_mask]):
  >           raise ValueError(
                  "The vectors in bvecs should be unit (The tolerance "
                  "can be modified as an input parameter)"
              )
  E           ValueError: The vectors in bvecs should be unit (The tolerance can be modified as an input parameter)
  ```

Add a test to check that DIPY models (e.g. DTI) can be instantiated and
that the returned prediction has the expected shape.

Author: jhlegarreta

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,7 @@ 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,33 @@ 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"""