32 lines
1.3 KiB
Python
32 lines
1.3 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
|
|
def ema(s: pd.Series, n: int) -> pd.Series:
|
|
return s.ewm(span=n, adjust=False).mean()
|
|
|
|
def rsi(s: pd.Series, n: int = 14) -> pd.Series:
|
|
d = s.diff()
|
|
up = d.clip(lower=0).ewm(alpha=1 / n, adjust=False).mean()
|
|
dn = (-d.clip(upper=0)).ewm(alpha=1 / n, adjust=False).mean()
|
|
rs = up / dn.replace(0, np.nan)
|
|
return 100 - 100 / (1 + rs)
|
|
|
|
def atr(df: pd.DataFrame, n: int = 14) -> pd.Series:
|
|
pc = df["close"].shift()
|
|
tr = pd.concat([df["high"] - df["low"], (df["high"] - pc).abs(), (df["low"] - pc).abs()], axis=1).max(axis=1)
|
|
return tr.ewm(alpha=1 / n, adjust=False).mean()
|
|
|
|
def bollinger(s: pd.Series, n: int = 20, dev: float = 2.0):
|
|
m = s.rolling(n).mean()
|
|
sd = s.rolling(n).std()
|
|
return m - dev * sd, m, m + dev * sd
|
|
|
|
def adx(df: pd.DataFrame, n: int = 14) -> pd.Series:
|
|
up, dn = df["high"].diff(), -df["low"].diff()
|
|
plus = pd.Series(np.where((up > dn) & (up > 0), up, 0.0), index=df.index)
|
|
minus = pd.Series(np.where((dn > up) & (dn > 0), dn, 0.0), index=df.index)
|
|
a = atr(df, n)
|
|
pdi = 100 * plus.ewm(alpha=1 / n, adjust=False).mean() / a
|
|
mdi = 100 * minus.ewm(alpha=1 / n, adjust=False).mean() / a
|
|
dx = 100 * (pdi - mdi).abs() / (pdi + mdi).replace(0, np.nan)
|
|
return dx.ewm(alpha=1 / n, adjust=False).mean()
|