Merge pull request #1199 from quantopian/boybands-factor

BollingerBands factor

tests/pipeline/test_technical.py

@@ -0,0 +1,140 @@
+import numpy as np
+import pandas as pd
+import talib
+
+from zipline.lib.adjusted_array import AdjustedArray
+from zipline.pipeline import TermGraph
+from zipline.pipeline.data import USEquityPricing
+from zipline.pipeline.engine import SimplePipelineEngine
+from zipline.pipeline.term import AssetExists
+from zipline.pipeline.factors import BollingerBands
+from zipline.testing import ExplodingObject, parameter_space
+from zipline.testing.fixtures import WithAssetFinder, ZiplineTestCase
+from zipline.testing.predicates import assert_equal
+
+
+class WithTechnicalFactor(WithAssetFinder):
+    """ZiplineTestCase fixture for testing technical factors.
+    """
+    ASSET_FINDER_EQUITY_SIDS = tuple(range(5))
+    START_DATE = pd.Timestamp('2014-01-01', tz='utc')
+
+    @classmethod
+    def init_class_fixtures(cls):
+        super(WithTechnicalFactor, cls).init_class_fixtures()
+        cls.ndays = ndays = 24
+        cls.nassets = nassets = len(cls.ASSET_FINDER_EQUITY_SIDS)
+        cls.dates = dates = pd.date_range(cls.START_DATE, periods=ndays)
+        cls.assets = pd.Index(cls.asset_finder.sids)
+        cls.engine = SimplePipelineEngine(
+            lambda column: ExplodingObject(),
+            dates,
+            cls.asset_finder,
+        )
+        cls.asset_exists = exists = np.full((ndays, nassets), True, dtype=bool)
+        cls.asset_exists_masked = masked = exists.copy()
+        masked[:, -1] = False
+
+    def run_graph(self, graph, initial_workspace, mask_sid):
+        initial_workspace.setdefault(
+            AssetExists(),
+            self.asset_exists_masked if mask_sid else self.asset_exists,
+        )
+        return self.engine.compute_chunk(
+            graph,
+            self.dates,
+            self.assets,
+            initial_workspace,
+        )
+
+
+class BollingerBandsTestCase(WithTechnicalFactor, ZiplineTestCase):
+    @classmethod
+    def init_class_fixtures(cls):
+        super(BollingerBandsTestCase, cls).init_class_fixtures()
+        cls._closes = closes = (
+            np.arange(cls.ndays, dtype=float)[:, np.newaxis] +
+            np.arange(cls.nassets, dtype=float) * 100
+        )
+        cls._closes_masked = masked = closes.copy()
+        masked[:, -1] = np.nan
+
+    def closes(self, masked):
+        return self._closes_masked if masked else self._closes
+
+    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.
+        """
+        lower_cols = []
+        middle_cols = []
+        upper_cols = []
+        for n in range(self.nassets):
+            close_col = closes[:, n]
+            if np.isnan(close_col).all():
+                # ta-lib doesn't deal well with all nans.
+                upper, middle, lower = [np.full(self.ndays, np.nan)] * 3
+            else:
+                upper, middle, lower = talib.BBANDS(
+                    close_col,
+                    window_length,
+                    k,
+                    k,
+                )
+
+            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({
+                'f': BollingerBands(
+                    window_length=window_length,
+                    k=k,
+                ),
+            }),
+            initial_workspace={
+                USEquityPricing.close: AdjustedArray(
+                    closes,
+                    np.full_like(closes, True, dtype=bool),
+                    {},
+                    np.nan,
+                ),
+            },
+            mask_sid=mask_sid,
+        )['f']
+
+        expected_upper, expected_middle, expected_lower = self.expected(
+            window_length,
+            k,
+            closes,
+        )
+
+        assert_equal(result.upper, expected_upper)
+        assert_equal(result.middle, expected_middle)
+        assert_equal(result.lower, expected_lower)
+
+    def test_bollinger_bands_output_ordering(self):
+        bbands = BollingerBands(window_length=5, k=2)
+        lower, middle, upper = bbands
+        self.assertIs(lower, bbands.lower)
+        self.assertIs(middle, bbands.middle)
+        self.assertIs(upper, bbands.upper)

tests/pipeline/test_term.py

@@ -27,6 +27,7 @@ from zipline.pipeline.data.testing import TestingDataSet
 from zipline.pipeline.term import AssetExists, NotSpecified
 from zipline.pipeline.expression import NUMEXPR_MATH_FUNCS
 from zipline.testing import parameter_space
+from zipline.testing.predicates import assert_equal, assert_raises
 from zipline.utils.numpy_utils import (
     bool_dtype,
     categorical_dtype,
@@ -358,20 +359,20 @@ class ObjectIdentityTestCase(TestCase):
             method = getattr(f, funcname)
             self.assertIs(method(), method())
 
+    class SomeFactorParameterized(SomeFactor):
+        params = ('a', 'b')
+
     def test_parameterized_term(self):
 
-        class SomeFactorParameterized(SomeFactor):
-            params = ('a', 'b')
-
-        f = SomeFactorParameterized(a=1, b=2)
+        f = self.SomeFactorParameterized(a=1, b=2)
         self.assertEqual(f.params, {'a': 1, 'b': 2})
 
-        g = SomeFactorParameterized(a=1, b=3)
-        h = SomeFactorParameterized(a=2, b=2)
+        g = self.SomeFactorParameterized(a=1, b=3)
+        h = self.SomeFactorParameterized(a=2, b=2)
         self.assertDifferentObjects(f, g, h)
 
-        f2 = SomeFactorParameterized(a=1, b=2)
-        f3 = SomeFactorParameterized(b=2, a=1)
+        f2 = self.SomeFactorParameterized(a=1, b=2)
+        f3 = self.SomeFactorParameterized(b=2, a=1)
         self.assertSameObject(f, f2, f3)
 
         self.assertEqual(f.params['a'], 1)
@@ -379,6 +380,31 @@ class ObjectIdentityTestCase(TestCase):
         self.assertEqual(f.window_length, SomeFactor.window_length)
         self.assertEqual(f.inputs, tuple(SomeFactor.inputs))
 
+    def test_parameterized_term_non_hashable_arg(self):
+        with assert_raises(TypeError) as e:
+            self.SomeFactorParameterized(a=[], b=1)
+        assert_equal(
+            str(e.exception),
+            "SomeFactorParameterized expected a hashable value for parameter"
+            " 'a', but got [] instead.",
+        )
+
+        with assert_raises(TypeError) as e:
+            self.SomeFactorParameterized(a=1, b=[])
+        assert_equal(
+            str(e.exception),
+            "SomeFactorParameterized expected a hashable value for parameter"
+            " 'b', but got [] instead.",
+        )
+
+        with assert_raises(TypeError) as e:
+            self.SomeFactorParameterized(a=[], b=[])
+        assert_equal(
+            str(e.exception),
+            "SomeFactorParameterized expected a hashable value for parameter"
+            " 'a', but got [] instead.",
+        )
+
     def test_bad_input(self):
 
         class SomeFactor(Factor):

tests/utils/test_metautils.py

@@ -0,0 +1,91 @@
+from zipline.testing.fixtures import ZiplineTestCase
+from zipline.testing.predicates import (
+    assert_equal,
+    assert_is,
+    assert_is_instance,
+    assert_is_subclass,
+    assert_true,
+)
+from zipline.utils.metautils import compose_types, with_metaclasses
+
+
+class C(object):
+    @staticmethod
+    def f():
+        return 'C.f'
+
+    def delegate(self):
+        return 'C.delegate', super(C, self).delegate()
+
+
+class D(object):
+    @staticmethod
+    def f():
+        return 'D.f'
+
+    @staticmethod
+    def g():
+        return 'D.g'
+
+    def delegate(self):
+        return 'D.delegate'
+
+
+class ComposeTypesTestCase(ZiplineTestCase):
+
+    def test_identity(self):
+        assert_is(
+            compose_types(C),
+            C,
+            msg='compose_types of a single class should be identity',
+        )
+
+    def test_compose(self):
+        composed = compose_types(C, D)
+
+        assert_is_subclass(composed, C)
+        assert_is_subclass(composed, D)
+
+    def test_compose_mro(self):
+        composed = compose_types(C, D)
+
+        assert_equal(composed.f(), C.f())
+        assert_equal(composed.g(), D.g())
+
+        assert_equal(composed().delegate(), ('C.delegate', 'D.delegate'))
+
+
+class M(type):
+    def __new__(mcls, name, bases, dict_):
+        dict_['M'] = True
+        return super(M, mcls).__new__(mcls, name, bases, dict_)
+
+
+class N(type):
+    def __new__(mcls, name, bases, dict_):
+        dict_['N'] = True
+        return super(N, mcls).__new__(mcls, name, bases, dict_)
+
+
+class WithMetaclassesTestCase(ZiplineTestCase):
+    def test_with_metaclasses_no_subclasses(self):
+        class E(with_metaclasses((M, N))):
+            pass
+
+        assert_true(E.M)
+        assert_true(E.N)
+
+        assert_is_instance(E, M)
+        assert_is_instance(E, N)
+
+    def test_with_metaclasses_with_subclasses(self):
+        class E(with_metaclasses((M, N), C, D)):
+            pass
+
+        assert_true(E.M)
+        assert_true(E.N)
+
+        assert_is_instance(E, M)
+        assert_is_instance(E, N)
+        assert_is_subclass(E, C)
+        assert_is_subclass(E, D)