From d1ca5aa7cb5bcab70b12dbb6551fba3bc47d8ab4 Mon Sep 17 00:00:00 2001
From: Oscar Esteban <code@oscaresteban.es>
Date: Wed, 23 Oct 2024 11:44:36 +0200
Subject: [PATCH] fix: update test to use Jon's testing utilities

---
 src/eddymotion/model/_dipy.py |  4 +--
 test/test_model.py            | 47 ++++++++++++++++++++++++-----------
 2 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/src/eddymotion/model/_dipy.py b/src/eddymotion/model/_dipy.py
index b4e1971d..e64b9e27 100644
--- a/src/eddymotion/model/_dipy.py
+++ b/src/eddymotion/model/_dipy.py
@@ -67,7 +67,7 @@ def gp_prediction(
 
     """
 
-    X = gtab.bvecs.T if hasattr("bvecs") else np.asarray(gtab)
+    X = gtab.bvecs.T if hasattr(gtab, "bvecs") else np.asarray(gtab)
 
     # Check it's fitted as they do in sklearn internally
     # https://github.com/scikit-learn/scikit-learn/blob/972e17fe1aa12d481b120ad4a3dc076bae736931/\
@@ -151,7 +151,7 @@ def fit(
 
         # Extract b-vecs: Scikit learn wants (n_samples, n_features)
         # where n_features is 3, and n_samples the different diffusion orientations.
-        X = gtab.bvecs.T if hasattr("bvecs") else np.asarray(gtab)
+        X = gtab.bvecs if hasattr(gtab, "bvecs") else np.asarray(gtab)
 
         # Data must be shapes (n_samples, n_targets) where n_samples is
         # the number of diffusion orientations, and n_targets is number of voxels.
diff --git a/test/test_model.py b/test/test_model.py
index 4bb8de65..2f50e6f9 100644
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -24,8 +24,7 @@
 
 import numpy as np
 import pytest
-from dipy.core.gradients import gradient_table
-from sklearn.datasets import make_regression
+from dipy.sims.voxel import single_tensor
 
 from eddymotion import model
 from eddymotion.data.dmri import DWI
@@ -33,6 +32,7 @@
 from eddymotion.exceptions import ModelNotFittedError
 from eddymotion.model._dipy import GaussianProcessModel
 from eddymotion.model.dmri import DEFAULT_MAX_S0, DEFAULT_MIN_S0
+from eddymotion.testing import simulations as _sim
 
 
 def test_trivial_model():
@@ -108,22 +108,39 @@ def test_average_model():
     assert np.all(tmodel_2000.predict([0, 0, 0]) == 1100)
 
 
-def test_gp_model():
-    gp = GaussianProcessModel("test")
-
+@pytest.mark.parametrize(
+    (
+        "bval_shell",
+        "S0",
+        "evals",
+    ),
+    [
+        (
+            1000,
+            100,
+            (0.0015, 0.0003, 0.0003),
+        )
+    ],
+)
+@pytest.mark.parametrize("snr", (10, 20))
+@pytest.mark.parametrize("hsph_dirs", (60, 30))
+def test_gp_model(evals, S0, snr, hsph_dirs, bval_shell):
+    # Simulate signal for a single tensor
+    evecs = _sim.create_single_fiber_evecs()
+    gtab = _sim.create_single_shell_gradient_table(hsph_dirs, bval_shell)
+    signal = single_tensor(gtab, S0=S0, evals=evals, evecs=evecs, snr=snr)
+
+    # Drop the initial b=0
+    gtab = gtab[1:]
+    data = signal[1:]
+
+    gp = GaussianProcessModel(kernel_model="spherical")
     assert isinstance(gp, model._dipy.GaussianProcessModel)
 
-    X, y = make_regression(n_samples=100, n_features=3, noise=0, random_state=0)
-
-    bvecs = X.T / np.linalg.norm(X.T, axis=0)
-    gtab = gradient_table([1000] * bvecs.shape[-1], bvecs)
-
-    gpfit = gp.fit(gtab, y)
-
-    X_qry = bvecs[:, :2].T
-    prediction = gpfit.predict(X_qry)
+    gpfit = gp.fit(data[:-2], gtab[:-2])
+    prediction = gpfit.predict(gtab.bvecs[-2:])
 
-    assert prediction.shape == (X_qry.shape[0],)
+    assert prediction.shape == (2,)
 
 
 def test_two_initialisations(datadir):