58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
from __future__ import annotations
|
|
import pandas as pd
|
|
import numpy as np
|
|
|
|
|
|
def _atr(high: pd.Series, low: pd.Series, close: pd.Series, period: int) -> pd.Series:
|
|
hl = (high - low).abs()
|
|
hc = (high - close.shift()).abs()
|
|
lc = (low - close.shift()).abs()
|
|
tr = pd.concat([hl, hc, lc], axis=1).max(axis=1)
|
|
return tr.rolling(period, min_periods=period).mean()
|
|
|
|
|
|
def supertrend_series(df: pd.DataFrame, length: int, multiplier: float) -> pd.Series:
|
|
atr = _atr(df["High"], df["Low"], df["Close"], length)
|
|
hl2 = (df["High"] + df["Low"]) / 2
|
|
upper = hl2 + multiplier * atr
|
|
lower = hl2 - multiplier * atr
|
|
|
|
trend = pd.Series(index=df.index, dtype=float)
|
|
dir_up = True
|
|
prev_upper = np.nan
|
|
prev_lower = np.nan
|
|
|
|
for i in range(len(df)):
|
|
if i == 0 or pd.isna(atr.iat[i]):
|
|
trend.iat[i] = np.nan
|
|
prev_upper = upper.iat[i]
|
|
prev_lower = lower.iat[i]
|
|
continue
|
|
|
|
cu = min(upper.iat[i], prev_upper) if dir_up else upper.iat[i]
|
|
cl = max(lower.iat[i], prev_lower) if not dir_up else lower.iat[i]
|
|
|
|
if df["Close"].iat[i] > cu:
|
|
dir_up = True
|
|
elif df["Close"].iat[i] < cl:
|
|
dir_up = False
|
|
|
|
prev_upper = cu if dir_up else upper.iat[i]
|
|
prev_lower = lower.iat[i] if dir_up else cl
|
|
trend.iat[i] = cl if dir_up else cu
|
|
|
|
return trend
|
|
|
|
|
|
def add_supertrends(df: pd.DataFrame, settings: list[tuple[int, float]]) -> pd.DataFrame:
|
|
out = df.copy()
|
|
for length, mult in settings:
|
|
col = f"supertrend_{length}_{mult}"
|
|
out[col] = supertrend_series(out, length, mult)
|
|
out[f"bull_{length}_{mult}"] = (out["Close"] >= out[col]).astype(int)
|
|
return out
|
|
|
|
|
|
def compute_meta_trend(df: pd.DataFrame, settings: list[tuple[int, float]]) -> pd.Series:
|
|
bull_cols = [f"bull_{l}_{m}" for l, m in settings]
|
|
return (df[bull_cols].sum(axis=1) == len(bull_cols)).astype(int) |