feat: sobol_indices() -- expose total-order Sobol indices (#120)

CHANGELOG.md

@@ -9,6 +9,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).
+- `sobol_indices()`: like `screen(method="sobol")`, but returns both S1 (first-order) and ST (total-order) per observable instead of discarding ST. `screen()` itself is unchanged for backward compatibility; `ST - S1` is the standard way to detect interaction effects that a first-order-only view misses entirely (#120).
 
 ## [0.2.0] — 2026-08-17
 

src/trade_study/__init__.py

@@ -13,6 +13,7 @@
     build_grid,
     reduce_factors,
     screen,
+    sobol_indices,
 )
 from .io import load_results, save_results
 from .protocols import (
@@ -83,6 +84,7 @@
     "score",
     "screen",
     "sensitivity_from_table",
+    "sobol_indices",
     "stack_bayesian",
     "stack_scores",
     "top_k_pareto_filter",

src/trade_study/design.py

@@ -501,18 +501,21 @@ def _screen_morris(
     return importance
 
 
-def _screen_sobol(
+def _sobol_sample_and_evaluate(
     run_fn: Callable[[dict[str, Any]], dict[str, float]],
     problem: dict[str, Any],
     n_samples: int,
     seed: int,
-) -> dict[str, NDArray[np.floating[Any]]]:
-    """Sobol variance-based sensitivity analysis.
+) -> dict[str, list[float]]:
+    """Draw a Saltelli design and evaluate ``run_fn`` at every point.
+
+    Shared by ``_screen_sobol`` (S1 only) and :func:`sobol_indices` (S1 and
+    ST) so both draw from -- and evaluate -- the exact same design.
 
     Returns:
-        Mapping from observable name to S1 (first-order) index array.
+        Mapping from observable name to its list of scores, one per
+        sampled design point, in sample order.
     """
-    from SALib.analyze import sobol as sobol_analyze
     from SALib.sample import sobol as sobol_sample
 
     param_values = sobol_sample.sample(problem, n_samples, seed=seed)
@@ -523,6 +526,23 @@ def _screen_sobol(
         scores = run_fn(cfg)
         for obs_name, val in scores.items():
             results_by_obs.setdefault(obs_name, []).append(val)
+    return results_by_obs
+
+
+def _screen_sobol(
+    run_fn: Callable[[dict[str, Any]], dict[str, float]],
+    problem: dict[str, Any],
+    n_samples: int,
+    seed: int,
+) -> dict[str, NDArray[np.floating[Any]]]:
+    """Sobol variance-based sensitivity analysis.
+
+    Returns:
+        Mapping from observable name to S1 (first-order) index array.
+    """
+    from SALib.analyze import sobol as sobol_analyze
+
+    results_by_obs = _sobol_sample_and_evaluate(run_fn, problem, n_samples, seed)
 
     importance: dict[str, NDArray[np.floating[Any]]] = {}
     for obs_name, vals in results_by_obs.items():
@@ -536,6 +556,63 @@ def _screen_sobol(
     return importance
 
 
+def sobol_indices(
+    run_fn: Callable[[dict[str, Any]], dict[str, float]],
+    factors: list[Factor],
+    *,
+    n_samples: int = 100,
+    seed: int = 42,
+) -> dict[str, tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]]]]:
+    """Sobol first- and total-order sensitivity indices (#120).
+
+    Like ``screen(method="sobol")``, but returns both S1 (first-order) and
+    ST (total-order) per observable instead of discarding ST. ``ST - S1``
+    is the standard way to detect interaction effects: a factor with small
+    S1 but large ST is interacting with other factors rather than acting
+    independently -- information ``screen()`` can't surface at all, since
+    it only keeps S1.
+
+    Args:
+        run_fn: Callable that takes a config dict and returns a dict of
+            observable name -> scalar score.
+        factors: Factor list to analyze; only continuous factors are
+            varied (as in ``screen()``).
+        n_samples: Base sample size *N* for the Saltelli design; total
+            evaluations are *N* x (2 x num_vars + 2).
+        seed: Random seed.
+
+    Returns:
+        Dictionary mapping observable name to an ``(S1, ST)`` tuple, each
+        an array of one value per continuous factor (in the order they
+        appear in ``factors``).
+
+    Raises:
+        ValueError: If no continuous factors are provided.
+    """
+    from SALib.analyze import sobol as sobol_analyze
+
+    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)
+
+    problem: dict[str, Any] = {
+        "num_vars": len(continuous),
+        "names": [f.name for f in continuous],
+        "bounds": [list(f.bounds) for f in continuous if f.bounds is not None],
+    }
+    results_by_obs = _sobol_sample_and_evaluate(run_fn, problem, n_samples, seed)
+
+    indices: dict[str, tuple[NDArray[np.floating[Any]], NDArray[np.floating[Any]]]] = {}
+    for obs_name, vals in results_by_obs.items():
+        si = sobol_analyze.analyze(problem, np.array(vals), seed=seed)
+        s1 = np.asarray(si["S1"], dtype=np.float64)
+        st = np.asarray(si["ST"], dtype=np.float64)
+        indices[obs_name] = (s1, st)
+
+    return indices
+
+
 def reduce_factors(
     factors: list[Factor],
     importance: dict[str, NDArray[np.floating[Any]]],

tests/test_design.py

@@ -14,6 +14,7 @@
     build_grid,
     reduce_factors,
     screen,
+    sobol_indices,
 )
 
 # ---------------------------------------------------------------------------
@@ -397,6 +398,72 @@ def multi_obs(cfg: dict[str, Any]) -> dict[str, float]:
     assert result["obs2"].shape == (2,)
 
 
+# ---------------------------------------------------------------------------
+# sobol_indices (#120)
+# ---------------------------------------------------------------------------
+
+
+def test_sobol_indices_returns_tuple_per_observable(
+    continuous_factors: list[Factor],
+) -> None:
+    result = sobol_indices(_linear_model, continuous_factors, n_samples=64, seed=0)
+    assert isinstance(result, dict)
+    assert "y" in result
+    s1, st = result["y"]
+    assert s1.shape == (2,)
+    assert st.shape == (2,)
+
+
+def test_sobol_indices_matches_screen_s1(continuous_factors: list[Factor]) -> None:
+    """S1 from sobol_indices should match screen(method="sobol")'s S1."""
+    screened = screen(
+        _linear_model, continuous_factors, method="sobol", n_trajectories=64, seed=0
+    )
+    result = sobol_indices(_linear_model, continuous_factors, n_samples=64, seed=0)
+    s1, _st = result["y"]
+    np.testing.assert_allclose(s1, screened["y"])
+
+
+def test_sobol_indices_total_order_at_least_first_order(
+    continuous_factors: list[Factor],
+) -> None:
+    """ST >= S1 holds (up to MC estimation noise) for a variance decomposition."""
+    result = sobol_indices(_linear_model, continuous_factors, n_samples=512, seed=0)
+    s1, st = result["y"]
+    # SALib's S1/ST are independently-estimated MC quantities, not computed
+    # from a shared exact decomposition, so a generous tolerance is needed
+    # even at a reasonable sample size -- this only guards against a
+    # systematic ST-vs-S1 mixup (e.g. an accidental swap), not tight
+    # numerical agreement.
+    assert np.all(st >= s1 - 0.05)
+
+
+def test_sobol_indices_detects_pure_interaction(
+    continuous_factors: list[Factor],
+) -> None:
+    """A pure product of two zero-mean factors has ~zero S1 but nonzero ST."""
+
+    def interaction_model(cfg: dict[str, Any]) -> dict[str, float]:
+        # Both factors centered to zero mean: a pure product term then has
+        # zero first-order effect for *both* factors (E[Y|A] = A*E[B] = 0
+        # and vice versa), so any measured S1 is interaction leaking in,
+        # while ST captures the interaction directly.
+        alpha_centered = (cfg["alpha"] - 0.5) / 0.5  # [0, 1] -> [-1, 1]
+        beta_centered = (cfg["beta"] - 15.0) / 5.0  # [10, 20] -> [-1, 1]
+        return {"y": alpha_centered * beta_centered}
+
+    result = sobol_indices(interaction_model, continuous_factors, n_samples=256, seed=0)
+    s1, st = result["y"]
+    assert np.all(s1 < 0.1)
+    assert np.all(st > 0.1)
+
+
+def test_sobol_indices_rejects_no_continuous() -> None:
+    factors = [Factor("m", FactorType.CATEGORICAL, levels=["a", "b"])]
+    with pytest.raises(ValueError, match="at least one continuous"):
+        sobol_indices(lambda _c: {"y": 0.0}, factors)
+
+
 # ---------------------------------------------------------------------------
 # reduce_factors (#10)
 # ---------------------------------------------------------------------------