feat: cross-validated accuracy for surrogates (fit_surrogate/RegimeSurrogate)

CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to this project will be documented in this file.
 ### Added
 
 - Replicated trials: `run_grid(..., n_reps=N)` evaluates each design point N times; simulators may opt in to per-replicate randomness via an optional `rep` keyword on `Simulator.generate` (detected by introspection). `ResultsTable.aggregate_replicates()` collapses replicate rows back to per-design-point means with `n_reps`/`score_std` metadata. `Phase.n_reps` forwards this into `Study`, and phase filtering now runs against aggregated design points rather than raw replicates when `n_reps>1` (#112).
+- Surrogate accuracy reporting: `fit_surrogate()`/`SurrogateModel` (and `fit_regime_surrogate()`/`RegimeSurrogate` via passthrough) now compute a uniform k-fold cross-validated `cv_r2`/`cv_rmse` per observable for both the `gp` and `rf` backends. Warns (`warn_below_r2`, default threshold `0.0`) at fit time and in `RegimeSurrogate.recommend()` when an observable's accuracy is too low to trust (#114).
 
 ## [0.2.0] — 2026-08-17
 

src/trade_study/regime.py

@@ -28,6 +28,7 @@
 
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
@@ -83,6 +84,26 @@ class RegimeSurrogate:
     regime_factors: list[Factor]
     factors: list[Factor]
 
+    @property
+    def cv_r2(self) -> dict[str, float]:
+        """Per-observable held-out cross-validated R^2 (#114).
+
+        Returns:
+            Mapping from observable name to cross-validated R^2, from the
+            underlying :attr:`inner` surrogate.
+        """
+        return self.inner.cv_r2
+
+    @property
+    def cv_rmse(self) -> dict[str, float]:
+        """Per-observable held-out cross-validated RMSE (#114).
+
+        Returns:
+            Mapping from observable name to cross-validated RMSE, from
+            the underlying :attr:`inner` surrogate.
+        """
+        return self.inner.cv_rmse
+
     def predict(
         self,
         regime: dict[str, Any],
@@ -146,6 +167,7 @@ def recommend(
         n_candidates: int = 512,
         seed: int = 0,
         candidates: Sequence[dict[str, Any]] | None = None,
+        warn_below_r2: float | None = 0.0,
     ) -> dict[str, Any]:
         """Recommend a design-factor config at a query regime.
 
@@ -163,6 +185,11 @@ def recommend(
             seed: Seed for the Sobol' sampler.
             candidates: Optional explicit list of design-factor configs
                 to score; if given, overrides ``n_candidates``.
+            warn_below_r2: Warn if ``objective``'s cross-validated R^2
+                (#114) is below this threshold, so a caller optimizing
+                against a poorly-fit surrogate gets a signal right at the
+                point of use, not just buried in fit-time logs. Pass
+                ``None`` to disable.
 
         Returns:
             The candidate config (a copy) achieving the best predicted
@@ -182,6 +209,15 @@ def recommend(
                 f"available: {self.inner.observable_names}"
             )
             raise ValueError(msg)
+        r2 = self.cv_r2.get(objective, float("nan"))
+        if warn_below_r2 is not None and np.isfinite(r2) and r2 < warn_below_r2:
+            warnings.warn(
+                f"recommend: objective {objective!r} has cross-validated "
+                f"R^2={r2:.3f} (< {warn_below_r2}); this recommendation may "
+                f"not be trustworthy.",
+                UserWarning,
+                stacklevel=2,
+            )
         pool = (
             list(candidates)
             if candidates is not None
@@ -200,14 +236,16 @@ def recommend(
         return dict(pool[idx])
 
 
-def fit_regime_surrogate(
+def fit_regime_surrogate(  # ruff: ignore[too-many-arguments]
     results: ResultsTable,
     regime_factors: list[Factor],
     factors: list[Factor],
     *,
     method: str = "gp",
     seed: int = 0,
     n_estimators: int = 200,
+    cv_folds: int = 5,
+    warn_below_r2: float | None = 0.0,
 ) -> RegimeSurrogate:
     """Fit a surrogate that conditions on regime features.
 
@@ -230,6 +268,10 @@ def fit_regime_surrogate(
             :func:`trade_study.fit_surrogate`.
         seed: Random seed forwarded to the backend estimators.
         n_estimators: Number of trees for the ``"rf"`` backend.
+        cv_folds: Cross-validation folds for the held-out accuracy check
+            (#114). See :func:`trade_study.fit_surrogate`.
+        warn_below_r2: Warn if any observable's cross-validated R^2 falls
+            below this threshold. See :func:`trade_study.fit_surrogate`.
 
     Returns:
         A fitted :class:`RegimeSurrogate`.
@@ -255,6 +297,8 @@ def fit_regime_surrogate(
         method=method,
         seed=seed,
         n_estimators=n_estimators,
+        cv_folds=cv_folds,
+        warn_below_r2=warn_below_r2,
     )
     return RegimeSurrogate(
         inner=inner,

src/trade_study/surrogate.py

@@ -19,6 +19,7 @@
 
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
@@ -128,12 +129,19 @@ class SurrogateModel:
         encoder: Factor encoder used at fit time.
         observable_names: Column names of the predicted observables.
         models: One fitted scikit-learn estimator per observable.
+        cv_r2: Per-observable held-out cross-validated R^2 (#114). Missing
+            for an observable if too few rows were available to run CV
+            (fewer than 2 folds); ``float("nan")`` in that case.
+        cv_rmse: Per-observable held-out cross-validated RMSE, companion
+            to ``cv_r2`` in the observable's original units.
     """
 
     method: str
     encoder: _FactorEncoder
     observable_names: list[str]
     models: list[Any]
+    cv_r2: dict[str, float] = field(default_factory=dict)
+    cv_rmse: dict[str, float] = field(default_factory=dict)
 
     def predict(self, config: dict[str, Any]) -> dict[str, float]:
         """Predict observables for a single config.
@@ -203,13 +211,23 @@ def fit_surrogate(
     method: str = "gp",
     seed: int = 0,
     n_estimators: int = 200,
+    cv_folds: int = 5,
+    warn_below_r2: float | None = 0.0,
 ) -> SurrogateModel:
     """Fit a per-observable surrogate over a :class:`ResultsTable`.
 
     Rows whose score column contains ``NaN`` are dropped on a
     per-observable basis (so a partially-evaluated trial still
     contributes to the observables it does have).
 
+    After fitting each observable's model on all its available rows,
+    also computes a held-out cross-validated R^2/RMSE (#114) so callers
+    can tell a well-fit surrogate from one that's effectively guessing
+    in a sparse or noisy region of the design space -- neither backend
+    reports this on its own (RF's cheap OOB score isn't available for
+    GP, so CV is used uniformly for both, at the cost of ``cv_folds``
+    extra fits per observable).
+
     Args:
         results: A :class:`ResultsTable` from a previous study run.
         factors: Factor definitions used to encode ``results.configs``.
@@ -219,6 +237,14 @@ def fit_surrogate(
         seed: Random seed forwarded to the backend estimators.
         n_estimators: Number of trees for the ``"rf"`` backend; ignored
             for ``"gp"``.
+        cv_folds: Number of cross-validation folds used to compute
+            ``cv_r2``/``cv_rmse``. Clamped down to the observable's
+            available row count when fewer than ``cv_folds`` rows exist;
+            skipped (``nan``) for an observable with fewer than 2 rows.
+        warn_below_r2: If not ``None``, emit a ``UserWarning`` naming any
+            observable whose ``cv_r2`` falls below this threshold --
+            0.0 (the default) flags a surrogate that predicts no better
+            than the training mean. Pass ``None`` to disable.
 
     Returns:
         A fitted :class:`SurrogateModel`.
@@ -242,31 +268,94 @@ def fit_surrogate(
 
     models: list[Any] = []
     fitted_obs: list[str] = []
+    cv_r2: dict[str, float] = {}
+    cv_rmse: dict[str, float] = {}
+    low_accuracy: list[str] = []
     for j, name in enumerate(results.observable_names):
         y = results.scores[:, j]
         mask = ~np.isnan(y)
-        if int(mask.sum()) < 2:
+        n_rows = int(mask.sum())
+        if n_rows < 2:
             continue
         model = _make_estimator(method, seed=seed, n_estimators=n_estimators)
         model.fit(x_full[mask], y[mask])
         models.append(model)
         fitted_obs.append(name)
 
+        r2, rmse = _cross_val_accuracy(
+            x_full[mask],
+            y[mask],
+            method=method,
+            seed=seed,
+            n_estimators=n_estimators,
+            n_folds=min(cv_folds, n_rows),
+        )
+        cv_r2[name] = r2
+        cv_rmse[name] = rmse
+        if warn_below_r2 is not None and np.isfinite(r2) and r2 < warn_below_r2:
+            low_accuracy.append(name)
+
     if not models:
         msg = (
             "fit_surrogate: no observable has at least 2 non-NaN training "
             "rows; nothing to fit"
         )
         raise ValueError(msg)
 
+    if low_accuracy:
+        warnings.warn(
+            f"fit_surrogate: cross-validated R^2 below {warn_below_r2} for "
+            f"{low_accuracy} -- predictions for these observables may not "
+            f"be trustworthy. See SurrogateModel.cv_r2 for exact values.",
+            UserWarning,
+            stacklevel=2,
+        )
+
     return SurrogateModel(
         method=method,
         encoder=encoder,
         observable_names=fitted_obs,
         models=models,
+        cv_r2=cv_r2,
+        cv_rmse=cv_rmse,
     )
 
 
+def _cross_val_accuracy(
+    x: NDArray[np.floating[Any]],
+    y: NDArray[np.floating[Any]],
+    *,
+    method: str,
+    seed: int,
+    n_estimators: int,
+    n_folds: int,
+) -> tuple[float, float]:
+    """K-fold cross-validated R^2/RMSE for one observable.
+
+    Returns:
+        ``(r2, rmse)``, both ``float("nan")`` if fewer than 2 folds are
+        possible.
+    """
+    if n_folds < 2:
+        return float("nan"), float("nan")
+
+    from sklearn.model_selection import KFold  # type: ignore[import-untyped]
+
+    kfold = KFold(n_splits=n_folds, shuffle=True, random_state=seed)
+    y_pred = np.empty_like(y)
+    for train_idx, test_idx in kfold.split(x):
+        fold_model = _make_estimator(method, seed=seed, n_estimators=n_estimators)
+        fold_model.fit(x[train_idx], y[train_idx])
+        y_pred[test_idx] = fold_model.predict(x[test_idx])
+
+    residuals = y - y_pred
+    rmse = float(np.sqrt(np.mean(residuals**2)))
+    ss_res = float(np.sum(residuals**2))
+    ss_tot = float(np.sum((y - y.mean()) ** 2))
+    r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan")
+    return r2, rmse
+
+
 def _make_estimator(method: str, *, seed: int, n_estimators: int) -> Any:  # ruff: ignore[any-type]
     """Construct an unfitted scikit-learn estimator for the requested method.
 

tests/test_regime.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -100,6 +101,68 @@ def test_fit_rejects_overlapping_names(
         fit_regime_surrogate(results, [regime_factor], [dup])
 
 
+# ---------------------------------------------------------------------------
+# Cross-validated accuracy passthrough (#114)
+# ---------------------------------------------------------------------------
+
+
+def test_cv_r2_rmse_passthrough(
+    regime_factor: Factor,
+    design_factor: Factor,
+) -> None:
+    results = _make_results(regime_factor, design_factor, n=64)
+    sur = fit_regime_surrogate(
+        results,
+        [regime_factor],
+        [design_factor],
+        method="rf",
+        seed=0,
+    )
+    assert sur.cv_r2 == sur.inner.cv_r2
+    assert sur.cv_rmse == sur.inner.cv_rmse
+    assert "loss" in sur.cv_r2
+    assert "loss" in sur.cv_rmse
+
+
+def test_recommend_warns_on_poor_fit(
+    regime_factor: Factor,
+    design_factor: Factor,
+) -> None:
+    results = _make_results(regime_factor, design_factor, n=16)
+    rng = np.random.default_rng(0)
+    results.scores[:, 0] = rng.standard_normal(len(results.configs))
+    sur = fit_regime_surrogate(
+        results,
+        [regime_factor],
+        [design_factor],
+        method="rf",
+        seed=0,
+        warn_below_r2=None,  # suppress the fit-time warning to isolate recommend()'s
+    )
+    with pytest.warns(UserWarning, match="cross-validated R\\^2"):
+        sur.recommend({"n": 5.0}, objective="loss")
+
+
+def test_recommend_warn_below_r2_none_disables(
+    regime_factor: Factor,
+    design_factor: Factor,
+) -> None:
+    results = _make_results(regime_factor, design_factor, n=16)
+    rng = np.random.default_rng(0)
+    results.scores[:, 0] = rng.standard_normal(len(results.configs))
+    sur = fit_regime_surrogate(
+        results,
+        [regime_factor],
+        [design_factor],
+        method="rf",
+        seed=0,
+        warn_below_r2=None,
+    )
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        sur.recommend({"n": 5.0}, objective="loss", warn_below_r2=None)
+
+
 # ---------------------------------------------------------------------------
 # Predict / uncertainty
 # ---------------------------------------------------------------------------