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)),
)