Merge branch 'pr/153' into development

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.
+    """