def setUp(self): self.constants = { # Every day, assume every stock starts at 2, goes down to 1, # goes up to 4, and finishes at 3. USEquityPricing.low: 1, USEquityPricing.open: 2, USEquityPricing.close: 3, USEquityPricing.high: 4, } self.assets = [1, 2, 3] self.dates = date_range('2014-01', '2014-03', freq='D', tz='UTC') self.loader = ConstantLoader( constants=self.constants, dates=self.dates, assets=self.assets, ) self.asset_info = make_simple_asset_info( self.assets, start_date=self.dates[0], end_date=self.dates[-1], ) environment = TradingEnvironment() environment.write_data(equities_df=self.asset_info) self.asset_finder = environment.asset_finder
def test_rolling_and_nonrolling(self): open_ = USEquityPricing.open close = USEquityPricing.close volume = USEquityPricing.volume # Test for thirty days up to the last day that we think all # the assets existed. dates_to_test = self.dates[-30:] constants = {open_: 1, close: 2, volume: 3} loader = ConstantLoader( constants=constants, dates=self.dates, assets=self.asset_ids, ) engine = SimplePipelineEngine( lambda column: loader, self.dates, self.asset_finder, ) sumdiff = RollingSumDifference() result = engine.run_pipeline( Pipeline(columns={ 'sumdiff': sumdiff, 'open': open_.latest, 'close': close.latest, 'volume': volume.latest, }, ), dates_to_test[0], dates_to_test[-1]) self.assertIsNotNone(result) self.assertEqual({'sumdiff', 'open', 'close', 'volume'}, set(result.columns)) result_index = self.asset_ids * len(dates_to_test) result_shape = (len(result_index), ) check_arrays( result['sumdiff'], Series( index=result_index, data=full(result_shape, -3, dtype=float), ), ) for name, const in [('open', 1), ('close', 2), ('volume', 3)]: check_arrays( result[name], Series( index=result_index, data=full(result_shape, const, dtype=float), ), )
def test_engine_with_multicolumn_loader(self): open_ = USEquityPricing.open close = USEquityPricing.close volume = USEquityPricing.volume # Test for thirty days up to the second to last day that we think all # the assets existed. If we test the last day of our calendar, no # assets will be in our output, because their end dates are all dates_to_test = self.dates[-32:-2] constants = {open_: 1, close: 2, volume: 3} loader = ConstantLoader( constants=constants, dates=self.dates, assets=self.assets, ) engine = SimplePipelineEngine(loader, self.dates, self.asset_finder) sumdiff = RollingSumDifference() result = engine.run_pipeline( Pipeline( columns={ 'sumdiff': sumdiff, 'open': open_.latest, 'close': close.latest, 'volume': volume.latest, }, ), dates_to_test[0], dates_to_test[-1] ) self.assertIsNotNone(result) self.assertEqual( {'sumdiff', 'open', 'close', 'volume'}, set(result.columns) ) result_index = self.assets * len(dates_to_test) result_shape = (len(result_index),) check_arrays( result['sumdiff'], Series(index=result_index, data=full(result_shape, -3)), ) for name, const in [('open', 1), ('close', 2), ('volume', 3)]: check_arrays( result[name], Series(index=result_index, data=full(result_shape, const)), )