feat: sensitivity_from_table() -- post-hoc Sobol/Morris from a table (#113)

CHANGELOG.md

@@ -8,6 +8,7 @@ All notable changes to this project will be documented in this file.
 
 - 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).
+- `sensitivity_from_table()`: post-hoc Sobol/Morris sensitivity from an already-collected `ResultsTable`, by fitting a cheap surrogate over it (`fit_surrogate`) and running `screen()`'s existing machinery against the surrogate instead of a fresh simulator. Unlike a marginal Spearman correlation, this correctly detects non-monotonic (e.g. U-shaped) factor effects. Returns a `TableSensitivity` with `importance` indices and the surrogate's `surrogate_cv_r2` (#114) so callers can judge whether to trust the result (#113).
 
 ## [0.2.0] — 2026-08-17
 

src/trade_study/__init__.py

@@ -28,6 +28,7 @@
 )
 from .regime import RegimeSurrogate, fit_regime_surrogate
 from .runner import run_adaptive, run_grid, run_hyperband, run_successive_halving
+from .sensitivity import TableSensitivity, sensitivity_from_table
 from .stacking import ensemble_predict, stack_bayesian, stack_scores
 from .study import (
     Phase,
@@ -55,6 +56,7 @@
     "Simulator",
     "Study",
     "SurrogateModel",
+    "TableSensitivity",
     "TrialResult",
     "__version__",
     "build_grid",
@@ -80,6 +82,7 @@
     "save_results",
     "score",
     "screen",
+    "sensitivity_from_table",
     "stack_bayesian",
     "stack_scores",
     "top_k_pareto_filter",

src/trade_study/sensitivity.py

@@ -0,0 +1,128 @@
+"""Post-hoc sensitivity analysis from an already-collected ResultsTable (#113).
+
+``screen()`` (see :mod:`trade_study.design`) needs a live, re-runnable
+simulator and a proper Saltelli/Morris sample design -- neither of which an
+arbitrary already-collected :class:`~trade_study.protocols.ResultsTable`
+has, so its Sobol/Morris indices can't be computed retroactively from
+whatever points happen to be in the table.
+
+:func:`sensitivity_from_table` bridges the gap: it fits a cheap surrogate
+over the table (:func:`trade_study.fit_surrogate`) and runs ``screen()``'s
+existing Sobol/Morris machinery against the surrogate's ``predict()``
+instead of a fresh, expensive simulator evaluation. Because the surrogate
+is cheap to query, this produces genuine variance-based sensitivity
+indices -- which, unlike a marginal Spearman correlation, correctly detect
+non-monotonic effects -- without any new simulator evaluations.
+
+The result is only as trustworthy as the surrogate it comes from. Check
+:attr:`TableSensitivity.surrogate_cv_r2` (see #114) before trusting the
+indices for an observable with a poor cross-validated fit.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from .design import FactorType, screen
+from .surrogate import fit_surrogate
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+
+    from .design import Factor
+    from .protocols import ResultsTable
+
+
+@dataclass(frozen=True)
+class TableSensitivity:
+    """Post-hoc sensitivity indices computed via a table-fit surrogate.
+
+    Attributes:
+        importance: Mapping from observable name to an array of factor
+            importances (mu_star for Morris, S1 for Sobol), one value per
+            continuous factor, in the same order as
+            :func:`trade_study.screen` reports (i.e. continuous factors
+            only, in the order they appear in the ``factors`` argument).
+        surrogate_cv_r2: Per-observable cross-validated R^2 of the
+            surrogate the indices were computed from (#114). A low value
+            means the sensitivity indices reflect a poorly learned
+            response surface, not necessarily the true system -- treat
+            such observables' indices as unreliable.
+    """
+
+    importance: dict[str, NDArray[np.floating[Any]]]
+    surrogate_cv_r2: dict[str, float]
+
+
+def sensitivity_from_table(  # ruff: ignore[too-many-arguments]
+    results: ResultsTable,
+    factors: list[Factor],
+    *,
+    method: str = "sobol",
+    surrogate_method: str = "rf",
+    n_trajectories: int = 100,
+    seed: int = 42,
+    n_estimators: int = 200,
+    warn_below_r2: float | None = 0.0,
+) -> TableSensitivity:
+    """Compute post-hoc Sobol/Morris sensitivity from a collected table.
+
+    Only continuous factors are screened, matching ``screen()``'s own
+    contract; the surrogate is fit on that same continuous subset, so any
+    non-continuous keys present in ``results.configs`` (categorical
+    factors, bookkeeping fields, etc.) are simply ignored rather than
+    causing an encoding mismatch.
+
+    Args:
+        results: A :class:`~trade_study.protocols.ResultsTable` from a
+            previous ``run_grid``/``Study``/etc. call.
+        factors: Factor definitions to screen. Non-continuous factors are
+            dropped (as in ``screen()``); at least one continuous factor
+            must remain.
+        method: ``"sobol"`` or ``"morris"``, forwarded to ``screen()``.
+        surrogate_method: ``"rf"`` or ``"gp"``, forwarded to
+            :func:`trade_study.fit_surrogate`.
+        n_trajectories: Forwarded to ``screen()`` (Morris trajectory
+            count, or Sobol base sample size).
+        seed: Random seed for both the surrogate fit and ``screen()``.
+        n_estimators: Forwarded to :func:`trade_study.fit_surrogate`
+            (``rf`` only).
+        warn_below_r2: Forwarded to :func:`trade_study.fit_surrogate`;
+            warns if any observable's cross-validated R^2 is too low to
+            trust its sensitivity indices. Pass ``None`` to disable.
+
+    Returns:
+        A :class:`TableSensitivity` with importance indices and the
+        surrogate's cross-validated accuracy per observable.
+
+    Raises:
+        ValueError: If ``factors`` has no continuous entries, or
+            propagated from ``fit_surrogate()`` (e.g. an empty table).
+    """
+    continuous = [f for f in factors if f.factor_type == FactorType.CONTINUOUS]
+    if not continuous:
+        msg = "Screening requires at least one continuous factor"
+        raise ValueError(msg)
+
+    surrogate = fit_surrogate(
+        results,
+        continuous,
+        method=surrogate_method,
+        seed=seed,
+        n_estimators=n_estimators,
+        warn_below_r2=warn_below_r2,
+    )
+
+    def run_fn(cfg: dict[str, Any]) -> dict[str, float]:
+        return surrogate.predict(cfg)
+
+    importance = screen(
+        run_fn,
+        continuous,
+        method=method,
+        n_trajectories=n_trajectories,
+        seed=seed,
+    )
+    return TableSensitivity(importance=importance, surrogate_cv_r2=surrogate.cv_r2)

tests/test_sensitivity.py

@@ -0,0 +1,138 @@
+"""Tests for post-hoc sensitivity from an existing ResultsTable (#113)."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pytest
+
+from trade_study import (
+    Direction,
+    Factor,
+    FactorType,
+    Observable,
+    TableSensitivity,
+    build_grid,
+    run_grid,
+    sensitivity_from_table,
+)
+
+if TYPE_CHECKING:
+    from trade_study.protocols import ResultsTable
+
+
+class _World:
+    """Trivial simulator: passes config through."""
+
+    def generate(
+        self, config: dict[str, float]
+    ) -> tuple[dict[str, float], dict[str, float]]:
+        return config, config
+
+
+class _NonMonotonicScorer:
+    """Scorer where ``y = (a - 0.5)**2 + 0.1*b``.
+
+    ``a`` has a U-shaped (non-monotonic) effect symmetric about its
+    midpoint -- a marginal Spearman correlation would show it as
+    near-zero despite it being the dominant driver of variance. ``b`` has
+    a small, purely linear effect.
+    """
+
+    def score(
+        self,
+        truth: object,
+        observations: dict[str, float],
+        config: dict[str, float],
+    ) -> dict[str, float]:
+        del truth, config
+        a = float(observations["a"])
+        b = float(observations["b"])
+        return {"y": (a - 0.5) ** 2 + 0.1 * b}
+
+
+@pytest.fixture
+def continuous_factors() -> list[Factor]:
+    return [
+        Factor("a", FactorType.CONTINUOUS, bounds=(0.0, 1.0)),
+        Factor("b", FactorType.CONTINUOUS, bounds=(0.0, 1.0)),
+    ]
+
+
+def _make_results(factors: list[Factor], n: int = 128, seed: int = 0) -> ResultsTable:
+    grid = build_grid(factors, method="sobol", n_samples=n, seed=seed)
+    obs = [Observable("y", Direction.MINIMIZE)]
+    return run_grid(_World(), _NonMonotonicScorer(), grid, obs)
+
+
+def test_returns_table_sensitivity(continuous_factors: list[Factor]) -> None:
+    results = _make_results(continuous_factors)
+    result = sensitivity_from_table(results, continuous_factors, seed=0)
+    assert isinstance(result, TableSensitivity)
+    assert "y" in result.importance
+    assert "y" in result.surrogate_cv_r2
+
+
+def test_importance_shape(continuous_factors: list[Factor]) -> None:
+    results = _make_results(continuous_factors)
+    result = sensitivity_from_table(results, continuous_factors, seed=0)
+    assert result.importance["y"].shape == (2,)
+
+
+def test_sobol_detects_nonmonotonic_effect(continuous_factors: list[Factor]) -> None:
+    """Sobol S1 (via the surrogate) should rank U-shaped 'a' above linear 'b'.
+
+    This is the exact failure mode #113 was filed over: a marginal
+    Spearman correlation misses this because 'a's effect is symmetric
+    around its midpoint.
+    """
+    results = _make_results(continuous_factors, n=128)
+    result = sensitivity_from_table(
+        results, continuous_factors, method="sobol", surrogate_method="rf", seed=0
+    )
+    importance = result.importance["y"]
+    assert importance[0] > importance[1]
+    # 'a' should be clearly dominant, not just marginally ahead.
+    assert importance[0] > 0.5
+
+
+def test_morris_method_runs(continuous_factors: list[Factor]) -> None:
+    results = _make_results(continuous_factors)
+    result = sensitivity_from_table(
+        results, continuous_factors, method="morris", n_trajectories=20, seed=0
+    )
+    assert result.importance["y"].shape == (2,)
+
+
+def test_surrogate_cv_r2_reflects_fit_quality(continuous_factors: list[Factor]) -> None:
+    results = _make_results(continuous_factors, n=128)
+    result = sensitivity_from_table(results, continuous_factors, seed=0)
+    assert result.surrogate_cv_r2["y"] > 0.5
+
+
+def test_drops_noncontinuous_factors(continuous_factors: list[Factor]) -> None:
+    """A categorical factor in the input list is dropped, not an error."""
+    mixed = [*continuous_factors, Factor("kind", FactorType.CATEGORICAL, levels=["x"])]
+    results = _make_results(continuous_factors)
+    result = sensitivity_from_table(results, mixed, seed=0)
+    # Only the two continuous factors are screened.
+    assert result.importance["y"].shape == (2,)
+
+
+def test_propagates_screen_no_continuous_error(
+    continuous_factors: list[Factor],
+) -> None:
+    factors = [Factor("kind", FactorType.CATEGORICAL, levels=["a", "b"])]
+    results = _make_results(continuous_factors, n=8)
+    with pytest.raises(ValueError, match="at least one continuous"):
+        sensitivity_from_table(results, factors, seed=0)
+
+
+def test_warns_on_poor_surrogate_fit(continuous_factors: list[Factor]) -> None:
+    """A surrogate fit to pure noise warns via the forwarded fit_surrogate call."""
+    results = _make_results(continuous_factors, n=16)
+    rng = np.random.default_rng(0)
+    results.scores[:, 0] = rng.standard_normal(len(results.configs))
+    with pytest.warns(UserWarning, match="cross-validated R\\^2 below"):
+        sensitivity_from_table(results, continuous_factors, seed=0)