mirror of
https://ghfast.top/https://github.com/aeroxw/easy-tdx.git
synced 2026-09-12 21:34:16 +08:00
feat(backtest): add PortfolioTracker with equity curve and drawdown
- Pre-allocate numpy arrays for performance (cash, position, avg_price) - apply_trades() processes buys/sells with commission and slippage - equity_curve returns DataFrame with drawdown calculation - positions returns DataFrame with market value and unrealized PnL - 12 unit tests covering all scenarios Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
16dc2e7da9
commit
a2aa319803
@@ -0,0 +1,153 @@
|
||||
"""回测引擎持仓追踪器。
|
||||
|
||||
负责应用交易记录、计算资金曲线和持仓快照。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from easy_tdx.backtest.types import Trade
|
||||
|
||||
|
||||
class PortfolioTracker:
|
||||
"""持仓追踪器。
|
||||
|
||||
预分配 numpy 数组存储每个 bar 的状态,遍历应用交易。
|
||||
|
||||
Attributes:
|
||||
_close: 收盘价数组
|
||||
_datetime: 时间戳数组
|
||||
_n: bar 数量
|
||||
_cash: 每个 bar 的现金
|
||||
_position: 每个 bar 的持仓数量
|
||||
_avg_price: 每个 bar 的平均成本价
|
||||
"""
|
||||
|
||||
def __init__(self, df: pd.DataFrame, initial_cash: float = 100000) -> None:
|
||||
"""初始化追踪器。
|
||||
|
||||
Args:
|
||||
df: 必须包含 close 和 datetime 列的 DataFrame
|
||||
initial_cash: 初始现金
|
||||
"""
|
||||
self._close = df["close"].to_numpy(dtype=float)
|
||||
self._datetime = df["datetime"].to_numpy()
|
||||
self._n = len(df)
|
||||
self._cash = np.full(self._n, initial_cash)
|
||||
self._position = np.zeros(self._n)
|
||||
self._avg_price = np.zeros(self._n)
|
||||
self._initial_cash = initial_cash
|
||||
|
||||
def apply_trades(self, trades: list[Trade]) -> None:
|
||||
"""应用交易记录,更新内部状态数组。
|
||||
|
||||
Args:
|
||||
trades: 交易列表
|
||||
"""
|
||||
# 构建 datetime → Trade 映射
|
||||
buy_map: dict[int, Trade] = {}
|
||||
sell_map: dict[int, Trade] = {}
|
||||
|
||||
for trade in trades:
|
||||
if trade.rejected:
|
||||
continue
|
||||
if trade.direction == "BUY":
|
||||
buy_map[trade.datetime] = trade
|
||||
else:
|
||||
sell_map[trade.datetime] = trade
|
||||
|
||||
# 遍历每个 bar
|
||||
for i in range(self._n):
|
||||
dt = self._datetime[i]
|
||||
|
||||
# 继承前一个 bar 的状态(除了第一个 bar)
|
||||
if i > 0:
|
||||
self._cash[i] = self._cash[i - 1]
|
||||
self._position[i] = self._position[i - 1]
|
||||
self._avg_price[i] = self._avg_price[i - 1]
|
||||
|
||||
# 处理买入
|
||||
if dt in buy_map:
|
||||
trade = buy_map[dt]
|
||||
cost = trade.size * trade.price + trade.commission + trade.slippage
|
||||
self._cash[i] -= cost
|
||||
|
||||
# 更新均价
|
||||
if self._position[i] > 0:
|
||||
total_cost = self._position[i] * self._avg_price[i] + trade.size * trade.price
|
||||
self._position[i] += trade.size
|
||||
self._avg_price[i] = total_cost / self._position[i]
|
||||
else:
|
||||
# 新开仓或从空仓开仓
|
||||
self._position[i] = trade.size
|
||||
self._avg_price[i] = trade.price
|
||||
|
||||
# 处理卖出
|
||||
if dt in sell_map:
|
||||
trade = sell_map[dt]
|
||||
proceeds = trade.size * trade.price - trade.commission - trade.slippage
|
||||
self._cash[i] += proceeds
|
||||
self._position[i] -= trade.size
|
||||
|
||||
# 清空持仓时归零
|
||||
if self._position[i] <= 0:
|
||||
self._position[i] = 0.0
|
||||
self._avg_price[i] = 0.0
|
||||
|
||||
@property
|
||||
def equity_curve(self) -> pd.DataFrame:
|
||||
"""计算资金曲线。
|
||||
|
||||
Returns:
|
||||
DataFrame 包含 datetime, cash, position_value, total, drawdown, drawdown_pct
|
||||
"""
|
||||
position_value = self._position * self._close
|
||||
total = self._cash + position_value
|
||||
peak = np.maximum.accumulate(total)
|
||||
drawdown = peak - total
|
||||
drawdown_pct = np.divide(
|
||||
drawdown,
|
||||
peak,
|
||||
out=np.zeros_like(drawdown),
|
||||
where=(peak != 0),
|
||||
)
|
||||
|
||||
return pd.DataFrame(
|
||||
{
|
||||
"datetime": self._datetime,
|
||||
"cash": self._cash,
|
||||
"position_value": position_value,
|
||||
"total": total,
|
||||
"drawdown": drawdown,
|
||||
"drawdown_pct": drawdown_pct,
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def positions(self) -> pd.DataFrame:
|
||||
"""计算持仓快照。
|
||||
|
||||
Returns:
|
||||
DataFrame 包含 datetime, size, avg_price, market_value, unrealized_pnl
|
||||
"""
|
||||
market_value = self._position * self._close
|
||||
unrealized_pnl = (self._close - self._avg_price) * self._position
|
||||
|
||||
return pd.DataFrame(
|
||||
{
|
||||
"datetime": self._datetime,
|
||||
"size": self._position,
|
||||
"avg_price": self._avg_price,
|
||||
"market_value": market_value,
|
||||
"unrealized_pnl": unrealized_pnl,
|
||||
}
|
||||
)
|
||||
|
||||
@property
|
||||
def initial_cash(self) -> float:
|
||||
"""初始现金。"""
|
||||
return self._initial_cash
|
||||
Reference in New Issue
Block a user