diff --git a/stockstats.py b/stockstats.py index 37767b6..2a2e9f8 100644 --- a/stockstats.py +++ b/stockstats.py @@ -603,6 +603,30 @@ def _get_supertrend(self, window=None): self['supertrend_lb'] = lb self['supertrend'] = st + def _get_aroon(self, window=None): + if window is None: + window = 25 + column_name = 'aroon' + else: + window = self.get_int_positive(window) + column_name = 'aroon_{}'.format(window) + + def _window_pct(s: pd.Series) -> pd.Series: + return (window - (window - (s + 1))) / window * 100 + + high_since = self['high'].rolling( + min_periods=1, + window=window, + center=False).apply(np.argmax) + low_since = self['low'].rolling( + min_periods=1, + window=window, + center=False).apply(np.argmin) + + aroon_up = _window_pct(high_since) + aroon_down = _window_pct(low_since) + self[column_name] = aroon_up - aroon_down + def _atr(self, window): tr = self._tr() return self._smma(tr, window) @@ -1278,6 +1302,7 @@ def handler(self): ('supertrend', 'supertrend_lb', 'supertrend_ub'): self._get_supertrend, + ('aroon',): self._get_aroon, } def __init_not_exist_column(self, key): diff --git a/test.py b/test.py index 33b0d18..c2b6524 100644 --- a/test.py +++ b/test.py @@ -726,3 +726,15 @@ def test_drop_tail(self): assert_that(ret.iloc[-1].name, equal_to(20040831)) assert_that(stock, has_length(20)) assert_that(stock.iloc[-1].name, equal_to(20040913)) + + def test_aroon(self): + stock = self._supor[:50] + _ = stock['aroon'] + assert_that(stock.loc[20040924, 'aroon'], equal_to(28)) + + _ = stock['aroon_25'] + assert_that(stock.loc[20040924, 'aroon_25'], equal_to(28)) + + _ = stock['aroon_5'] + assert_that(stock.loc[20040924, 'aroon_5'], equal_to(40)) + assert_that(stock.loc[20041020, 'aroon_5'], equal_to(-80))