diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cc6036..376275c 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/trade_study/__init__.py b/src/trade_study/__init__.py index 1f48f5d..9ad3b3a 100644 --- a/src/trade_study/__init__.py +++ b/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", diff --git a/src/trade_study/design.py b/src/trade_study/design.py index d7ed354..5db69e4 100644 --- a/src/trade_study/design.py +++ b/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]]], diff --git a/tests/test_design.py b/tests/test_design.py index 304f0f8..58d7c17 100644 --- a/tests/test_design.py +++ b/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) # ---------------------------------------------------------------------------