diff --git a/src/easy_tdx/backtest/portfolio_engine.py b/src/easy_tdx/backtest/portfolio_engine.py new file mode 100644 index 0000000..d684945 --- /dev/null +++ b/src/easy_tdx/backtest/portfolio_engine.py @@ -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, + } diff --git a/tests/unit/test_portfolio_engine.py b/tests/unit/test_portfolio_engine.py new file mode 100644 index 0000000..ad5e9cc --- /dev/null +++ b/tests/unit/test_portfolio_engine.py @@ -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