refactor: split recommend_bucketed_config into search + aggregation

CHANGELOG.md

@@ -12,6 +12,7 @@ All notable changes to this project will be documented in this file.
 - `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.
+- `recommend_bucketed_config()` split into `recommend_per_regime()` (the expensive adaptive search) and `aggregate_bucketed_config()` (cheap post-processing), with `recommend_bucketed_config()` now a thin wrapper composing them. Lets a caller experiment with different `bucket_fn` groupings against the same search results without re-running `run_adaptive` for every attempt -- found necessary in practice deriving VBPCApy buckets, where a first grouping choice performed poorly and needed re-grouping without repeating an hours-long search.
 
 ## [0.2.0] — 2026-08-17
 

src/trade_study/__init__.py

@@ -27,7 +27,13 @@
     Simulator,
     TrialResult,
 )
-from .regime import RegimeSurrogate, fit_regime_surrogate, recommend_bucketed_config
+from .regime import (
+    RegimeSurrogate,
+    aggregate_bucketed_config,
+    fit_regime_surrogate,
+    recommend_bucketed_config,
+    recommend_per_regime,
+)
 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
@@ -60,6 +66,7 @@
     "TableSensitivity",
     "TrialResult",
     "__version__",
+    "aggregate_bucketed_config",
     "build_grid",
     "coverage_curve",
     "ensemble_predict",
@@ -76,6 +83,7 @@
     "plot_parallel",
     "plot_scores",
     "recommend_bucketed_config",
+    "recommend_per_regime",
     "reduce_factors",
     "run_adaptive",
     "run_grid",

src/trade_study/regime.py

@@ -325,9 +325,8 @@ def _aggregate_factor_values(factor: Factor, values: list[Any]) -> Any:  # ruff:
     return type(values[0])(np.median(values))
 
 
-def recommend_bucketed_config(  # ruff: ignore[too-many-arguments]
+def recommend_per_regime(  # 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],
@@ -338,24 +337,19 @@ def recommend_bucketed_config(  # ruff: ignore[too-many-arguments]
     n_reps: int = 1,
     seed: int = 42,
 ) -> dict[str, dict[str, Any]]:
-    """Recommend a config per named bucket via per-regime adaptive search (#123).
+    """Find each regime's best config via independent 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.
+    Runs :func:`~trade_study.run_adaptive` (NSGA-II) independently per
+    regime and picks each regime's best-found config by ``primary``. This
+    is the expensive half of :func:`recommend_bucketed_config` (the other
+    half, :func:`aggregate_bucketed_config`, is pure post-processing) --
+    split out so a caller can experiment with different bucket groupings
+    (different ``bucket_fn``s) against the same search results without
+    re-running it.
 
     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).
@@ -373,14 +367,14 @@ def recommend_bucketed_config(  # ruff: ignore[too-many-arguments]
         seed: Random seed forwarded to each regime's ``run_adaptive`` call.
 
     Returns:
-        Mapping from bucket name to its aggregated recommended config.
+        Mapping from regime name to its best-found 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"
+        msg = "recommend_per_regime: regimes must be non-empty"
         raise ValueError(msg)
     matching = [o for o in observables if o.name == primary]
     if not matching:
@@ -389,7 +383,6 @@ def recommend_bucketed_config(  # ruff: ignore[too-many-arguments]
     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(
@@ -408,6 +401,39 @@ def recommend_bucketed_config(  # ruff: ignore[too-many-arguments]
             else int(np.argmax(table.scores[:, primary_col]))
         )
         per_regime_best[name] = table.configs[best_i]
+    return per_regime_best
+
+
+def aggregate_bucketed_config(
+    per_regime_best: dict[str, dict[str, Any]],
+    regimes: dict[str, dict[str, Any]],
+    bucket_fn: Callable[[str, dict[str, Any]], str],
+    factors: list[Factor],
+) -> dict[str, dict[str, Any]]:
+    """Aggregate per-regime best configs into named buckets (#123).
+
+    Pure post-processing over :func:`recommend_per_regime`'s output --
+    groups regimes into named buckets via ``bucket_fn`` and aggregates
+    each bucket's members' best configs (median for continuous/discrete
+    factors, mode for categorical). Cheap enough to call repeatedly with
+    different ``bucket_fn``s (e.g. finer-grained groupings) against the
+    same search results.
+
+    Args:
+        per_regime_best: Output of :func:`recommend_per_regime`.
+        regimes: The same regime dict passed to :func:`recommend_per_regime`
+            (``bucket_fn`` receives each regime's descriptor, not just its
+            name).
+        bucket_fn: Maps ``(regime_name, regime_dict)`` to a bucket name;
+            regimes sharing a bucket name have their best configs
+            aggregated together.
+        factors: The same factor list passed to :func:`recommend_per_regime`.
+
+    Returns:
+        Mapping from bucket name to its aggregated recommended config.
+    """
+    bucket_members: dict[str, list[str]] = defaultdict(list)
+    for name, regime in regimes.items():
         bucket_members[bucket_fn(name, regime)].append(name)
 
     factors_by_name = {f.name: f for f in factors}
@@ -421,3 +447,72 @@ def recommend_bucketed_config(  # ruff: ignore[too-many-arguments]
         }
         for bucket, members in bucket_members.items()
     }
+
+
+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.
+
+    A convenience wrapper composing :func:`recommend_per_regime` (the
+    expensive search) and :func:`aggregate_bucketed_config` (cheap
+    post-processing) -- call those directly if you want to try more than
+    one ``bucket_fn`` without repeating the search.
+
+    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.
+    """
+    per_regime_best = recommend_per_regime(
+        regimes,
+        world_factory,
+        scorer,
+        factors,
+        observables,
+        primary=primary,
+        n_trials=n_trials,
+        n_reps=n_reps,
+        seed=seed,
+    )
+    return aggregate_bucketed_config(per_regime_best, regimes, bucket_fn, factors)

tests/test_regime.py

@@ -14,9 +14,11 @@
     FactorType,
     Observable,
     RegimeSurrogate,
+    aggregate_bucketed_config,
     build_grid,
     fit_regime_surrogate,
     recommend_bucketed_config,
+    recommend_per_regime,
     run_grid,
 )
 
@@ -521,3 +523,71 @@ def test_recommend_bucketed_config_rejects_unknown_primary(
             observables=bucketed_observables,
             primary="bogus",
         )
+
+
+# ---------------------------------------------------------------------------
+# recommend_per_regime / aggregate_bucketed_config split (#123 follow-up)
+# ---------------------------------------------------------------------------
+
+
+def test_split_functions_match_combined_wrapper(
+    bucketed_factors: list[Factor],
+    bucketed_observables: list[Observable],
+) -> None:
+    regimes = {
+        "r1": {"target": 0.2, "method": "a"},
+        "r2": {"target": 0.8, "method": "a"},
+    }
+    kwargs = {
+        "world_factory": lambda r: _TargetWorld(target=r["target"], method=r["method"]),
+        "scorer": _TargetScorer(),
+        "factors": bucketed_factors,
+        "observables": bucketed_observables,
+        "primary": "cost",
+        "n_trials": 30,
+        "seed": 0,
+    }
+    combined = recommend_bucketed_config(
+        regimes, bucket_fn=lambda _n, _r: "only", **kwargs
+    )
+    per_regime = recommend_per_regime(regimes, **kwargs)
+    split = aggregate_bucketed_config(
+        per_regime, regimes, bucket_fn=lambda _n, _r: "only", factors=bucketed_factors
+    )
+    assert combined == split
+
+
+def test_aggregate_bucketed_config_reuses_search_for_different_groupings(
+    bucketed_factors: list[Factor],
+    bucketed_observables: list[Observable],
+) -> None:
+    """Re-bucketing with a different bucket_fn needs no new search calls."""
+    regimes = {
+        "r1": {"target": 0.1, "method": "a"},
+        "r2": {"target": 0.5, "method": "a"},
+        "r3": {"target": 0.9, "method": "a"},
+    }
+    per_regime = recommend_per_regime(
+        regimes,
+        world_factory=lambda r: _TargetWorld(target=r["target"], method=r["method"]),
+        scorer=_TargetScorer(),
+        factors=bucketed_factors,
+        observables=bucketed_observables,
+        primary="cost",
+        n_trials=30,
+        seed=0,
+    )
+    assert set(per_regime) == {"r1", "r2", "r3"}
+
+    one_bucket = aggregate_bucketed_config(
+        per_regime, regimes, bucket_fn=lambda _n, _r: "only", factors=bucketed_factors
+    )
+    assert set(one_bucket) == {"only"}
+
+    per_regime_buckets = aggregate_bucketed_config(
+        per_regime, regimes, bucket_fn=lambda name, _r: name, factors=bucketed_factors
+    )
+    assert set(per_regime_buckets) == {"r1", "r2", "r3"}
+    # each regime's own bucket should just be its own best config verbatim.
+    for name in regimes:
+        assert per_regime_buckets[name] == per_regime[name]