256 lines
8.7 KiB
Python
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),
|
|
}
|