Merge pull request #264 from alan-turing-institute/plotting

fix wrong R2's in plot_cv due to not subsetting training data

Author: mastoffel

autoemulate/compare.py

@@ -422,8 +422,8 @@ def plot_cv(
         )
         figure = _plot_cv(
             self.cv_results,
-            self.X,
-            self.y,
+            self.X[self.train_idxs],
+            self.y[self.train_idxs],
             model_name=model_name,
             n_cols=n_cols,
             style=style,

autoemulate/plotting.py

@@ -3,6 +3,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from sklearn.metrics import PredictionErrorDisplay
+from sklearn.metrics import r2_score
 from sklearn.pipeline import Pipeline
 
 from autoemulate.utils import _ensure_2d
@@ -482,6 +483,7 @@ def _plot_Xy(X, y, y_pred, y_std=None, ax=None, title="Xy"):
     """
     Plots observed and predicted values vs. features, including 2σ error bands where available.
     """
+
     # Sort the data
     sort_idx = np.argsort(X).flatten()
     X_sorted = X[sort_idx]
@@ -558,10 +560,10 @@ def _plot_Xy(X, y, y_pred, y_std=None, ax=None, title="Xy"):
             columnspacing=0,
             ncol=2,
         )
+
     # Calculate R2 score
-    r2 = 1 - np.sum((y_sorted - y_pred_sorted) ** 2) / np.sum(
-        (y_sorted - np.mean(y_sorted)) ** 2
-    )
+    r2 = r2_score(y, y_pred)
+
     ax.text(
         0.05,
         0.05,