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()
|