From c128b69a91d95bd7cf0cecf52448c918616f3e46 Mon Sep 17 00:00:00 2001 From: Joe Jevnik Date: Thu, 12 May 2016 17:00:30 -0400 Subject: [PATCH] STY: get Scott to stop yelling at me --- tests/pipeline/test_technical.py | 139 +++++++------------------------ 1 file changed, 29 insertions(+), 110 deletions(-) diff --git a/tests/pipeline/test_technical.py b/tests/pipeline/test_technical.py index 5ac08fd6..c2b1dfc6 100644 --- a/tests/pipeline/test_technical.py +++ b/tests/pipeline/test_technical.py @@ -1,13 +1,8 @@ -from abc import abstractmethod - import numpy as np -from operator import itemgetter import pandas as pd import talib -from toolz import compose, excepts from zipline.lib.adjusted_array import AdjustedArray -from zipline.lib.adjustment import Float64Add from zipline.pipeline import TermGraph from zipline.pipeline.data import USEquityPricing from zipline.pipeline.engine import SimplePipelineEngine @@ -52,22 +47,8 @@ class WithTechnicalFactor(WithAssetFinder): initial_workspace, ) - @abstractmethod - def test_without_adjustments(self): - raise NotImplementedError('test_without_adjustments') - - @abstractmethod - def test_with_adjustments(self): - raise NotImplementedError('test_with_adjustments') - class BollingerBandsTestCase(WithTechnicalFactor, ZiplineTestCase): - cases = parameter_space( - window_length={5, 10, 20}, - k={1.5, 2, 2.5}, - mask_sid={True, False}, - ) - @classmethod def init_class_fixtures(cls): super(BollingerBandsTestCase, cls).init_class_fixtures() @@ -81,45 +62,47 @@ class BollingerBandsTestCase(WithTechnicalFactor, ZiplineTestCase): def closes(self, masked): return self._closes_masked if masked else self._closes - def _allnans(self, e): - """Handler for toolz.excepts that returns three all nan arrays. - """ - nans = np.full(self.ndays, np.nan) - return nans, nans, nans - def expected(self, window_length, k, closes): """Compute the expected data (without adjustments) for the given window, k, and closes array. This uses talib.BBANDS to generate the expected data. """ - return map( - # Stack all of our uppers, middles, lowers into three 2d arrays - # whose columns are the sids. After that, slice off only the - # window care about. - compose( - itemgetter(np.s_[window_length - 1:]), - np.column_stack, - ), - # Take our sequence of [(uppers, middles, lowers)] per sid and - # turn it into three sequences of: - # uppers for all sids, middles for all sids, lowers for all sids. - zip(*( - # talib breaks when the input array is all nan and raises - # an instance of Exception. We catch that here and return - # three all nan arrays instead. - excepts(Exception, talib.BBANDS, self._allnans)( + lower_cols = [] + middle_cols = [] + upper_cols = [] + for n in range(self.nassets): + try: + upper, middle, lower = talib.BBANDS( closes[:, n], window_length, k, k, ) - for n in range(self.nassets) - )), - ) + except Exception: + # If the input array is all nan then talib raises an instance + # of Exception. + upper, middle, lower = [np.full(self.ndays, np.nan)] * 3 - @cases - def test_without_adjustments(self, window_length, k, mask_sid): + upper_cols.append(upper) + middle_cols.append(middle) + lower_cols.append(lower) + + # Stack all of our uppers, middles, lowers into three 2d arrays + # whose columns are the sids. After that, slice off only the + # rows we care about. + where = np.s_[window_length - 1:] + uppers = np.column_stack(upper_cols)[where] + middles = np.column_stack(middle_cols)[where] + lowers = np.column_stack(lower_cols)[where] + return uppers, middles, lowers + + @parameter_space( + window_length={5, 10, 20}, + k={1.5, 2, 2.5}, + mask_sid={True, False}, + ) + def test_bollinger_bands(self, window_length, k, mask_sid): closes = self.closes(mask_sid) result = self.run_graph( TermGraph({ @@ -157,67 +140,3 @@ class BollingerBandsTestCase(WithTechnicalFactor, ZiplineTestCase): result.lower, expected_lower, ) - - @cases - def test_with_adjustments(self, window_length, k, mask_sid): - closes = self.closes(mask_sid) - adjustment_offset = 5 - adjustment_idx = self.ndays - adjustment_offset - result = self.run_graph( - TermGraph({ - 'f': BollingerBands( - window_length=window_length, - k=k, - ), - }), - initial_workspace={ - USEquityPricing.close: AdjustedArray( - closes, - np.full_like(closes, True, dtype=bool), - { - adjustment_idx: [ - Float64Add( - first_row=0, - last_row=adjustment_idx, - first_col=0, - last_col=self.nassets - 1, - value=1000, - ), - ], - }, - np.nan, - ), - }, - mask_sid=mask_sid, - )['f'] - - # Get the uppers, middles, and lowers without the adjustment applied. - bases = self.expected(window_length, k, closes) - adjusted_closes = closes.copy() - adjusted_closes[:adjustment_idx + 1] += 1000 - # Get the uppers, middles, and lowers with the adjustment applied. - adjusted = self.expected(window_length, k, adjusted_closes) - - # Create the actual expected data by using the unadjusted results up - # to the adjument offset, then use the adjusted data for all indices - # past the adjustment offset. - expected_upper, expected_middle, expected_lower = ( - np.vstack(( - base[:-adjustment_offset], - adjusted[-adjustment_offset:], - )) - for base, adjusted in zip(bases, adjusted) - ) - - assert_equal( - result.upper, - expected_upper, - ) - assert_equal( - result.middle, - expected_middle, - ) - assert_equal( - result.lower, - expected_lower, - )