diff --git a/README.md b/README.md index 2af9ed1..7c7cc6d 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![codecov](https://codecov.io/gh/jealous/stockstats/branch/master/graph/badge.svg?token=IFMD1pVJ7T)](https://codecov.io/gh/jealous/stockstats) [![pypi](https://img.shields.io/pypi/v/stockstats.svg)](https://pypi.python.org/pypi/stockstats) -VERSION: 0.5.0 +VERSION: 0.5.1 ## Introduction diff --git a/stockstats.py b/stockstats.py index 4902118..eda718d 100644 --- a/stockstats.py +++ b/stockstats.py @@ -759,7 +759,7 @@ def _get_kdj_default(self): self['kdjd'] = self['kdjd_{}'.format(self.KDJ_WINDOW)] self['kdjj'] = self['kdjj_{}'.format(self.KDJ_WINDOW)] - def _get_cr(self, window=26): + def _get_cr(self, windows=None): """ Energy Index (Intermediate Willingness Index) https://support.futunn.com/en/topic167/?lang=en-us @@ -767,9 +767,14 @@ def _get_cr(self, window=26): yesterday's middle price to reflect the market's willingness to buy and sell. - :param window: window of the moving sum + :param windows: window of the moving sum :return: None """ + if windows is None: + window = 26 + else: + window = self.get_int_positive(windows) + middle = self._tp() last_middle = self._shift(middle, -1) ym = self._shift(middle, -1) @@ -779,10 +784,22 @@ def _get_cr(self, window=26): p2_m = pd.concat((last_middle, low), axis=1).min(axis=1) p1 = self._mov_sum(high - p1_m, window) p2 = self._mov_sum(ym - p2_m, window) - self['cr'] = cr = p1 / p2 * 100 - self['cr-ma1'] = self._shifted_cr_sma(cr, self.CR_MA1) - self['cr-ma2'] = self._shifted_cr_sma(cr, self.CR_MA2) - self['cr-ma3'] = self._shifted_cr_sma(cr, self.CR_MA3) + + if windows is None: + cr = 'cr' + cr_ma1 = 'cr-ma1' + cr_ma2 = 'cr-ma2' + cr_ma3 = 'cr-ma3' + else: + cr = 'cr_{}'.format(window) + cr_ma1 = 'cr_{}-ma1'.format(window) + cr_ma2 = 'cr_{}-ma2'.format(window) + cr_ma3 = 'cr_{}-ma3'.format(window) + + self[cr] = cr = p1 / p2 * 100 + self[cr_ma1] = self._shifted_cr_sma(cr, self.CR_MA1) + self[cr_ma2] = self._shifted_cr_sma(cr, self.CR_MA2) + self[cr_ma3] = self._shifted_cr_sma(cr, self.CR_MA3) def _shifted_cr_sma(self, cr, window): cr_sma = self._sma(cr, window) @@ -1108,7 +1125,10 @@ def _get_kama(self, column, windows, fasts=None, slows=None): # start with simple moving average kama = self._sma(col, window) - last_kama = kama.iloc[window - 1] + if len(kama) >= window: + last_kama = kama.iloc[window - 1] + else: + last_kama = 0.0 for i in range(window, len(kama)): cur = smoothing.iloc[i] * (col.iloc[i] - last_kama) + last_kama kama.iloc[i] = cur diff --git a/test.py b/test.py index 4886bf8..e7d7b21 100644 --- a/test.py +++ b/test.py @@ -161,6 +161,12 @@ def test_cr(self): assert_that(stock['cr-ma2'].loc[20110331], near_to(117.1488)) assert_that(stock['cr-ma3'].loc[20110331], near_to(111.5195)) + stock.get('cr_26') + assert_that(stock['cr_26'].loc[20110331], near_to(178.1714)) + assert_that(stock['cr_26-ma1'].loc[20110331], near_to(120.0364)) + assert_that(stock['cr_26-ma2'].loc[20110331], near_to(117.1488)) + assert_that(stock['cr_26-ma3'].loc[20110331], near_to(111.5195)) + def test_column_permutation(self): stock = self.get_stock_20day() amount_p = stock['volume_-1_d_-3,-2,-1_p']