From 5acbf732f86c34bfe663e1ff48f2a5651ebe4c5b Mon Sep 17 00:00:00 2001 From: Joshua Date: Thu, 20 Aug 2026 11:57:48 -0400 Subject: [PATCH] feat: recommend_bucketed_config() -- discrete regime-bucketed adaptive recommendation (#123) RegimeSurrogate interpolates continuously across regime descriptors, but needs reasonably dense training coverage -- RF surrogates extrapolate poorly (effectively flat) outside their training range. There was no equivalent for the sparse case: a handful of named regimes, each far from the others and far from any existing training data, where you just want "find a good config for each one directly, then aggregate into named buckets." Found deriving VBPCA hyperparameter buckets for extreme-aspect-ratio data shapes (yoavram-lab/VBPCApy#116) -- the pattern that actually worked (run_adaptive per regime + median aggregation into buckets) was hand-rolled there; this formalizes it as a reusable trade-study function. recommend_bucketed_config() runs run_adaptive per regime (respecting n_reps, #122), selects each regime's best trial by a primary observable (respecting its Direction), groups regimes via a caller-supplied bucket_fn, and aggregates each bucket's per-regime best configs -- median for continuous/discrete factors, mode for categorical. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 1 + src/trade_study/__init__.py | 3 +- src/trade_study/regime.py | 122 ++++++++++++++++++++++- tests/test_regime.py | 194 ++++++++++++++++++++++++++++++++++++ 4 files changed, 316 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7fd7abb..282cb5c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ All notable changes to this project will be documented in this file. - `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). - Replicate averaging in `run_adaptive()` and `screen()`/`sobol_indices()` (`n_reps`, #122): the same `run_grid(..., n_reps=N)` convention (#112), now applied consistently across every entry point that repeatedly evaluates a simulator/`run_fn`. `run_adaptive` detects an opt-in `rep` keyword on `Simulator.generate` and averages each trial's objective(s) over `n_reps` draws before Optuna sees them; `screen()`/`sobol_indices()` detect the same convention on the bare `run_fn` callable they take instead. Without this, adaptive optimization or sensitivity screening against a stochastic simulator could select a "best" config or report an "important" factor that's actually just a single lucky/unlucky data draw. Default `n_reps=1` preserves prior behavior. +- `recommend_bucketed_config()`: the discrete counterpart to `RegimeSurrogate` (#123). For a handful of named regimes too sparse (and too far outside any existing training data) for a surrogate to extrapolate across sensibly, runs `run_adaptive` independently per regime, picks each regime's best-found config by a primary objective, groups regimes into named buckets, and aggregates each bucket's per-regime best configs (median for continuous/discrete factors, mode for categorical) into one recommended config per bucket. ## [0.2.0] — 2026-08-17 diff --git a/src/trade_study/__init__.py b/src/trade_study/__init__.py index 9ad3b3a..cba0a95 100644 --- a/src/trade_study/__init__.py +++ b/src/trade_study/__init__.py @@ -27,7 +27,7 @@ Simulator, TrialResult, ) -from .regime import RegimeSurrogate, fit_regime_surrogate +from .regime import RegimeSurrogate, fit_regime_surrogate, recommend_bucketed_config 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 @@ -75,6 +75,7 @@ "plot_front", "plot_parallel", "plot_scores", + "recommend_bucketed_config", "reduce_factors", "run_adaptive", "run_grid", diff --git a/src/trade_study/regime.py b/src/trade_study/regime.py index d5bea98..d5c5795 100644 --- a/src/trade_study/regime.py +++ b/src/trade_study/regime.py @@ -29,20 +29,23 @@ from __future__ import annotations import warnings +from collections import Counter, defaultdict from dataclasses import dataclass from typing import TYPE_CHECKING, Any import numpy as np -from .design import Factor, build_grid +from .design import Factor, FactorType, build_grid +from .protocols import Direction +from .runner import run_adaptive from .surrogate import SurrogateModel, fit_surrogate if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence from numpy.typing import NDArray - from .protocols import ResultsTable + from .protocols import Observable, ResultsTable, Scorer, Simulator _SUPPORTED_MODES: frozenset[str] = frozenset({"min", "max"}) @@ -305,3 +308,116 @@ def fit_regime_surrogate( # ruff: ignore[too-many-arguments] regime_factors=list(regime_factors), factors=list(factors), ) + + +def _aggregate_factor_values(factor: Factor, values: list[Any]) -> Any: # ruff: ignore[any-type] + """Aggregate one factor's per-regime best values into a bucket value. + + Continuous/discrete (numeric) factors are aggregated by median; + categorical factors by mode (most common value, ties broken by + first occurrence). + + Returns: + The aggregated value for this factor. + """ + if factor.factor_type == FactorType.CATEGORICAL: + return Counter(values).most_common(1)[0][0] + return type(values[0])(np.median(values)) + + +def recommend_bucketed_config( # ruff: ignore[too-many-arguments] + regimes: dict[str, dict[str, Any]], + bucket_fn: Callable[[str, dict[str, Any]], str], + world_factory: Callable[[dict[str, Any]], Simulator], + scorer: Scorer, + factors: list[Factor], + observables: list[Observable], + *, + primary: str, + n_trials: int = 30, + n_reps: int = 1, + seed: int = 42, +) -> dict[str, dict[str, Any]]: + """Recommend a config per named bucket via per-regime adaptive search (#123). + + The discrete counterpart to :func:`fit_regime_surrogate`: for a + handful of named regimes too sparse (and too far outside any + existing training data) for a surrogate to extrapolate across + sensibly, this instead runs :func:`~trade_study.run_adaptive` (NSGA-II) + independently per regime, picks each regime's best-found config by + ``primary``, groups regimes into named buckets via ``bucket_fn``, and + aggregates each bucket's per-regime best configs (median for + continuous/discrete factors, mode for categorical) into one + recommended config per bucket. + + Args: + regimes: Mapping from regime name to a regime descriptor dict + (whatever ``world_factory`` needs to fix that regime). + bucket_fn: Maps ``(regime_name, regime_dict)`` to a bucket name; + regimes sharing a bucket name have their best configs + aggregated together. + world_factory: Builds a regime-scoped :class:`Simulator` from a + regime dict, e.g. ``lambda r: MySimulator(regime_defaults=r)``. + scorer: Scorer for observables (shared across all regimes). + factors: Tunable factors searched by ``run_adaptive`` at each + regime (the regime itself is fixed via ``world_factory``, not + part of this search space). + observables: Observable definitions passed to ``run_adaptive``. + primary: Name of the observable used to pick each regime's single + best trial (the one minimizing/maximizing it, per that + observable's ``direction``) from ``run_adaptive``'s Pareto set. + n_trials: Optuna trials per regime. + n_reps: Replicate draws averaged per trial (#122) -- see + :func:`~trade_study.run_adaptive`'s ``n_reps`` for the full + rationale. Default 1 (a single draw per trial). + seed: Random seed forwarded to each regime's ``run_adaptive`` call. + + Returns: + Mapping from bucket name to its aggregated recommended config. + + Raises: + ValueError: If ``regimes`` is empty, or ``primary`` doesn't match + any name in ``observables``. + """ + if not regimes: + msg = "recommend_bucketed_config: regimes must be non-empty" + raise ValueError(msg) + matching = [o for o in observables if o.name == primary] + if not matching: + msg = f"primary={primary!r} not found in observables" + raise ValueError(msg) + minimize = matching[0].direction == Direction.MINIMIZE + + per_regime_best: dict[str, dict[str, Any]] = {} + bucket_members: dict[str, list[str]] = defaultdict(list) + for name, regime in regimes.items(): + world = world_factory(regime) + table = run_adaptive( + world, + scorer, + factors, + observables, + n_trials=n_trials, + n_reps=n_reps, + seed=seed, + ) + primary_col = table.observable_names.index(primary) + best_i = ( + int(np.argmin(table.scores[:, primary_col])) + if minimize + else int(np.argmax(table.scores[:, primary_col])) + ) + per_regime_best[name] = table.configs[best_i] + bucket_members[bucket_fn(name, regime)].append(name) + + factors_by_name = {f.name: f for f in factors} + return { + bucket: { + key: _aggregate_factor_values( + factors_by_name[key], + [per_regime_best[member][key] for member in members], + ) + for key in factors_by_name + } + for bucket, members in bucket_members.items() + } diff --git a/tests/test_regime.py b/tests/test_regime.py index c8199df..0f3f015 100644 --- a/tests/test_regime.py +++ b/tests/test_regime.py @@ -16,6 +16,7 @@ RegimeSurrogate, build_grid, fit_regime_surrogate, + recommend_bucketed_config, run_grid, ) @@ -327,3 +328,196 @@ def test_recommend_rejects_empty_candidates( ) with pytest.raises(ValueError, match="no candidates"): sur.recommend({"n": 5.0}, objective="loss", candidates=[]) + + +# --------------------------------------------------------------------------- +# recommend_bucketed_config (#123) +# --------------------------------------------------------------------------- + + +class _TargetWorld: + """Simulator fixed to one regime's target alpha via regime_defaults.""" + + def __init__(self, *, target: float, method: str) -> None: + self._target = target + self._method = method + + def generate( + self, config: dict[str, object] + ) -> tuple[dict[str, object], dict[str, object]]: + merged = {**config, "target": self._target, "true_method": self._method} + return merged, merged + + +class _TargetScorer: + """cost = (alpha - target)**2; bonus reward when method matches target's.""" + + def score( + self, + truth: object, + observations: dict[str, object], + config: dict[str, object], + ) -> dict[str, float]: + del truth, config + alpha = float(observations["alpha"]) + target = float(observations["target"]) + method_ok = observations["method"] == observations["true_method"] + return { + "cost": (alpha - target) ** 2, + "reward": 1.0 if method_ok else 0.0, + } + + +@pytest.fixture +def bucketed_factors() -> list[Factor]: + """One continuous and one categorical tunable factor. + + Returns: + List with ``alpha`` (continuous) and ``method`` (categorical). + """ + return [ + Factor("alpha", FactorType.CONTINUOUS, bounds=(0.0, 1.0)), + Factor("method", FactorType.CATEGORICAL, levels=["a", "b"]), + ] + + +@pytest.fixture +def bucketed_observables() -> list[Observable]: + """Cost (minimize) and reward (maximize) observables. + + Returns: + List of two Observable instances. + """ + return [ + Observable("cost", Direction.MINIMIZE), + Observable("reward", Direction.MAXIMIZE), + ] + + +def test_recommend_bucketed_config_aggregates_median( + bucketed_factors: list[Factor], + bucketed_observables: list[Observable], +) -> None: + regimes = { + "r1": {"target": 0.2, "method": "a"}, + "r2": {"target": 0.8, "method": "a"}, + } + result = recommend_bucketed_config( + regimes, + bucket_fn=lambda _name, _r: "only", + world_factory=lambda r: _TargetWorld(target=r["target"], method=r["method"]), + scorer=_TargetScorer(), + factors=bucketed_factors, + observables=bucketed_observables, + primary="cost", + n_trials=40, + seed=0, + ) + assert set(result.keys()) == {"only"} + # each regime's best alpha should land near its own target; median of + # two well-optimized targets (0.2, 0.8) should land near their midpoint. + assert result["only"]["alpha"] == pytest.approx(0.5, abs=0.2) + + +def test_recommend_bucketed_config_separate_buckets_stay_separate( + bucketed_factors: list[Factor], + bucketed_observables: list[Observable], +) -> None: + regimes = { + "r1": {"target": 0.1, "method": "a"}, + "r2": {"target": 0.9, "method": "a"}, + } + result = recommend_bucketed_config( + regimes, + bucket_fn=lambda name, _r: name, # each regime is its own bucket + world_factory=lambda r: _TargetWorld(target=r["target"], method=r["method"]), + scorer=_TargetScorer(), + factors=bucketed_factors, + observables=bucketed_observables, + primary="cost", + n_trials=40, + seed=0, + ) + assert set(result.keys()) == {"r1", "r2"} + assert result["r1"]["alpha"] == pytest.approx(0.1, abs=0.2) + assert result["r2"]["alpha"] == pytest.approx(0.9, abs=0.2) + + +def test_recommend_bucketed_config_categorical_uses_mode( + bucketed_factors: list[Factor], + bucketed_observables: list[Observable], +) -> None: + regimes = { + "r1": {"target": 0.5, "method": "a"}, + "r2": {"target": 0.5, "method": "a"}, + "r3": {"target": 0.5, "method": "b"}, + } + result = recommend_bucketed_config( + regimes, + bucket_fn=lambda _name, _r: "only", + world_factory=lambda r: _TargetWorld(target=r["target"], method=r["method"]), + scorer=_TargetScorer(), + factors=bucketed_factors, + observables=bucketed_observables, + primary="reward", + n_trials=30, + seed=0, + ) + # two of three regimes reward method="a"; mode should pick it. + assert result["only"]["method"] == "a" + + +def test_recommend_bucketed_config_respects_maximize_direction( + bucketed_factors: list[Factor], + bucketed_observables: list[Observable], +) -> None: + regimes = {"r1": {"target": 0.5, "method": "a"}} + result = recommend_bucketed_config( + regimes, + bucket_fn=lambda _name, _r: "only", + world_factory=lambda r: _TargetWorld(target=r["target"], method=r["method"]), + scorer=_TargetScorer(), + factors=bucketed_factors, + observables=bucketed_observables, + primary="reward", + n_trials=20, + seed=0, + ) + assert result["only"]["method"] == "a" + + +def test_recommend_bucketed_config_rejects_empty_regimes( + bucketed_factors: list[Factor], + bucketed_observables: list[Observable], +) -> None: + with pytest.raises(ValueError, match="regimes must be non-empty"): + recommend_bucketed_config( + {}, + bucket_fn=lambda _n, _r: "b", + world_factory=lambda r: _TargetWorld( + target=r["target"], method=r["method"] + ), + scorer=_TargetScorer(), + factors=bucketed_factors, + observables=bucketed_observables, + primary="cost", + ) + + +def test_recommend_bucketed_config_rejects_unknown_primary( + bucketed_factors: list[Factor], + bucketed_observables: list[Observable], +) -> None: + regimes = {"r1": {"target": 0.5, "method": "a"}} + with pytest.raises(ValueError, match="not found in observables"): + recommend_bucketed_config( + regimes, + bucket_fn=lambda _n, _r: "b", + world_factory=lambda r: _TargetWorld( + target=r["target"], method=r["method"] + ), + scorer=_TargetScorer(), + factors=bucketed_factors, + observables=bucketed_observables, + primary="bogus", + )