feat(backtest): add TWAPExecution + VWAPExecution

src/easy_tdx/backtest/execution.py

@@ -185,3 +185,339 @@ class ImmediateExecution(ExecutionModel):
             ]
 
         return []
+
+
+class TWAPExecution(ExecutionModel):
+    """时间加权平均价格执行。
+
+    将订单均匀拆分为 n_bars 份，在连续 n_bars 根 K 线上执行。
+    """
+
+    def __init__(self, n_bars: int = 5) -> None:
+        self._n_bars = max(1, n_bars)
+
+    def execute(
+        self,
+        signal: Signal,
+        df: pd.DataFrame,
+        bar_idx: int,
+        cash: float,
+        position: float,
+        position_mode: str,
+        commission: float,
+        min_commission: float,
+        stamp_tax: float,
+        slippage_model: SlippageModel | None,
+    ) -> list[Trade]:
+        if signal.direction == "BUY":
+            return self._execute_buy(
+                signal,
+                df,
+                bar_idx,
+                cash,
+                position_mode,
+                commission,
+                min_commission,
+                stamp_tax,
+                slippage_model,
+            )
+        return self._execute_sell(
+            signal,
+            df,
+            bar_idx,
+            position,
+            commission,
+            min_commission,
+            stamp_tax,
+            slippage_model,
+        )
+
+    def _execute_buy(
+        self,
+        signal: Signal,
+        df: pd.DataFrame,
+        bar_idx: int,
+        cash: float,
+        position_mode: str,
+        commission: float,
+        min_commission: float,
+        stamp_tax: float,
+        slippage_model: SlippageModel | None,
+    ) -> list[Trade]:
+        first_price = float(df["open"].iloc[bar_idx + 1]) if bar_idx + 1 < len(df) else 0
+        if first_price <= 0:
+            return []
+        total_size = self._calc_buy_size(
+            signal.size,
+            first_price,
+            cash,
+            position_mode,
+            commission,
+        )
+        if total_size <= 0:
+            return []
+
+        sub_size = int(total_size / self._n_bars / 100) * 100
+        if sub_size <= 0:
+            sub_size = 100
+
+        trades: list[Trade] = []
+        for i in range(self._n_bars):
+            exec_idx = bar_idx + 1 + i
+            if exec_idx >= len(df):
+                break
+            price = float(df["close"].iloc[exec_idx])
+            remaining = total_size - sum(t.size for t in trades)
+            actual_size = min(sub_size, remaining)
+            actual_size = int(actual_size / 100) * 100
+            if actual_size <= 0:
+                break
+            comm = self._calc_commission(
+                actual_size,
+                price,
+                False,
+                commission,
+                min_commission,
+                stamp_tax,
+            )
+            slip = self._calc_slippage(actual_size, price, False, slippage_model, df)
+            trades.append(
+                Trade(
+                    datetime=self._get_datetime_int(df, exec_idx),
+                    direction="BUY",
+                    size=float(actual_size),
+                    price=price,
+                    commission=comm,
+                    slippage=slip,
+                )
+            )
+        return trades
+
+    def _execute_sell(
+        self,
+        signal: Signal,
+        df: pd.DataFrame,
+        bar_idx: int,
+        position: float,
+        commission: float,
+        min_commission: float,
+        stamp_tax: float,
+        slippage_model: SlippageModel | None,
+    ) -> list[Trade]:
+        total_size = signal.size if signal.size > 0 else position
+        if total_size <= 0:
+            return []
+
+        sub_size = int(total_size / self._n_bars / 100) * 100
+        if sub_size <= 0:
+            sub_size = 100
+
+        trades: list[Trade] = []
+        for i in range(self._n_bars):
+            exec_idx = bar_idx + 1 + i
+            if exec_idx >= len(df):
+                break
+            price = float(df["close"].iloc[exec_idx])
+            remaining = total_size - sum(t.size for t in trades)
+            actual_size = min(sub_size, remaining)
+            actual_size = int(actual_size / 100) * 100
+            if actual_size <= 0:
+                break
+            comm = self._calc_commission(
+                actual_size,
+                price,
+                True,
+                commission,
+                min_commission,
+                stamp_tax,
+            )
+            slip = self._calc_slippage(actual_size, price, True, slippage_model, df)
+            trades.append(
+                Trade(
+                    datetime=self._get_datetime_int(df, exec_idx),
+                    direction="SELL",
+                    size=float(actual_size),
+                    price=price,
+                    commission=comm,
+                    slippage=slip,
+                )
+            )
+        return trades
+
+
+class VWAPExecution(ExecutionModel):
+    """成交量加权平均价格执行。
+
+    按历史成交量分布比例拆分订单。
+    """
+
+    def __init__(self, n_bars: int = 5, volume_lookback: int = 20) -> None:
+        self._n_bars = max(1, n_bars)
+        self._volume_lookback = max(1, volume_lookback)
+
+    def execute(
+        self,
+        signal: Signal,
+        df: pd.DataFrame,
+        bar_idx: int,
+        cash: float,
+        position: float,
+        position_mode: str,
+        commission: float,
+        min_commission: float,
+        stamp_tax: float,
+        slippage_model: SlippageModel | None,
+    ) -> list[Trade]:
+        if signal.direction == "BUY":
+            return self._execute_buy(
+                signal,
+                df,
+                bar_idx,
+                cash,
+                position_mode,
+                commission,
+                min_commission,
+                stamp_tax,
+                slippage_model,
+            )
+        return self._execute_sell(
+            signal,
+            df,
+            bar_idx,
+            position,
+            commission,
+            min_commission,
+            stamp_tax,
+            slippage_model,
+        )
+
+    def _get_volume_weights(self, df: pd.DataFrame, bar_idx: int) -> list[float]:
+        """获取成交量权重分布。"""
+        start = max(0, bar_idx - self._volume_lookback + 1)
+        lookback = df.iloc[start : bar_idx + 1]
+        if "volume" not in lookback.columns or len(lookback) == 0:
+            return [1.0 / self._n_bars] * self._n_bars
+
+        volumes = lookback["volume"].to_numpy()
+        total_vol = float(volumes.sum())
+        if total_vol <= 0:
+            return [1.0 / self._n_bars] * self._n_bars
+
+        weights: list[float] = []
+        for i in range(self._n_bars):
+            idx = max(0, len(volumes) - 1 - (i % max(1, len(volumes))))
+            weights.append(float(volumes[idx]) / total_vol)
+        total_w = sum(weights)
+        if total_w <= 0:
+            return [1.0 / self._n_bars] * self._n_bars
+        return [w / total_w for w in weights]
+
+    def _execute_buy(
+        self,
+        signal: Signal,
+        df: pd.DataFrame,
+        bar_idx: int,
+        cash: float,
+        position_mode: str,
+        commission: float,
+        min_commission: float,
+        stamp_tax: float,
+        slippage_model: SlippageModel | None,
+    ) -> list[Trade]:
+        first_price = float(df["open"].iloc[bar_idx + 1]) if bar_idx + 1 < len(df) else 0
+        if first_price <= 0:
+            return []
+        total_size = self._calc_buy_size(
+            signal.size,
+            first_price,
+            cash,
+            position_mode,
+            commission,
+        )
+        if total_size <= 0:
+            return []
+
+        weights = self._get_volume_weights(df, bar_idx)
+        trades: list[Trade] = []
+        for i in range(self._n_bars):
+            exec_idx = bar_idx + 1 + i
+            if exec_idx >= len(df):
+                break
+            price = float(df["close"].iloc[exec_idx])
+            w = weights[i] if i < len(weights) else 1.0 / self._n_bars
+            target = int(total_size * w / 100) * 100
+            remaining = total_size - sum(t.size for t in trades)
+            actual_size = min(target, remaining)
+            actual_size = int(actual_size / 100) * 100
+            if actual_size <= 0:
+                continue
+            comm = self._calc_commission(
+                actual_size,
+                price,
+                False,
+                commission,
+                min_commission,
+                stamp_tax,
+            )
+            slip = self._calc_slippage(actual_size, price, False, slippage_model, df)
+            trades.append(
+                Trade(
+                    datetime=self._get_datetime_int(df, exec_idx),
+                    direction="BUY",
+                    size=float(actual_size),
+                    price=price,
+                    commission=comm,
+                    slippage=slip,
+                )
+            )
+        return trades
+
+    def _execute_sell(
+        self,
+        signal: Signal,
+        df: pd.DataFrame,
+        bar_idx: int,
+        position: float,
+        commission: float,
+        min_commission: float,
+        stamp_tax: float,
+        slippage_model: SlippageModel | None,
+    ) -> list[Trade]:
+        total_size = signal.size if signal.size > 0 else position
+        if total_size <= 0:
+            return []
+
+        weights = self._get_volume_weights(df, bar_idx)
+        trades: list[Trade] = []
+        for i in range(self._n_bars):
+            exec_idx = bar_idx + 1 + i
+            if exec_idx >= len(df):
+                break
+            price = float(df["close"].iloc[exec_idx])
+            w = weights[i] if i < len(weights) else 1.0 / self._n_bars
+            target = int(total_size * w / 100) * 100
+            remaining = total_size - sum(t.size for t in trades)
+            actual_size = min(target, remaining)
+            actual_size = int(actual_size / 100) * 100
+            if actual_size <= 0:
+                continue
+            comm = self._calc_commission(
+                actual_size,
+                price,
+                True,
+                commission,
+                min_commission,
+                stamp_tax,
+            )
+            slip = self._calc_slippage(actual_size, price, True, slippage_model, df)
+            trades.append(
+                Trade(
+                    datetime=self._get_datetime_int(df, exec_idx),
+                    direction="SELL",
+                    size=float(actual_size),
+                    price=price,
+                    commission=comm,
+                    slippage=slip,
+                )
+            )
+        return trades

tests/unit/test_backtest_execution.py

@@ -5,7 +5,12 @@ from __future__ import annotations
 import pandas as pd
 import pytest
 
-from easy_tdx.backtest.execution import ExecutionModel, ImmediateExecution
+from easy_tdx.backtest.execution import (
+    ExecutionModel,
+    ImmediateExecution,
+    TWAPExecution,
+    VWAPExecution,
+)
 from easy_tdx.backtest.types import Signal
 
 
@@ -167,3 +172,150 @@ class TestImmediateExecution:
         )
         assert len(trades) == 1
         assert trades[0].size == 100
+
+
+class TestTWAPExecution:
+    """时间加权平均价格执行。"""
+
+    def test_split_buy_into_3_bars(self) -> None:
+        df = _make_df(20)
+        model = TWAPExecution(n_bars=3)
+        signal = Signal(datetime=20240101, direction="BUY", size=300)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=0,
+            cash=100000,
+            position=0,
+            position_mode="fixed",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) == 3
+        total_size = sum(t.size for t in trades)
+        assert total_size <= 300
+        prices = [t.price for t in trades]
+        assert prices[0] != prices[1]
+
+    def test_split_sell_into_2_bars(self) -> None:
+        df = _make_df(20)
+        model = TWAPExecution(n_bars=2)
+        signal = Signal(datetime=20240101, direction="SELL", size=200)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=0,
+            cash=0,
+            position=500,
+            position_mode="fixed",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) == 2
+        assert sum(t.size for t in trades) == 200.0
+
+    def test_truncates_at_data_end(self) -> None:
+        df = _make_df(5)
+        model = TWAPExecution(n_bars=10)
+        signal = Signal(datetime=20240101, direction="BUY", size=1000)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=0,
+            cash=100000,
+            position=0,
+            position_mode="fixed",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) <= 4
+
+    def test_full_position_mode(self) -> None:
+        df = _make_df(20)
+        model = TWAPExecution(n_bars=3)
+        signal = Signal(datetime=20240101, direction="BUY", size=0)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=0,
+            cash=60000,
+            position=0,
+            position_mode="full",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) == 3
+        assert all(t.size > 0 for t in trades)
+
+
+class TestVWAPExecution:
+    """成交量加权平均价格执行。"""
+
+    def test_basic_buy(self) -> None:
+        df = _make_df(20)
+        model = VWAPExecution(n_bars=3, volume_lookback=10)
+        signal = Signal(datetime=20240101, direction="BUY", size=300)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=5,
+            cash=100000,
+            position=0,
+            position_mode="fixed",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) == 3
+        total_size = sum(t.size for t in trades)
+        assert total_size <= 300
+
+    def test_volume_weighted_split(self) -> None:
+        df = _make_df(20)
+        df.loc[6, "volume"] = 50000
+        df.loc[7, "volume"] = 50000
+        df.loc[8, "volume"] = 50000
+        model = VWAPExecution(n_bars=3, volume_lookback=5)
+        signal = Signal(datetime=20240105, direction="BUY", size=300)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=5,
+            cash=100000,
+            position=0,
+            position_mode="fixed",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) == 3
+        sizes = [t.size for t in trades]
+        assert sum(sizes) <= 300
+
+    def test_truncates_at_data_end(self) -> None:
+        df = _make_df(5)
+        model = VWAPExecution(n_bars=10, volume_lookback=3)
+        signal = Signal(datetime=20240101, direction="BUY", size=1000)
+        trades = model.execute(
+            signal=signal,
+            df=df,
+            bar_idx=0,
+            cash=100000,
+            position=0,
+            position_mode="fixed",
+            commission=0.0003,
+            min_commission=5.0,
+            stamp_tax=0.001,
+            slippage_model=None,
+        )
+        assert len(trades) <= 4