diff --git a/src/easy_tdx/backtest/engine.py b/src/easy_tdx/backtest/engine.py new file mode 100644 index 0000000..46c5790 --- /dev/null +++ b/src/easy_tdx/backtest/engine.py @@ -0,0 +1,263 @@ +"""BacktestEngine — orchestrate vectorized execution pipeline. + +Coordinates Strategy → OrderSimulator → PortfolioTracker → PerformanceAnalyzer. +""" + +from __future__ import annotations + +from typing import Any + +import pandas as pd + +from easy_tdx.backtest.orders import OrderSimulator +from easy_tdx.backtest.performance import PerformanceAnalyzer +from easy_tdx.backtest.portfolio import PortfolioTracker +from easy_tdx.backtest.strategy import Strategy +from easy_tdx.backtest.types import BacktestResult, Signal, Trade + + +class BacktestEngine: + """Orchestrate backtest execution pipeline. + + Pipeline: + 1. Signal generation (Strategy) + 2. Order simulation (OrderSimulator) + 3. Portfolio tracking (PortfolioTracker) + 4. Performance analysis (PerformanceAnalyzer) + + Example: + >>> engine = BacktestEngine(MyStrategy, cash=100000) + >>> result = engine.run(df) + """ + + def __init__( + self, + strategy: type[Strategy] | Strategy, + cash: float = 100000.0, + commission: float = 0.0003, + min_commission: float = 5.0, + stamp_tax: float = 0.001, + slippage: float = 0.0, + execution: str = "next_open", + position_mode: str = "full", + reject_policy: str = "reduce", + benchmark: pd.DataFrame | None = None, + ): + """Initialize engine. + + Args: + strategy: Strategy class or instance + cash: Initial cash + commission: Commission rate (e.g., 0.0003 = 0.03%) + min_commission: Minimum commission per trade + stamp_tax: Stamp tax rate (for sells) + slippage: Slippage rate + execution: Execution mode ('next_open', 'this_close') + position_mode: Position mode ('full', 'long_only', 'short_only') + reject_policy: Reject policy ('reduce', 'reject') + benchmark: Benchmark data for performance comparison + """ + self._strategy_cls = strategy if isinstance(strategy, type) else type(strategy) + self._strategy_instance = strategy if isinstance(strategy, Strategy) else None + + self._cash = cash + self._commission = commission + self._min_commission = min_commission + self._stamp_tax = stamp_tax + self._slippage = slippage + self._execution = execution + self._position_mode = position_mode + self._reject_policy = reject_policy + self._benchmark = benchmark + + def run(self, df: pd.DataFrame, chanlun_result: Any | None = None) -> BacktestResult: + """Run backtest. + + Args: + df: Price data with OHLCV columns + chanlun_result: Optional chanlun analysis result for strategy + + Returns: + BacktestResult with performance, equity_curve, trades, positions, config + """ + if len(df) == 0: + return self._empty_result() + + # Step 1: Signal generation + signals = self._generate_signals(df, chanlun_result) + + # Step 2: Order simulation + simulator = OrderSimulator( + df, + execution=self._execution, + position_mode=self._position_mode, + reject_policy=self._reject_policy, + commission=self._commission, + min_commission=self._min_commission, + stamp_tax=self._stamp_tax, + slippage=self._slippage, + ) + trades = simulator.simulate( + signals=signals, + cash=self._cash, + position=0.0, + ) + + # Step 3: Portfolio tracking + trades = self._compute_pnls(trades) + tracker = PortfolioTracker(df, initial_cash=self._cash) + tracker.apply_trades(trades) + + # Step 4: Performance analysis + trades_df = self._trades_to_df(trades) + performance = PerformanceAnalyzer( + tracker.equity_curve, + trades_df, + risk_free_rate=0.03, + ).compute() + + # Config snapshot + config = { + "cash": self._cash, + "commission": self._commission, + "execution": self._execution, + "position_mode": self._position_mode, + "reject_policy": self._reject_policy, + "future_leak_warning": simulator.future_leak_warning, + } + + return BacktestResult( + performance=performance, + equity_curve=tracker.equity_curve, + trades=trades_df, + positions=tracker.positions, + config=config, + ) + + def _generate_signals(self, df: pd.DataFrame, chanlun_result: Any | None) -> list[Signal]: + """Generate signals from strategy. + + Args: + df: Price data + chanlun_result: Optional chanlun analysis result + + Returns: + List of signals + """ + # Instantiate strategy if needed + strat = ( + self._strategy_instance if self._strategy_instance is not None else self._strategy_cls() + ) + + # Bind data + strat._bind_data(df) + + # Inject chanlun result if provided + if chanlun_result is not None: + strat._chanlun_result = chanlun_result + + # Call init + strat._call_init() + + # Generate signals bar by bar + all_signals: list[Signal] = [] + for i in range(len(df)): + strat._set_bar_index(i) + strat._call_next() + bar_signals = strat._clear_signals() + all_signals.extend(bar_signals) + + return all_signals + + def _compute_pnls(self, trades: list[Trade]) -> list[Trade]: + """Compute realized PnL for sell trades. + + Args: + trades: List of trades + + Returns: + Trades with PnL computed + """ + position_cost = 0.0 + position_size = 0.0 + + for trade in trades: + if not trade.rejected: + if trade.direction == "BUY": + position_cost += trade.size * trade.price + trade.commission + position_size += trade.size + trade.pnl = 0.0 + elif trade.direction == "SELL": + if position_size > 0: + avg_cost = position_cost / position_size + trade.pnl = (trade.price - avg_cost) * trade.size - trade.commission + else: + trade.pnl = 0.0 + position_cost -= avg_cost * trade.size + position_size -= trade.size + + return trades + + def _trades_to_df(self, trades: list[Trade]) -> pd.DataFrame: + """Convert trades to DataFrame. + + Args: + trades: List of trades + + Returns: + DataFrame with trade data + """ + if not trades: + return pd.DataFrame( + columns=[ + "datetime", + "direction", + "size", + "price", + "commission", + "pnl", + "rejected", + ] + ) + + data = [ + { + "datetime": t.datetime, + "direction": t.direction, + "size": t.size, + "price": t.price, + "commission": t.commission, + "pnl": t.pnl, + "rejected": t.rejected, + } + for t in trades + ] + return pd.DataFrame(data) + + def _empty_result(self) -> BacktestResult: + """Return empty result for empty input. + + Returns: + BacktestResult with empty DataFrames + """ + return BacktestResult( + performance={}, + equity_curve=pd.DataFrame( + columns=["datetime", "cash", "position_value", "total", "drawdown"] + ), + trades=pd.DataFrame( + columns=[ + "datetime", + "direction", + "size", + "price", + "commission", + "pnl", + "rejected", + ] + ), + positions=pd.DataFrame( + columns=["datetime", "size", "avg_price", "market_value", "unrealized_pnl"] + ), + config={}, + ) diff --git a/src/easy_tdx/backtest/types.py b/src/easy_tdx/backtest/types.py index eba4391..6c208a4 100644 --- a/src/easy_tdx/backtest/types.py +++ b/src/easy_tdx/backtest/types.py @@ -121,7 +121,19 @@ class BacktestResult: def to_json(self) -> str: """将结果序列化为 JSON 字符串。""" - return json.dumps(self.to_dict(), ensure_ascii=False, indent=2) + d = self.to_dict() + return json.dumps(d, ensure_ascii=False, indent=2, default=self._json_default) + + @staticmethod + def _json_default(obj: Any) -> Any: + """JSON serializer for objects not serializable by default json code.""" + if hasattr(obj, "item"): + # numpy types + return obj.item() + if hasattr(obj, "isoformat"): + # datetime/timestamp objects + return obj.isoformat() + raise TypeError(f"Object of type {type(obj)} is not JSON serializable") def summary(self) -> None: """打印回测概要(标准输出)。""" diff --git a/tests/unit/test_backtest_engine.py b/tests/unit/test_backtest_engine.py new file mode 100644 index 0000000..eea2c0e --- /dev/null +++ b/tests/unit/test_backtest_engine.py @@ -0,0 +1,288 @@ +"""Test BacktestEngine orchestration.""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + +from easy_tdx import MyTT +from easy_tdx.backtest.engine import BacktestEngine +from easy_tdx.backtest.strategy import Strategy + + +def _make_df(n: int = 100, seed: int = 42) -> pd.DataFrame: + """Generate synthetic OHLCV data.""" + rng = np.random.default_rng(seed) + close = 100.0 + np.cumsum(rng.normal(0, 1, n)) + high = close + rng.uniform(0, 1, n) + low = close - rng.uniform(0, 1, n) + open_ = low + rng.uniform(0, high - low, n) + volume = rng.integers(1000000, 10000000, n) + + dates = pd.date_range("2024-01-01", periods=n, freq="D") + return pd.DataFrame( + { + "datetime": dates, + "open": open_, + "high": high, + "low": low, + "close": close, + "volume": volume, + } + ) + + +class MACrossStrategy(Strategy): + """Simple MA crossover strategy.""" + + def init(self): + self.ma5 = self.I(MyTT.MA, self.data.close, 5) + self.ma20 = self.I(MyTT.MA, self.data.close, 20) + self.cross_up = False + self.cross_down = False + + def next(self): + # Check if crossing happened on this bar + if self._bar_index > 0: + prev_ma5 = self.ma5[self._bar_index - 1] + prev_ma20 = self.ma20[self._bar_index - 1] + curr_ma5 = self.ma5[self._bar_index] + curr_ma20 = self.ma20[self._bar_index] + + # Golden cross: ma5 crosses above ma20 + if prev_ma5 <= prev_ma20 and curr_ma5 > curr_ma20: + self.buy(size=0) + # Death cross: ma5 crosses below ma20 + elif prev_ma5 >= prev_ma20 and curr_ma5 < curr_ma20: + self.sell(size=0) + + +class FixedBuyStrategy(Strategy): + """Strategy with fixed buy/sell at specific bars.""" + + def init(self): + pass + + def next(self): + if self._bar_index == 5: + self.buy(size=100) + if self._bar_index == 50: + self.sell(size=100) + + +class ChanlunStrategy(Strategy): + """Strategy that uses chanlun result.""" + + def init(self): + pass + + def next(self): + if self._bar_index == 10 and hasattr(self, "chanlun"): + # Access chanlun result + _ = self.chanlun + self.buy(size=50) + + +class PrecomputedIndicatorStrategy(Strategy): + """Strategy that uses precomputed indicator columns.""" + + def init(self): + # Assume BOLL_UPPER already exists in df + if hasattr(self.data, "BOLL_UPPER"): + self.boll_upper = self.data.BOLL_UPPER + else: + self.boll_upper = None + + def next(self): + if self.boll_upper is not None and self._bar_index == 20: + _ = self.boll_upper[self._bar_index] + self.buy(size=10) + + +def test_basic_run(): + """Test basic engine run with MACrossStrategy.""" + df = _make_df(n=200) + engine = BacktestEngine(MACrossStrategy, cash=100000) + result = engine.run(df) + + # Check performance metrics + assert result.performance is not None + assert "total_return" in result.performance + + # Check equity curve length + assert len(result.equity_curve) == 200 + + # Check columns + assert "datetime" in result.equity_curve.columns + assert "total" in result.equity_curve.columns + + +def test_fixed_strategy(): + """Test FixedBuyStrategy produces trades.""" + df = _make_df(n=100) + engine = BacktestEngine(FixedBuyStrategy, cash=100000) + result = engine.run(df) + + # Should have at least 2 trades + assert len(result.trades) >= 2, f"Expected at least 2 trades, got {len(result.trades)}" + + # Check buy at bar 5 + buy_trades = result.trades[result.trades["direction"] == "BUY"] + assert len(buy_trades) >= 1, "No buy trades found" + + # Check sell at bar 50 + sell_trades = result.trades[result.trades["direction"] == "SELL"] + assert len(sell_trades) >= 1, "No sell trades found" + + +def test_result_columns(): + """Test BacktestResult has correct columns.""" + df = _make_df(n=100) + engine = BacktestEngine(MACrossStrategy) + result = engine.run(df) + + # Equity curve columns + expected_ec_cols = ["datetime", "cash", "position_value", "total"] + for col in expected_ec_cols: + assert col in result.equity_curve.columns + + # Trades columns + expected_trade_cols = ["datetime", "direction", "size", "price", "pnl"] + for col in expected_trade_cols: + assert col in result.trades.columns + + +def test_to_dict(): + """Test BacktestResult is serializable.""" + df = _make_df(n=50) + engine = BacktestEngine(MACrossStrategy) + result = engine.run(df) + + # to_dict should not raise + d = result.to_dict() + assert "performance" in d + assert "equity_curve" in d + assert "trades" in d + + # to_json should not raise + json_str = result.to_json() + assert len(json_str) > 0 + + +def test_chanlun_injection(): + """Test chanlun result injection.""" + df = _make_df(n=50) + + # Mock chanlun result + chanlun_result = {"test": "data"} + + engine = BacktestEngine(ChanlunStrategy) + result = engine.run(df, chanlun_result=chanlun_result) + + # Should have trades + assert len(result.trades) >= 1 + + +def test_this_close_warning_in_config(): + """Test future_leak_warning in config when using this_close.""" + df = _make_df(n=50) + engine = BacktestEngine(MACrossStrategy, execution="this_close") + result = engine.run(df) + + # Config should have future_leak_warning + # Note: MACrossStrategy may not generate signals, so warning might be False + assert "future_leak_warning" in result.config + + +def test_config_snapshot(): + """Test config contains correct cash and commission.""" + df = _make_df(n=50) + engine = BacktestEngine(MACrossStrategy, cash=50000, commission=0.0005, execution="next_open") + result = engine.run(df) + + # Check config + assert result.config["cash"] == 50000 + assert result.config["commission"] == 0.0005 + assert result.config["execution"] == "next_open" + + +def test_precomputed_indicator_columns(): + """Test strategy works with precomputed indicator columns.""" + df = _make_df(n=50) + # Add precomputed BOLL_UPPER column + df["BOLL_UPPER"] = df["close"] * 1.05 + + engine = BacktestEngine(PrecomputedIndicatorStrategy) + result = engine.run(df) + + # Should not crash and should have trades + assert len(result.equity_curve) == 50 + + +def test_empty_df(): + """Test engine with empty DataFrame.""" + df = pd.DataFrame(columns=["datetime", "open", "high", "low", "close", "volume"]) + engine = BacktestEngine(MACrossStrategy) + result = engine.run(df) + + # Should return empty result + assert len(result.equity_curve) == 0 + assert len(result.trades) == 0 + + +def test_strategy_instance_vs_class(): + """Test engine accepts both strategy class and instance.""" + df = _make_df(n=50) + + # Test with class + engine1 = BacktestEngine(MACrossStrategy) + result1 = engine1.run(df) + assert len(result1.equity_curve) == 50 + + # Test with instance + strat = MACrossStrategy() + engine2 = BacktestEngine(strat) + result2 = engine2.run(df) + assert len(result2.equity_curve) == 50 + + +def test_commission_calculation(): + """Test commission is correctly applied.""" + df = _make_df(n=100) + engine = BacktestEngine( + FixedBuyStrategy, + cash=100000, + commission=0.001, + min_commission=10.0, + ) + result = engine.run(df) + + # Should have trades + assert len(result.trades) >= 2 + + # Check trades have commission + assert (result.trades["commission"] > 0).all() + + +def test_pnl_calculation(): + """Test PnL is calculated for sell trades.""" + df = _make_df(n=100, seed=123) # Use specific seed for predictable prices + engine = BacktestEngine(FixedBuyStrategy, cash=100000, commission=0.0) + result = engine.run(df) + + # Should have trades + assert len(result.trades) >= 2 + + # Get trades - should have at least one BUY and one SELL + buy_trades = result.trades[result.trades["direction"] == "BUY"] + sell_trades = result.trades[result.trades["direction"] == "SELL"] + + assert len(buy_trades) >= 1 + assert len(sell_trades) >= 1 + + # PnL is calculated for sell trades + # Check that sell trades have PnL computed + assert (sell_trades["pnl"] != 0).any() or len(sell_trades) == 0 + + # For buy trades, PnL should be 0 + assert (buy_trades["pnl"] == 0).all()