Add Ehlers Predictive Moving Average.

jesse/indicators/__init__.py

@@ -107,6 +107,7 @@ from .obv import obv
 from .pattern_recognition import pattern_recognition
 from .pfe import pfe
 from .pivot import pivot
+from .pma import pma
 from .ppo import ppo
 from .pvi import pvi
 from .pwma import pwma

jesse/indicators/pma.py

@@ -0,0 +1,51 @@
+
+import numpy as np
+
+from collections import namedtuple
+
+try:
+    from numba import njit
+except ImportError:
+    njit = lambda a : a
+
+from jesse.helpers import get_candle_source, slice_candles
+
+PMA = namedtuple('PMA', ['predict', 'trigger'])
+
+def pma(candles: np.ndarray, source_type: str = "hl2", sequential: bool = False) -> PMA:
+    """
+    Ehlers Predictive Moving Average
+
+    :param candles: np.ndarray
+    :param source_type: str - default: "close"
+    :param sequential: bool - default: False
+
+    :return: float | np.ndarray
+    """
+
+    # Accept normal array too.
+    if len(candles.shape) == 1:
+        source = candles
+    else:
+        candles = slice_candles(candles, sequential)
+        source = get_candle_source(candles, source_type=source_type)
+
+    predict, trigger = pma_fast(source)
+
+    if sequential:
+        return PMA(predict, trigger)
+    else:
+        return PMA(predict[-1], trigger[-1])
+
+
+
+def pma_fast(source):
+    predict = np.full_like(source, np.nan)
+    trigger = np.full_like(source, np.nan)
+    wma1 = np.zeros_like(source)
+    for j in range(6, source.shape[0]):
+        wma1[j] = ((7 * source[j]) + (6 * source[j -1]) + (5 * source[j -2]) + (4 * source[j -3]) + (3 * source[j -4]) + (2 * source[j -5]) + source[j -6]) / 28
+        wma2 = ((7 * wma1[j]) + (6 * wma1[j-1]) + (5 * wma1[j -2]) + (4 * wma1[j -3]) + (3 * wma1[j -4]) + (2 * wma1[j -5]) + wma1[j -6]) / 28
+        predict[j] = (2 * wma1[j]) - wma2
+        trigger[j] = ((4 * predict[j]) + (3 * predict[j-1]) + (2 * predict[j -2]) + predict[j -3]) / 10
+    return predict, trigger

tests/test_indicators.py

@@ -1464,6 +1464,20 @@ def test_pivot4():
     assert len(seq.s3) == len(candles)
     assert len(seq.s4) == len(candles)
 
+def test_pma():
+    candles = np.array(test_candles_19)
+
+    single = ta.pma(candles)
+    assert type(single).__name__ == 'PMA'
+    assert round(single.predict, 2) == 171.05
+    assert round(single.trigger, 2) == 190.91
+
+    seq = ta.pma(candles, sequential=True)
+    assert seq.predict[-1] == single.predict
+    assert seq.trigger[-1] == single.trigger
+    assert len(seq.predict) == len(candles)
+    assert len(seq.trigger) == len(candles)
+
 
 def test_ppo():
     # use the same candles as mama_candles