diff --git a/pandas_ta/candles/ha.py b/pandas_ta/candles/ha.py index 9bd436c..6f7184a 100644 --- a/pandas_ta/candles/ha.py +++ b/pandas_ta/candles/ha.py @@ -1,8 +1,33 @@ # -*- coding: utf-8 -*- +import numpy as np from pandas import DataFrame, Series -from pandas_ta._typing import DictLike, Int +from pandas_ta._typing import DictLike, Int, List from pandas_ta.utils import v_offset, v_series +try: + from numba import njit +except ImportError: + def njit(_): return _ + + +@njit +def heiken_ashi_numpy(c_open, c_high, c_low, c_close): + """fast version using numpy and optionally numba""" + ha_close = (c_open + c_high + c_low + c_close) / 4 + ha_open = np.empty_like(ha_close) + ha_open[0] = (c_open[0] + c_close[0]) / 2 + for i in range(1, len(c_close)): + ha_open[i] = (ha_open[i - 1] + ha_close[i - 1]) / 2 + ha_high = np.maximum(np.maximum(ha_open, ha_close), c_high) + ha_low = np.minimum(np.minimum(ha_open, ha_close), c_low) + return ha_open, ha_high, ha_low, ha_close + +def heiken_ashi_series(open: Series, high: Series, low: Series, close: Series) -> List[Series]: + """takes in series, outputs series.""" + inputs = [open, high, low, close] + outs = heiken_ashi_numpy(open.values, high.values, low.values, close.values) + outs = [Series(outs[i], index=inputs[i].index) for i in range(len(inputs))] + return outs def ha( open_: Series, high: Series, low: Series, close: Series, @@ -47,21 +72,15 @@ def ha( return # Calculate + ha_open, ha_high, ha_low, ha_close = heiken_ashi_series(open_, high, low, close) m = close.size df = DataFrame({ - "HA_open": 0.5 * (open_.iloc[0] + close.iloc[0]), - "HA_high": high, - "HA_low": low, - "HA_close": 0.25 * (open_ + high + low + close), + "HA_open": ha_open, + "HA_high": ha_high, + "HA_low": ha_low, + "HA_close": ha_close, }) - for i in range(1, m): - df["HA_open"].iloc[i] = 0.5 * (df["HA_open"].iloc[i - 1] \ - + df["HA_close"].iloc[i - 1]) - - df["HA_high"] = df[["HA_open", "HA_high", "HA_close"]].max(axis=1) - df["HA_low"] = df[["HA_open", "HA_low", "HA_close"]].min(axis=1) - # Offset if offset != 0: df = df.shift(offset)