diff --git a/tests/modelling/test_engine.py b/tests/modelling/test_engine.py index cf67354e..c647b50f 100644 --- a/tests/modelling/test_engine.py +++ b/tests/modelling/test_engine.py @@ -16,6 +16,7 @@ from pandas import ( Int64Index, rolling_mean, Timestamp, + Series, ) from pandas.util.testing import assert_frame_equal from testfixtures import TempDirectory @@ -48,6 +49,7 @@ from zipline.utils.test_utils import ( make_rotating_asset_info, make_simple_asset_info, product_upper_triangle, + check_arrays, ) @@ -430,3 +432,42 @@ class SyntheticBcolzTestCase(TestCase): columns=assets, ), ) + + +class MultiColumnLoaderTestCase(TestCase): + def setUp(self): + self.assets = [1, 2, 3] + self.dates = date_range('2014-01-01', '2014-02-01', freq='D', tz='UTC') + + asset_info = make_simple_asset_info( + self.assets, + start_date=self.dates[0], + end_date=self.dates[-1], + ) + self.asset_finder = AssetFinder(asset_info) + + def test_engine_with_multicolumn_loader(self): + open_, close = USEquityPricing.open, USEquityPricing.close + + loader = MultiColumnLoader({ + open_: ConstantLoader(dates=self.dates, + assets=self.assets, + constants={open_: 1}), + close: ConstantLoader(dates=self.dates, + assets=self.assets, + constants={close: 2}) + }) + + engine = SimpleFFCEngine(loader, self.dates, self.asset_finder) + + factor = RollingSumDifference() + + result = engine.factor_matrix({'f': factor}, + self.dates[2], + self.dates[-1]) + self.assertIsNotNone(result) + self.assertEqual({'f'}, set(result.columns)) + # (close - open) * window = (1 - 2) * 3 = -3 + # skipped 2 from the start, so that the window is full + check_arrays(result['f'], + Series([-3] * len(self.assets) * (len(self.dates) - 2))) diff --git a/tests/modelling/test_frameload.py b/tests/modelling/test_frameload.py index c99101f3..47512edf 100644 --- a/tests/modelling/test_frameload.py +++ b/tests/modelling/test_frameload.py @@ -79,7 +79,7 @@ class DataFrameFFCLoaderTestCase(TestCase): dates_slice = slice(None, 10, None) sids_slice = slice(1, 3, None) - adj_array = loader.load_adjusted_array( + [adj_array] = loader.load_adjusted_array( [USEquityPricing.close], self.mask.iloc[dates_slice, sids_slice] ) diff --git a/zipline/data/ffc/frame.py b/zipline/data/ffc/frame.py index 7f78a5e6..b8a19ed5 100644 --- a/zipline/data/ffc/frame.py +++ b/zipline/data/ffc/frame.py @@ -176,10 +176,10 @@ class DataFrameFFCLoader(FFCLoader): good_dates = (date_indexer != -1) good_assets = (assets_indexer != -1) - return adjusted_array( + return [adjusted_array( # Pull out requested columns/rows from our baseline data. data=self.baseline[ix_(date_indexer, assets_indexer)], # Mask out requested columns/rows that didnt match. mask=(good_assets & good_dates[:, None]) & mask_values, adjustments=self.format_adjustments(dates, assets), - ) + )] diff --git a/zipline/data/ffc/synthetic.py b/zipline/data/ffc/synthetic.py index 19e8cb33..ef3ba13f 100644 --- a/zipline/data/ffc/synthetic.py +++ b/zipline/data/ffc/synthetic.py @@ -58,7 +58,7 @@ class MultiColumnLoader(FFCLoader): loader = self._loaders[column] except KeyError: raise ValueError("Couldn't find loader for %s" % column) - out.append(loader.load_adjusted_array([column], mask)) + out.extend(loader.load_adjusted_array([column], mask)) return out