diff --git a/pandas_ta/volatility/atr.py b/pandas_ta/volatility/atr.py index 08e92da..c682241 100644 --- a/pandas_ta/volatility/atr.py +++ b/pandas_ta/volatility/atr.py @@ -1,17 +1,19 @@ # -*- coding: utf-8 -*- -from pandas_ta.overlap import ema, rma +from pandas_ta.overlap import ema, rma, sma +from pandas_ta.overlap.wma import wma from .true_range import true_range from pandas_ta.utils import get_drift, get_offset, verify_series -def atr(high, low, close, length=None, mamode=None, drift=None, offset=None, **kwargs): +def atr(high, low, close, length=None, mamode='rma', drift=None, offset=None, **kwargs): """Indicator: Average True Range (ATR)""" # Validate arguments high = verify_series(high) low = verify_series(low) close = verify_series(close) length = int(length) if length and length > 0 else 14 - mamode = mamode.lower() if mamode else "ema" + mamode = str(mamode).lower() + drift = get_drift(drift) offset = get_offset(offset) @@ -20,10 +22,14 @@ def atr(high, low, close, length=None, mamode=None, drift=None, offset=None, **k if mamode == "ema": # alpha = (1.0 / length) if length > 0 else 0.5 # atr = tr.ewm(alpha=alpha).mean() - atr = rma(tr, length=length) - else: + atr = ema(tr, length=length) + elif mamode == "sma": # atr = tr.rolling(length).mean() atr = sma(tr, length=length) + elif mamode == "wma": + atr = wma(tr, length=length) + else: + atr = rma(tr, length=length) percentage = kwargs.pop("percent", False) if percentage: @@ -47,43 +53,43 @@ def atr(high, low, close, length=None, mamode=None, drift=None, offset=None, **k atr.__doc__ = \ -"""Average True Range (ATR) - -Averge True Range is used to measure volatility, especially -volatility caused by gaps or limit moves. - -Sources: - https://www.tradingview.com/wiki/Average_True_Range_(ATR) - -Calculation: - Default Inputs: - length=14, drift=1, percent=False - SMA = Simple Moving Average - EMA = Exponential Moving Average - TR = True Range - tr = TR(high, low, close, drift) - if 'ema': - ATR = EMA(tr, length) - else: - ATR = SMA(tr, length) - - if percent: - ATR *= 100 / close - -Args: - high (pd.Series): Series of 'high's - low (pd.Series): Series of 'low's - close (pd.Series): Series of 'close's - length (int): It's period. Default: 14 - mamode (str): Two options: None or 'ema'. Default: 'ema' - drift (int): The difference period. Default: 1 - offset (int): How many periods to offset the result. Default: 0 - -Kwargs: - percent (bool, optional): Return as percentage. Default: False - fillna (value, optional): pd.DataFrame.fillna(value) - fill_method (value, optional): Type of fill method - -Returns: - pd.Series: New feature generated. -""" + """Average True Range (ATR) + Averge True Range is used to measure volatility, especially + volatility caused by gaps or limit moves. + Sources: + https://www.tradingview.com/wiki/Average_True_Range_(ATR) + Calculation: + Default Inputs: + length=14, drift=1, percent=False + SMA = Simple Moving Average + EMA = Exponential Moving Average + WMA = Weighted Moving Average + RMA = Running Moving Average + TR = True Range + tr = TR(high, low, close, drift) + if 'ema': + ATR = EMA(tr, length) + elif 'sma': + ATR = SMA(tr, length) + elif 'wma': + ATR = WMA(tr, length) + else: + ATR = RMA(tr, length) + + if percent: + ATR *= 100 / close + Args: + high (pd.Series): Series of 'high's + low (pd.Series): Series of 'low's + close (pd.Series): Series of 'close's + length (int): It's period. Default: 14 + mamode (str): 'sma', 'ema', 'wma' or 'rma'. Default: 'rma' + drift (int): The difference period. Default: 1 + offset (int): How many periods to offset the result. Default: 0 + Kwargs: + percent (bool, optional): Return as percentage. Default: False + fillna (value, optional): pd.DataFrame.fillna(value) + fill_method (value, optional): Type of fill method + Returns: + pd.Series: New feature generated. + """