mirror of
https://ghfast.top/https://github.com/aeroxw/easy-tdx.git
synced 2026-09-12 14:34:15 +08:00
feat(backtest): add BacktestEngine with vectorized execution pipeline
- Implement BacktestEngine orchestrator with 4-step pipeline: 1. Signal generation (Strategy) 2. Order simulation (OrderSimulator) 3. Portfolio tracking (PortfolioTracker) 4. Performance analysis (PerformanceAnalyzer) - Support both strategy class and instance initialization - Add PnL calculation for sell trades - Add JSON serialization with numpy/timestamp support - Include comprehensive test coverage (12 tests, all passing) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
94fabccef8
commit
371915a5f9
@@ -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={},
|
||||
)
|
||||
@@ -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:
|
||||
"""打印回测概要(标准输出)。"""
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user