feat: multi-stock portfolio backtest engine

- Add PortfolioBacktestEngine for shared-capital multi-stock backtesting
- Support equal allocation mode (total_cash / N per stock)
- Individual BacktestEngine per stock with allocated capital
- Aggregate performance via capital-weighted returns
- Add StockData, PortfolioResult data classes
- Add 4 tests: basic run, equal allocation, empty stocks, serialization
This commit is contained in:
Justin Gu
2026-06-11 02:31:43 +08:00
parent fa5bd2d44d
commit 9c39ad054d
2 changed files with 322 additions and 0 deletions
+206
View File
@@ -0,0 +1,206 @@
"""多标的组合回测引擎。
支持同时回测多只股票,共享资金池,按策略信号分配资金。
每只标的独立产生信号,引擎统一管理仓位和资金。
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import pandas as pd
from easy_tdx.backtest.engine import BacktestEngine
from easy_tdx.backtest.strategy import Strategy
from easy_tdx.backtest.types import BacktestResult
@dataclass
class StockData:
"""单只标的的数据和标识。
Attributes:
code: 股票代码(如 "000001"
market: 市场(如 "SZ"
df: K线 DataFrame
"""
code: str
market: str
df: pd.DataFrame
@dataclass
class PortfolioResult:
"""组合回测结果。
Attributes:
total_performance: 组合整体绩效指标
individual_results: 每只标的的独立回测结果
equity_allocation: 每只标的的资金分配比例
"""
total_performance: dict[str, float]
individual_results: dict[str, BacktestResult]
equity_allocation: dict[str, float]
def to_dict(self) -> dict[str, Any]:
"""转为可序列化字典。"""
return {
"total_performance": self.total_performance,
"individual_results": {k: v.to_dict() for k, v in self.individual_results.items()},
"equity_allocation": self.equity_allocation,
}
class PortfolioBacktestEngine:
"""多标的组合回测引擎。
管理多只股票的共享资金池,独立运行策略,
按均等或自定义比例分配资金。
用法::
engine = PortfolioBacktestEngine(
strategy_cls=MyStrategy,
stocks=[
StockData("000001", "SZ", df1),
StockData("600000", "SH", df2),
],
total_cash=200000,
)
result = engine.run()
print(result.total_performance)
"""
def __init__(
self,
strategy_cls: type[Strategy],
stocks: list[StockData],
total_cash: float = 200_000.0,
allocation: str = "equal",
commission: float = 0.0003,
min_commission: float = 5.0,
stamp_tax: float = 0.001,
slippage: float = 0.0,
execution: str = "next_open",
chanlun_level: str | None = None,
) -> None:
"""初始化组合回测引擎。
Args:
strategy_cls: 策略类
stocks: 标的列表(StockData
total_cash: 总资金
allocation: 资金分配方式
- "equal": 均等分配
- "capitalization": 按市值加权(需额外数据)
commission: 佣金率
min_commission: 最低佣金
stamp_tax: 印花税
slippage: 滑点
execution: 执行模式
chanlun_level: 缠论级别(可选)
"""
self._strategy_cls = strategy_cls
self._stocks = stocks
self._total_cash = total_cash
self._allocation = allocation
self._commission = commission
self._min_commission = min_commission
self._stamp_tax = stamp_tax
self._slippage = slippage
self._execution = execution
self._chanlun_level = chanlun_level
def _compute_allocations(self) -> dict[str, float]:
"""计算每只标的的资金分配。"""
n = len(self._stocks)
if n == 0:
return {}
if self._allocation == "equal":
per_stock_cash = self._total_cash / n
return {f"{s.market}{s.code}": per_stock_cash for s in self._stocks}
# 默认均等分配
per_stock_cash = self._total_cash / n
return {f"{s.market}{s.code}": per_stock_cash for s in self._stocks}
def run(self) -> PortfolioResult:
"""运行组合回测。
对每只标的独立运行回测,按分配的资金量计算收益,
最终汇总为组合整体绩效。
Returns:
PortfolioResult 包含整体绩效和各标的详细结果
"""
allocations = self._compute_allocations()
individual_results: dict[str, BacktestResult] = {}
for stock in self._stocks:
key = f"{stock.market}{stock.code}"
cash = allocations.get(key, 0)
engine = BacktestEngine(
strategy=self._strategy_cls,
cash=cash,
commission=self._commission,
min_commission=self._min_commission,
stamp_tax=self._stamp_tax,
slippage=self._slippage,
execution=self._execution,
chanlun_level=self._chanlun_level,
)
result = engine.run(stock.df)
individual_results[key] = result
# 汇总整体绩效
total_perf = self._aggregate_performance(individual_results, allocations)
# 计算资金占比
total_alloc = sum(allocations.values())
equity_pct = {k: v / total_alloc if total_alloc > 0 else 0 for k, v in allocations.items()}
return PortfolioResult(
total_performance=total_perf,
individual_results=individual_results,
equity_allocation=equity_pct,
)
def _aggregate_performance(
self,
results: dict[str, BacktestResult],
allocations: dict[str, float],
) -> dict[str, float]:
"""汇总所有标的的绩效为组合整体绩效。
使用资金加权方式计算组合收益率。
Args:
results: 各标的回测结果
allocations: 各标的资金分配
Returns:
组合整体绩效指标
"""
total_cash = sum(allocations.values())
if total_cash == 0:
return {"total_return": 0.0, "annual_return": 0.0}
# 资金加权收益率
weighted_return = 0.0
for key, result in results.items():
alloc = allocations.get(key, 0)
weight = alloc / total_cash
ret = result.performance.get("total_return", 0.0)
weighted_return += weight * ret
return {
"total_return": weighted_return,
"annual_return": weighted_return, # 简化,实际应根据周期年化
"total_stocks": len(results),
"total_cash": total_cash,
}
+116
View File
@@ -0,0 +1,116 @@
"""单元测试:多标的组合回测引擎."""
from __future__ import annotations
import numpy as np
import pandas as pd
from easy_tdx.backtest.portfolio_engine import (
PortfolioBacktestEngine,
StockData,
)
from easy_tdx.backtest.strategy import Strategy
class SimpleBuyStrategy(Strategy):
"""简单策略:bar 5 买入,bar 30 卖出."""
def init(self) -> None:
pass
def next(self) -> None:
if self._bar_index == 5 and self.position["size"] == 0:
self.buy(size=0)
elif self._bar_index == 30 and self.position["size"] > 0:
self.sell(size=0)
def _make_df(n: int = 100, seed: int = 42) -> pd.DataFrame:
"""生成随机 OHLCV DataFrame."""
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)
vol = rng.integers(1000000, 10000000, n).astype(float)
return pd.DataFrame(
{
"datetime": pd.date_range("2024-01-01", periods=n, freq="D"),
"open": open_,
"high": high,
"low": low,
"close": close,
"vol": vol,
"amount": vol * close,
}
)
class TestPortfolioBacktest:
"""测试组合回测引擎."""
def test_basic_portfolio_run(self) -> None:
"""基本组合回测应正常完成."""
stocks = [
StockData("000001", "SZ", _make_df(100, seed=42)),
StockData("600000", "SH", _make_df(100, seed=99)),
]
engine = PortfolioBacktestEngine(
strategy_cls=SimpleBuyStrategy,
stocks=stocks,
total_cash=200000,
)
result = engine.run()
assert result.total_performance is not None
assert "total_return" in result.total_performance
assert len(result.individual_results) == 2
assert result.total_performance["total_stocks"] == 2
def test_equal_allocation(self) -> None:
"""均等分配:每只标的资金应为总资金/标的数."""
stocks = [
StockData("000001", "SZ", _make_df(100, seed=42)),
StockData("000002", "SZ", _make_df(100, seed=99)),
]
engine = PortfolioBacktestEngine(
strategy_cls=SimpleBuyStrategy,
stocks=stocks,
total_cash=100000,
allocation="equal",
)
result = engine.run()
# 每只标的分配 50000
assert result.equity_allocation["SZ000001"] == 0.5
assert result.equity_allocation["SZ000002"] == 0.5
def test_empty_stocks(self) -> None:
"""空标的列表应返回零绩效."""
engine = PortfolioBacktestEngine(
strategy_cls=SimpleBuyStrategy,
stocks=[],
total_cash=100000,
)
result = engine.run()
assert result.total_performance["total_return"] == 0.0
assert len(result.individual_results) == 0
def test_to_dict_serializable(self) -> None:
"""结果应可序列化为字典."""
stocks = [StockData("000001", "SZ", _make_df(100))]
engine = PortfolioBacktestEngine(
strategy_cls=SimpleBuyStrategy,
stocks=stocks,
total_cash=100000,
)
result = engine.run()
d = result.to_dict()
assert "total_performance" in d
assert "individual_results" in d
assert "equity_allocation" in d