mirror of
https://github.com/manifoldbt/manifoldbt.git
synced 2026-08-24 14:38:04 +00:00
141 lines
4.3 KiB
Python
141 lines
4.3 KiB
Python
"""Tests for bt.choice() — sweeping a CHOICE of expression, not just a number.
|
|
|
|
The contract under test: `choice("sel", {...})` resolves to exactly one branch
|
|
per combo BEFORE simulation, so a sweep over the selector must be
|
|
bit-identical to running each branch inlined by hand. The selector must count
|
|
as a declared parameter (otherwise `_validate_swept_params` would reject the
|
|
sweep), and an unknown branch must fail with a message naming the known ones.
|
|
"""
|
|
import json
|
|
import os
|
|
|
|
import pytest
|
|
|
|
import manifoldbt as bt
|
|
from manifoldbt.indicators import close, sma
|
|
|
|
pd = pytest.importorskip("pandas")
|
|
np = pytest.importorskip("numpy")
|
|
|
|
N_BARS = 3_000
|
|
|
|
|
|
def _store(tmp_path):
|
|
ts = pd.date_range("2022-01-01", periods=N_BARS, freq="1min", tz="UTC")
|
|
rng = np.random.default_rng(11)
|
|
px = 100.0 + np.cumsum(np.sin(np.arange(N_BARS) / 90.0) * 0.3 + rng.normal(0, 0.2, N_BARS)) * 0.05
|
|
px = np.maximum(px, 1.0)
|
|
df = pd.DataFrame(
|
|
{
|
|
"timestamp": ts,
|
|
"open": px,
|
|
"high": px * 1.0005,
|
|
"low": px * 0.9995,
|
|
"close": px,
|
|
"volume": [1000.0] * N_BARS,
|
|
}
|
|
)
|
|
root = tmp_path / "choice_store"
|
|
return bt.import_dataframe(
|
|
df,
|
|
symbol="ZC",
|
|
symbol_id=1,
|
|
interval="1m",
|
|
asset_class="equity",
|
|
exchange="TEST",
|
|
data_root=str(root / "data"),
|
|
metadata_db=str(root / "metadata.sqlite"),
|
|
)
|
|
|
|
|
|
def _config():
|
|
start, end = bt.time_range("2022-01-01", "2022-01-03")
|
|
cfg = bt.BacktestConfig(
|
|
universe=[1],
|
|
time_range_start=start,
|
|
time_range_end=end,
|
|
initial_capital=10_000.0,
|
|
provider="TEST",
|
|
bar_interval=bt.Interval.minutes(1),
|
|
symbol_names={"ZC": 1},
|
|
)
|
|
cfg.warmup_bars = 0
|
|
return cfg
|
|
|
|
|
|
def _strategy_with_choice():
|
|
band = bt.choice(
|
|
"pick",
|
|
{
|
|
"fast": sma(close, 5),
|
|
"slow": sma(close, 20),
|
|
},
|
|
)
|
|
return (
|
|
bt.Strategy.create("choice_e2e")
|
|
.signal("band", band)
|
|
.size(bt.when(close > bt.col("band"), 1.0, 0.0))
|
|
)
|
|
|
|
|
|
def _strategy_inlined(period):
|
|
return (
|
|
bt.Strategy.create(f"inline_{period}")
|
|
.signal("band", sma(close, period))
|
|
.size(bt.when(close > bt.col("band"), 1.0, 0.0))
|
|
)
|
|
|
|
|
|
def test_serializes_as_ordered_pairs():
|
|
"""serde expects Choice(String, Vec<(String, Expr)>): a list of pairs,
|
|
order preserved — the first branch is the compile-time default."""
|
|
e = bt.choice("pick", {"a": close, "b": sma(close, 3)})
|
|
payload = json.loads(json.dumps(e.to_json()))
|
|
assert list(payload) == ["Choice"]
|
|
name, branches = payload["Choice"]
|
|
assert name == "pick"
|
|
assert [k for k, _ in branches] == ["a", "b"]
|
|
|
|
|
|
def test_empty_branches_rejected():
|
|
with pytest.raises(ValueError, match="at least one branch"):
|
|
bt.choice("pick", {})
|
|
|
|
|
|
def test_selector_counts_as_declared_parameter():
|
|
"""Sweeping the selector must pass strategy-side validation: choice()
|
|
declares it via _param_meta exactly like param() does."""
|
|
strat = _strategy_with_choice()
|
|
assert "pick" in (strat.to_json_dict().get("parameters") or {})
|
|
|
|
|
|
def test_sweep_over_choice_matches_inlined_branches(tmp_path):
|
|
"""The money test: each combo of the selector sweep is bit-identical to
|
|
the strategy with that branch written directly."""
|
|
store = _store(tmp_path)
|
|
cfg = _config()
|
|
|
|
sweep = bt.run_sweep_lite(
|
|
_strategy_with_choice(), {"pick": ["fast", "slow"]}, cfg, store, device="cpu"
|
|
)
|
|
assert len(sweep) == 2
|
|
|
|
by_branch = dict(zip(["fast", "slow"], sweep))
|
|
for name, period in (("fast", 5), ("slow", 20)):
|
|
ref = bt.run_sweep_lite(
|
|
_strategy_inlined(period), {}, cfg, store, device="cpu"
|
|
)[0]
|
|
got, want = by_branch[name].metrics, ref.metrics
|
|
for key in ("total_return", "sharpe", "max_drawdown"):
|
|
assert got.get(key) == want.get(key), (
|
|
f"branch {name!r}: {key} diverged ({got.get(key)} vs {want.get(key)})"
|
|
)
|
|
|
|
|
|
def test_unknown_branch_names_the_known_ones(tmp_path):
|
|
store = _store(tmp_path)
|
|
with pytest.raises(Exception, match="fast"):
|
|
bt.run_sweep_lite(
|
|
_strategy_with_choice(), {"pick": ["nope"]}, _config(), store, device="cpu"
|
|
)
|