From ea93c4ade96e0ce8617dbceb77e03cc5a69ea25e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?=
 <jon.haitz.legarreta@gmail.com>
Date: Sun, 29 Sep 2024 13:43:20 -0400
Subject: [PATCH] WIP: Use scikit-learn GP returned std error value

Use scikit-learn GP returned std error value.
---
 scripts/dwi_estimation_error_analysis.py | 25 ++++++++++++------------
 src/eddymotion/model/_dipy.py            |  2 +-
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/scripts/dwi_estimation_error_analysis.py b/scripts/dwi_estimation_error_analysis.py
index 530394ce..8b74dfea 100644
--- a/scripts/dwi_estimation_error_analysis.py
+++ b/scripts/dwi_estimation_error_analysis.py
@@ -164,16 +164,16 @@ def perform_experiment(gtab, S0, evals1, evecs, snr, repeats, kfold):
 
         # Simulate the fitting a number of times: every time the signal created will be a little
         # different
-        for _ in range(repeats):
-            # Create the DWI signal using a single tensor
-            signal = single_tensor(gtab, S0=S0, evals=evals1, evecs=evecs, snr=snr, rng=rng)
-            # Fit the Gaussian process
-            gpfit = gp_model.fit(signal[train_mask], gtab[train_mask])
+        # for _ in range(repeats):
+        # Create the DWI signal using a single tensor
+        signal = single_tensor(gtab, S0=S0, evals=evals1, evecs=evecs, snr=snr, rng=rng)
+        # Fit the Gaussian process
+        gpfit = gp_model.fit(signal[train_mask], gtab[train_mask])
 
-            # Predict the signal
-            X_qry, idx_qry = get_query_vectors(gtab, train_mask)
-            _y_pred = gpfit.predict(X_qry)
-            data.append((idx_qry, signal[idx_qry], _y_pred))
+        # Predict the signal
+        X_qry, idx_qry = get_query_vectors(gtab, train_mask)
+        _y_pred, _y_std = gpfit.predict(X_qry)
+        data.append((idx_qry, signal[idx_qry], _y_pred, _y_std))
 
     return data
 
@@ -186,10 +186,11 @@ def compute_error(data, repeats, kfolds):
 
     # Loop over the range of indices that were predicted
     for n in range(len(kfolds)):
+        repeats = 1
         _data = np.array(data[n * repeats : n * repeats + repeats])
-        _rmse = list(map(root_mean_squared_error, _data[:, 1], _data[:, 2]))
-        _std_dev = np.std(_rmse)
-        mean_rmse.append(np.mean(_rmse))
+        _rmse = root_mean_squared_error(_data[0][1], _data[0][2])
+        _std_dev = np.mean(_data[0][3])  #  np.std(_rmse)
+        mean_rmse.append(_rmse)
         std_dev.append(_std_dev)
 
     return np.asarray(mean_rmse), np.asarray(std_dev)
diff --git a/src/eddymotion/model/_dipy.py b/src/eddymotion/model/_dipy.py
index 7a97420e..fb91f330 100644
--- a/src/eddymotion/model/_dipy.py
+++ b/src/eddymotion/model/_dipy.py
@@ -138,7 +138,7 @@ def gp_prediction(
         raise RuntimeError("Model is not yet fitted.")
 
     # Extract orientations from gtab, and highly likely, the b-value too.
-    return model.predict(gtab, return_std=False)
+    return model.predict(gtab, return_std=True)
 
 
 class GaussianProcessModel(ReconstModel):