Files
2026-07-13 12:36:27 +08:00

256 lines
8.7 KiB
Python

from __future__ import annotations
import math
from dataclasses import dataclass
from datetime import datetime, timedelta
import numpy as np
import pandas as pd
import yfinance as yf
_INITIAL_CAPITAL = 10_000.0
_COST_PER_SIDE = 1.99
_DISTANCE_THRESH = 6.0
_WHIPSAW_DAYS = 10
_CASH_YIELD_ANNUAL = 0.04
_CASH_YIELD_DAILY = (1.0 + _CASH_YIELD_ANNUAL) ** (1.0 / 252.0)
_PERIOD_SPECS: dict[str, dict] = {
"2y": {"label": "近 2 年", "fetch_type": "period", "dl_period": "3y", "target_days": 504},
"3y": {"label": "近 3 年", "fetch_type": "period", "dl_period": "4y", "target_days": 756},
"5y": {"label": "近 5 年", "fetch_type": "period", "dl_period": "6y", "target_days": 1260},
"10y": {"label": "近 10 年", "fetch_type": "range", "dl_start": None, "dl_end": None, "target_days": 2520},
"2015_2020": {
"label": "2015—2020",
"fetch_type": "range",
"dl_start": "2014-01-01",
"dl_end": "2020-12-31",
"eff_start": "2015-01-02",
},
"2010_2015": {
"label": "2010—2015",
"fetch_type": "range",
"dl_start": "2009-01-01",
"dl_end": "2015-12-31",
"eff_start": None,
},
}
_PERIOD_ORDER = ["2y", "3y", "5y", "10y", "2015_2020", "2010_2015"]
def _accrue_cash_yield(cash: float) -> float:
return cash * _CASH_YIELD_DAILY
def _fetch_by_period(dl_period: str) -> tuple[pd.DataFrame, pd.DataFrame]:
qqq = yf.download("QQQ", period=dl_period, auto_adjust=True, progress=False)
tqqq = yf.download("TQQQ", period=dl_period, auto_adjust=True, progress=False)
for df in (qqq, tqqq):
if isinstance(df.columns, pd.MultiIndex):
df.columns = df.columns.get_level_values(0)
return qqq, tqqq
def _fetch_by_range(start: str, end: str) -> tuple[pd.DataFrame, pd.DataFrame]:
qqq = yf.download("QQQ", start=start, end=end, auto_adjust=True, progress=False)
tqqq = yf.download("TQQQ", start=start, end=end, auto_adjust=True, progress=False)
for df in (qqq, tqqq):
if isinstance(df.columns, pd.MultiIndex):
df.columns = df.columns.get_level_values(0)
return qqq, tqqq
def _fetch(period_key: str) -> tuple[pd.DataFrame, pd.DataFrame]:
spec = _PERIOD_SPECS[period_key]
if spec["fetch_type"] == "period":
return _fetch_by_period(spec["dl_period"])
start = spec.get("dl_start")
end = spec.get("dl_end")
if start is None:
end_dt = datetime.now()
start_dt = end_dt - timedelta(days=365 * 11)
start = start_dt.strftime("%Y-%m-%d")
end = end_dt.strftime("%Y-%m-%d")
return _fetch_by_range(start, end)
def _get_effective_start(qqq_full: pd.Series, tqqq_full: pd.Series, period_key: str) -> pd.Timestamp:
spec = _PERIOD_SPECS[period_key]
if spec.get("eff_start"):
ts = pd.Timestamp(spec["eff_start"])
return max(ts, qqq_full.index[0])
if period_key == "2010_2015":
return tqqq_full.index[0]
n = spec.get("target_days", 504)
offset = max(0, len(qqq_full) - n)
return qqq_full.index[offset]
def _streak(s: pd.Series) -> pd.Series:
groups = (s != s.shift()).cumsum()
return (s.groupby(groups).cumcount() + 1).where(s, 0)
def _build_signals(qqq_close: pd.Series) -> pd.DataFrame:
ma200 = qqq_close.rolling(200).mean()
a1 = qqq_close > ma200
slope = (ma200 - ma200.shift(20)) / ma200.shift(20) * 100
a2 = slope > -0.5
a3 = _streak(a1) >= 3
signal = a1 & a2 & a3
entry_signal = signal & ~signal.shift(1).fillna(False)
below = qqq_close < ma200
exit_signal = below & below.shift(1).fillna(False)
cross = a1 & ~a1.shift(1).fillna(False)
return pd.DataFrame(
{
"qqq": qqq_close,
"ma200": ma200,
"a1": a1,
"a2": a2,
"a3": a3,
"signal": signal,
"entry_signal": entry_signal,
"exit_signal": exit_signal,
"cross": cross,
}
)
def _allocation(dist: float) -> tuple[str, float]:
if dist <= _DISTANCE_THRESH:
return "Direct Buy", 2 / 3
if dist <= 15.0:
return "Split", 1 / 3
return "Split(>15%)", 1 / 3
@dataclass
class Trade:
entry_date: pd.Timestamp
entry_price: float
entry_mode: str
shares: float
cost_buy: float
exit_date: pd.Timestamp | None = None
exit_price: float | None = None
cost_sell: float = 0.0
forced_exit: bool = False
@property
def holding_days(self) -> int:
if self.exit_date is None:
return 0
return (self.exit_date - self.entry_date).days
@property
def gross_pnl(self) -> float:
if self.exit_price is None:
return 0.0
return (self.exit_price - self.entry_price) * self.shares
@property
def net_pnl(self) -> float:
return self.gross_pnl - self.cost_buy - self.cost_sell
@property
def return_pct(self) -> float:
invested = self.entry_price * self.shares + self.cost_buy
return (self.net_pnl / invested * 100) if invested > 0 else 0.0
@property
def is_whipsaw(self) -> bool:
return self.holding_days <= _WHIPSAW_DAYS and self.net_pnl < 0
def _simulate(sig: pd.DataFrame, tqqq_close: pd.Series, trade_start: pd.Timestamp | None = None) -> tuple[list[Trade], pd.Series]:
tqqq = tqqq_close.reindex(sig.index, method="ffill")
cross_dates = sig.index[sig["cross"]].tolist()
cash = _INITIAL_CAPITAL
shares = 0.0
in_pos = False
current: Trade | None = None
trades: list[Trade] = []
equity: dict = {}
first_day = True
for date in sig.index:
if trade_start is not None and date < trade_start:
continue
cash = _accrue_cash_yield(cash)
raw = tqqq.get(date)
if raw is None or (isinstance(raw, float) and math.isnan(raw)):
equity[date] = cash + shares * (float(tqqq.dropna().iloc[-1]) if shares > 0 else 0.0)
first_day = False
continue
price = float(raw)
if in_pos and current is not None and bool(sig.at[date, "exit_signal"]):
cash += shares * price - _COST_PER_SIDE
current.exit_date = date
current.exit_price = price
current.cost_sell = _COST_PER_SIDE
trades.append(current)
shares = 0.0
in_pos = False
current = None
is_entry = bool(sig.at[date, "entry_signal"])
if first_day and not in_pos and bool(sig.at[date, "signal"]):
is_entry = True
first_day = False
if not in_pos and is_entry:
prior = [d for d in cross_dates if d <= date]
cross_date = prior[-1] if prior else date
cross_raw = tqqq.get(cross_date)
cross_price = float(cross_raw) if cross_raw is not None and not math.isnan(float(cross_raw)) else price
dist = (price - cross_price) / cross_price * 100 if cross_price > 0 else 0.0
mode, frac = _allocation(dist)
invest = cash * frac - _COST_PER_SIDE
if invest > 0:
n = math.floor(invest / price)
if n >= 1:
cash -= n * price + _COST_PER_SIDE
shares = float(n)
in_pos = True
current = Trade(entry_date=date, entry_price=price, entry_mode=mode, shares=shares, cost_buy=_COST_PER_SIDE)
equity[date] = cash + shares * price
if in_pos and current is not None and shares > 0:
last_date = sig.index[-1]
last_price = float(tqqq.iloc[-1])
cash += shares * last_price - _COST_PER_SIDE
current.exit_date = last_date
current.exit_price = last_price
current.cost_sell = _COST_PER_SIDE
current.forced_exit = True
trades.append(current)
equity[last_date] = cash
return trades, pd.Series(equity)
def _max_dd(equity: pd.Series) -> float:
peak = equity.cummax()
return float(((equity - peak) / peak).min() * 100)
def _metrics(trades: list[Trade], equity: pd.Series, qqq_close: pd.Series) -> dict:
n = len(trades)
wins = [t for t in trades if t.net_pnl > 0]
loss = [t for t in trades if t.net_pnl <= 0]
wsaw = [t for t in trades if t.is_whipsaw]
avg = float(np.mean([t.holding_days for t in trades])) if trades else 0.0
return {
"strategy_return": (equity.iloc[-1] / _INITIAL_CAPITAL - 1) * 100,
"qqq_return": (qqq_close.iloc[-1] / qqq_close.iloc[0] - 1) * 100,
"final_value": float(equity.iloc[-1]),
"n_trades": n,
"n_wins": len(wins),
"n_losses": len(loss),
"win_rate": (len(wins) / n * 100) if n > 0 else 0.0,
"max_dd": _max_dd(equity),
"avg_days": avg,
"n_whipsaws": len(wsaw),
"total_cost": sum(t.cost_buy + t.cost_sell for t in trades),
}