From 758c7d51c9be2d033a7e5028d94dd86567f55f2a Mon Sep 17 00:00:00 2001 From: Stefan Jansen Date: Thu, 23 Jul 2026 19:53:53 -0400 Subject: [PATCH] Support opening risk rules and entry-bar trailing state --- src/ml4t/backtest/broker.py | 24 ++++- src/ml4t/backtest/config.py | 7 ++ src/ml4t/backtest/core/order_book.py | 3 + src/ml4t/backtest/core/risk_engine.py | 34 ++++--- src/ml4t/backtest/core/shared.py | 3 + src/ml4t/backtest/engine.py | 4 + src/ml4t/backtest/strategy.py | 10 ++ tests/test_broker.py | 27 +++++ tests/test_pre_risk_strategy.py | 139 ++++++++++++++++++++++++++ 9 files changed, 234 insertions(+), 17 deletions(-) create mode 100644 tests/test_pre_risk_strategy.py diff --git a/src/ml4t/backtest/broker.py b/src/ml4t/backtest/broker.py index 67799c18..6db6ebf0 100644 --- a/src/ml4t/backtest/broker.py +++ b/src/ml4t/backtest/broker.py @@ -66,6 +66,7 @@ def __init__( stop_fill_mode: StopFillMode = StopFillMode.STOP_PRICE, stop_level_basis: StopLevelBasis = StopLevelBasis.FILL_PRICE, trail_hwm_source: WaterMarkSource = WaterMarkSource.CLOSE, + trail_include_entry_bar_extremes: bool = False, initial_hwm_source: InitialHwmSource = InitialHwmSource.FILL_PRICE, trail_stop_timing: TrailStopTiming = TrailStopTiming.LAGGED, allow_short_selling: bool = False, @@ -120,6 +121,7 @@ def __init__( self.stop_fill_mode = stop_fill_mode self.stop_level_basis = stop_level_basis self.trail_hwm_source = trail_hwm_source + self.trail_include_entry_bar_extremes = trail_include_entry_bar_extremes self.initial_hwm_source = initial_hwm_source self.trail_stop_timing = trail_stop_timing self.share_type = share_type @@ -360,6 +362,7 @@ def from_config( stop_fill_mode=config.stop_fill_mode, stop_level_basis=config.stop_level_basis, trail_hwm_source=config.trail_hwm_source, + trail_include_entry_bar_extremes=config.trail_include_entry_bar_extremes, initial_hwm_source=config.initial_hwm_source, trail_stop_timing=config.trail_stop_timing, allow_short_selling=config.allow_short_selling, @@ -767,18 +770,27 @@ def last_rejection_reason(self) -> str | None: # === Risk Management === - def set_position_rules(self, rules: PositionRule, asset: str | None = None) -> None: + def set_position_rules( + self, + rules: PositionRule | None, + asset: str | None = None, + ) -> None: """Set position rules globally or per-asset. Args: - rules: PositionRule or RuleChain to apply + rules: PositionRule or RuleChain to apply. ``None`` explicitly + disables rules for the selected scope. asset: If provided, apply only to this asset; otherwise global """ - if asset: + if asset is not None: self._position_rules_by_asset[asset] = rules else: self._position_rules = rules + def clear_position_rules(self, asset: str | None = None) -> None: + """Disable position rules globally or for one asset.""" + self.set_position_rules(None, asset=asset) + def update_position_context(self, asset: str, context: dict) -> None: """Update context data for a position (used by signal-based rules). @@ -1725,6 +1737,8 @@ def _update_water_marks(self): Water mark source configuration: - trail_hwm_source == BAR_EXTREME: Update HWM from high, LWM from low (VBT Pro OHLC mode) - trail_hwm_source == CLOSE: Update HWM/LWM from close only (default) + - trail_include_entry_bar_extremes: Include a completed entry bar's + extreme in the watermark used from the next bar onward """ for asset, pos in self.positions.items(): if asset in self._current_prices: @@ -1732,7 +1746,9 @@ def _update_water_marks(self): # VBT Pro only updates water marks from bar extremes on the bar AFTER entry is_new_position = asset in self._positions_created_this_bar # BAR_EXTREME: use HIGH for HWM (longs), LOW for LWM (shorts) - use_extremes = self.trail_hwm_source.value == "bar_extreme" and not is_new_position + use_extremes = self.trail_hwm_source.value == "bar_extreme" and ( + not is_new_position or self.trail_include_entry_bar_extremes + ) pos.update_water_marks( current_price=self._current_prices[asset], bar_high=self._current_highs.get(asset), diff --git a/src/ml4t/backtest/config.py b/src/ml4t/backtest/config.py index d168c4f4..c6aa3dc2 100644 --- a/src/ml4t/backtest/config.py +++ b/src/ml4t/backtest/config.py @@ -512,6 +512,7 @@ class BacktestConfig: stop_fill_mode: StopFillMode = StopFillMode.STOP_PRICE stop_level_basis: StopLevelBasis = StopLevelBasis.FILL_PRICE trail_hwm_source: WaterMarkSource = WaterMarkSource.CLOSE + trail_include_entry_bar_extremes: bool = False initial_hwm_source: InitialHwmSource = InitialHwmSource.FILL_PRICE trail_stop_timing: TrailStopTiming = TrailStopTiming.LAGGED @@ -840,6 +841,7 @@ def to_dict(self) -> dict: "stop_fill_mode": self.stop_fill_mode.value, "stop_level_basis": self.stop_level_basis.value, "trail_hwm_source": self.trail_hwm_source.value, + "trail_include_entry_bar_extremes": self.trail_include_entry_bar_extremes, "initial_hwm_source": self.initial_hwm_source.value, "trail_stop_timing": self.trail_stop_timing.value, }, @@ -946,6 +948,7 @@ def from_dict( "stop_fill_mode", "stop_level_basis", "trail_hwm_source", + "trail_include_entry_bar_extremes", "initial_hwm_source", "trail_stop_timing", }, @@ -1073,6 +1076,9 @@ def from_dict( stop_fill_mode=StopFillMode(stops_cfg.get("stop_fill_mode", "stop_price")), stop_level_basis=StopLevelBasis(stops_cfg.get("stop_level_basis", "fill_price")), trail_hwm_source=WaterMarkSource(stops_cfg.get("trail_hwm_source", "close")), + trail_include_entry_bar_extremes=stops_cfg.get( + "trail_include_entry_bar_extremes", False + ), initial_hwm_source=InitialHwmSource(stops_cfg.get("initial_hwm_source", "fill_price")), trail_stop_timing=TrailStopTiming(stops_cfg.get("trail_stop_timing", "lagged")), # Sizing @@ -1337,6 +1343,7 @@ def describe(self) -> str: f" Fill mode: {self.stop_fill_mode.value}", f" Level basis: {self.stop_level_basis.value}", f" Trail HWM source: {self.trail_hwm_source.value}", + f" Include entry-bar extremes: {self.trail_include_entry_bar_extremes}", f" Trail timing: {self.trail_stop_timing.value}", "", "Position Sizing:", diff --git a/src/ml4t/backtest/core/order_book.py b/src/ml4t/backtest/core/order_book.py index 8822bb88..8ff7dcb5 100644 --- a/src/ml4t/backtest/core/order_book.py +++ b/src/ml4t/backtest/core/order_book.py @@ -62,6 +62,9 @@ def submit_order( order_id=f"ORD-{broker._order_counter}", created_at=broker._current_time, _created_bar_index=broker._bar_index, + _risk_exit_reason=options.risk_exit_reason if options is not None else None, + _exit_reason=options.exit_reason if options is not None else None, + _risk_fill_price=options.risk_fill_price if options is not None else None, ) order._signal_price = broker._current_prices.get(asset) diff --git a/src/ml4t/backtest/core/risk_engine.py b/src/ml4t/backtest/core/risk_engine.py index f462d10f..a4216b8f 100644 --- a/src/ml4t/backtest/core/risk_engine.py +++ b/src/ml4t/backtest/core/risk_engine.py @@ -42,12 +42,14 @@ def evaluate_position_rules(self): asset, -pos.quantity, order_type=OrderType.MARKET, - _options=SubmitOrderOptions(eligible_in_next_bar_mode=True), + _options=SubmitOrderOptions( + eligible_in_next_bar_mode=True, + risk_exit_reason=action.reason, + exit_reason=reason_to_exit_reason(action.reason), + risk_fill_price=action.fill_price, + ), ) if order: - order._risk_exit_reason = action.reason - order._exit_reason = reason_to_exit_reason(action.reason) - order._risk_fill_price = action.fill_price exit_orders.append(order) broker._stop_exits_this_bar.add(asset) @@ -69,19 +71,23 @@ def evaluate_position_rules(self): asset, actual_qty, order_type=OrderType.MARKET, - _options=SubmitOrderOptions(eligible_in_next_bar_mode=True), + _options=SubmitOrderOptions( + eligible_in_next_bar_mode=True, + risk_exit_reason=action.reason, + exit_reason=reason_to_exit_reason(action.reason), + risk_fill_price=action.fill_price, + ), ) if order: - order._risk_exit_reason = action.reason - order._exit_reason = reason_to_exit_reason(action.reason) - order._risk_fill_price = action.fill_price exit_orders.append(order) return exit_orders def _get_position_rules(self, asset: str): broker = self.broker - return broker._position_rules_by_asset.get(asset) or broker._position_rules + if asset in broker._position_rules_by_asset: + return broker._position_rules_by_asset[asset] + return broker._position_rules def _build_position_state(self, pos, current_price: float): broker = self.broker @@ -144,12 +150,14 @@ def process_pending_exits(self): asset, -exit_qty, order_type=OrderType.MARKET, - _options=SubmitOrderOptions(eligible_in_next_bar_mode=True), + _options=SubmitOrderOptions( + eligible_in_next_bar_mode=True, + risk_exit_reason=pending["reason"], + exit_reason=reason_to_exit_reason(pending["reason"]), + risk_fill_price=fill_price, + ), ) if order: - order._risk_exit_reason = pending["reason"] - order._exit_reason = reason_to_exit_reason(pending["reason"]) - order._risk_fill_price = fill_price exit_orders.append(order) del broker._pending_exits[asset] diff --git a/src/ml4t/backtest/core/shared.py b/src/ml4t/backtest/core/shared.py index 9e825c7f..6694bd77 100644 --- a/src/ml4t/backtest/core/shared.py +++ b/src/ml4t/backtest/core/shared.py @@ -22,6 +22,9 @@ class SubmitOrderOptions: eligible_in_next_bar_mode: bool = False rebalance_id: str | None = None + risk_exit_reason: str | None = None + exit_reason: ExitReason | None = None + risk_fill_price: float | None = None def is_exit_order(order: Order, positions: dict[str, Position]) -> bool: diff --git a/src/ml4t/backtest/engine.py b/src/ml4t/backtest/engine.py index 439e6817..47b57b06 100644 --- a/src/ml4t/backtest/engine.py +++ b/src/ml4t/backtest/engine.py @@ -222,6 +222,10 @@ def run(self) -> BacktestResult: # This must happen BEFORE evaluate_position_rules() to clear deferred exits self.broker._process_pending_exits() + # Optional strategy phase for opening orders that must receive risk + # protection during the current bar. Existing strategies inherit a no-op. + self.strategy.on_before_risk(timestamp, assets_data, context, self.broker) + # Evaluate position rules (stops, trails, etc.) - generates exit orders self.broker.evaluate_position_rules() diff --git a/src/ml4t/backtest/strategy.py b/src/ml4t/backtest/strategy.py index b9c7606c..3cefb765 100644 --- a/src/ml4t/backtest/strategy.py +++ b/src/ml4t/backtest/strategy.py @@ -9,6 +9,16 @@ class Strategy(ABC): """Base strategy class.""" + def on_before_risk( + self, + timestamp: datetime, + data: dict[str, dict], + context: dict[str, Any], + broker: Any, + ) -> None: + """Called before position rules are evaluated for the current bar.""" + return None + @abstractmethod def on_data( self, diff --git a/tests/test_broker.py b/tests/test_broker.py index 7138f013..e6f660f1 100644 --- a/tests/test_broker.py +++ b/tests/test_broker.py @@ -1309,6 +1309,33 @@ def test_set_position_rules_global(self): # Should apply to assets without explicit overrides assert broker._position_rules_by_asset.get("AAPL") is None + def test_clear_position_rules_for_one_asset(self): + """Clearing one asset disables its rule without changing other assets.""" + from ml4t.backtest.risk.position.static import StopLoss + + broker = Broker(100000.0, NoCommission(), NoSlippage()) + stop_rule = StopLoss(pct=0.05) + broker.set_position_rules(stop_rule, asset="AAPL") + broker.set_position_rules(stop_rule, asset="MSFT") + + broker.clear_position_rules(asset="AAPL") + + assert broker._risk_engine._get_position_rules("AAPL") is None + assert broker._risk_engine._get_position_rules("MSFT") == stop_rule + + def test_asset_rule_can_explicitly_disable_global_rule(self): + """An explicit per-asset None overrides, rather than falls back to, a global rule.""" + from ml4t.backtest.risk.position.static import StopLoss + + broker = Broker(100000.0, NoCommission(), NoSlippage()) + stop_rule = StopLoss(pct=0.05) + broker.set_position_rules(stop_rule) + + broker.set_position_rules(None, asset="AAPL") + + assert broker._risk_engine._get_position_rules("AAPL") is None + assert broker._risk_engine._get_position_rules("MSFT") == stop_rule + def test_update_position_context(self): """Test updating position context.""" broker = Broker(100000.0, NoCommission(), NoSlippage()) diff --git a/tests/test_pre_risk_strategy.py b/tests/test_pre_risk_strategy.py new file mode 100644 index 00000000..0e014fbd --- /dev/null +++ b/tests/test_pre_risk_strategy.py @@ -0,0 +1,139 @@ +"""End-to-end tests for strategy work that must run before position risk.""" + +from datetime import datetime + +import polars as pl + +from ml4t.backtest import ( + BacktestConfig, + Broker, + DataFeed, + Engine, + ExecutionMode, + StopLoss, + Strategy, + TrailingStop, +) +from ml4t.backtest.config import ExecutionPrice, WaterMarkSource + + +class OpeningTargetWithStop(Strategy): + """Enter at the session open and protect the new position on that bar.""" + + def on_before_risk(self, timestamp, data, context, broker) -> None: + if broker.get_position("SPY") is None: + broker.set_position_rules(StopLoss(pct=0.05), asset="SPY") + broker.submit_order("SPY", 100) + + def on_data(self, timestamp, data, context, broker) -> None: + pass + + +class OpeningTargetWithTrailingStop(Strategy): + """Enter once and retain a trailing stop across sessions.""" + + def on_before_risk(self, timestamp, data, context, broker) -> None: + if not broker.fills: + broker.set_position_rules(TrailingStop(pct=0.05), asset="SPY") + broker.submit_order("SPY", 100) + + def on_data(self, timestamp, data, context, broker) -> None: + pass + + +def test_pre_risk_entry_can_trigger_stop_on_entry_bar(): + """A position entered at the open receives stop protection on the same bar.""" + prices = pl.DataFrame( + { + "timestamp": [datetime(2026, 8, 3)], + "asset": ["SPY"], + "open": [100.0], + "high": [101.0], + "low": [94.0], + "close": [98.0], + "volume": [1_000_000.0], + } + ) + config = BacktestConfig( + initial_cash=100_000.0, + execution_mode=ExecutionMode.SAME_BAR, + execution_price=ExecutionPrice.OPEN, + immediate_fill=True, + ) + + result = Engine(DataFeed(prices_df=prices), OpeningTargetWithStop(), config).run() + + assert [(fill.side.value, fill.price) for fill in result.fills] == [ + ("buy", 100.0), + ("sell", 95.0), + ] + assert result.trades[0].exit_reason == "stop_loss" + + +def test_entry_bar_extreme_becomes_next_bar_trailing_watermark_when_enabled(): + """Entry-bar highs are available to a lagged trail only after that bar completes.""" + prices = pl.DataFrame( + { + "timestamp": [datetime(2026, 8, 3), datetime(2026, 8, 4)], + "asset": ["SPY", "SPY"], + "open": [100.0, 105.0], + "high": [110.0, 106.0], + "low": [99.0, 103.0], + "close": [105.0, 104.0], + "volume": [1_000_000.0, 1_000_000.0], + } + ) + config = BacktestConfig( + initial_cash=100_000.0, + execution_mode=ExecutionMode.SAME_BAR, + execution_price=ExecutionPrice.OPEN, + immediate_fill=True, + trail_hwm_source=WaterMarkSource.BAR_EXTREME, + trail_include_entry_bar_extremes=True, + ) + + result = Engine(DataFeed(prices_df=prices), OpeningTargetWithTrailingStop(), config).run() + + assert [(fill.side.value, fill.price) for fill in result.fills] == [ + ("buy", 100.0), + ("sell", 104.5), + ] + assert result.fills[1].timestamp == datetime(2026, 8, 4) + assert result.trades[0].exit_reason == "trailing_stop" + + +def test_entry_bar_extreme_remains_excluded_by_default(): + """The opt-in does not change the existing entry-bar watermark contract.""" + prices = pl.DataFrame( + { + "timestamp": [datetime(2026, 8, 3), datetime(2026, 8, 4)], + "asset": ["SPY", "SPY"], + "open": [100.0, 105.0], + "high": [110.0, 106.0], + "low": [99.0, 103.0], + "close": [105.0, 104.0], + "volume": [1_000_000.0, 1_000_000.0], + } + ) + config = BacktestConfig( + initial_cash=100_000.0, + execution_mode=ExecutionMode.SAME_BAR, + execution_price=ExecutionPrice.OPEN, + immediate_fill=True, + trail_hwm_source=WaterMarkSource.BAR_EXTREME, + ) + + result = Engine(DataFeed(prices_df=prices), OpeningTargetWithTrailingStop(), config).run() + + assert [(fill.side.value, fill.price) for fill in result.fills] == [("buy", 100.0)] + + +def test_entry_bar_extreme_option_roundtrips_and_reaches_broker(): + """Serialized configs preserve the opt-in and Broker.from_config receives it.""" + config = BacktestConfig(trail_include_entry_bar_extremes=True) + + restored = BacktestConfig.from_dict(config.to_dict()) + broker = Broker.from_config(restored) + + assert restored.trail_include_entry_bar_extremes is True + assert broker.trail_include_entry_bar_extremes is True