def setUpClass(cls): cls.first_asset_start = Timestamp("2015-04-01", tz="UTC") cls.env = TradingEnvironment() cls.trading_day = day = cls.env.trading_day cls.calendar = date_range("2015", "2015-08", tz="UTC", freq=day) cls.asset_info = make_rotating_equity_info( num_assets=6, first_start=cls.first_asset_start, frequency=day, periods_between_starts=4, asset_lifetime=8 ) cls.last_asset_end = cls.asset_info["end_date"].max() cls.all_assets = cls.asset_info.index cls.env.write_data(equities_df=cls.asset_info) cls.finder = cls.env.asset_finder cls.temp_dir = TempDirectory() cls.temp_dir.create() try: cls.writer = SyntheticDailyBarWriter( asset_info=cls.asset_info[["start_date", "end_date"]], calendar=cls.calendar ) table = cls.writer.write(cls.temp_dir.getpath("testdata.bcolz"), cls.calendar, cls.all_assets) cls.pipeline_loader = USEquityPricingLoader(BcolzDailyBarReader(table), NullAdjustmentReader()) except: cls.temp_dir.cleanup() raise
def test_compute_lifetimes(self): num_assets = 4 trading_day = self.env.trading_day first_start = pd.Timestamp('2015-04-01', tz='UTC') frame = make_rotating_equity_info( num_assets=num_assets, first_start=first_start, frequency=self.env.trading_day, periods_between_starts=3, asset_lifetime=5 ) self.env.write_data(equities_df=frame) finder = self.env.asset_finder all_dates = pd.date_range( start=first_start, end=frame.end_date.max(), freq=trading_day, ) for dates in all_subindices(all_dates): expected_with_start_raw = full( shape=(len(dates), num_assets), fill_value=False, dtype=bool, ) expected_no_start_raw = full( shape=(len(dates), num_assets), fill_value=False, dtype=bool, ) for i, date in enumerate(dates): it = frame[['start_date', 'end_date']].itertuples() for j, start, end in it: # This way of doing the checks is redundant, but very # clear. if start <= date <= end: expected_with_start_raw[i, j] = True if start < date: expected_no_start_raw[i, j] = True expected_with_start = pd.DataFrame( data=expected_with_start_raw, index=dates, columns=frame.index.values, ) result = finder.lifetimes(dates, include_start_date=True) assert_frame_equal(result, expected_with_start) expected_no_start = pd.DataFrame( data=expected_no_start_raw, index=dates, columns=frame.index.values, ) result = finder.lifetimes(dates, include_start_date=False) assert_frame_equal(result, expected_no_start)
def setUpClass(cls): cls.first_asset_start = Timestamp('2015-04-01', tz='UTC') cls.env = TradingEnvironment() cls.trading_day = day = cls.env.trading_day cls.calendar = date_range('2015', '2015-08', tz='UTC', freq=day) cls.asset_info = make_rotating_equity_info( num_assets=6, first_start=cls.first_asset_start, frequency=day, periods_between_starts=4, asset_lifetime=8, ) cls.last_asset_end = cls.asset_info['end_date'].max() cls.all_asset_ids = cls.asset_info.index cls.env.write_data(equities_df=cls.asset_info) cls.finder = cls.env.asset_finder cls.temp_dir = TempDirectory() cls.temp_dir.create() try: cls.writer = SyntheticDailyBarWriter( asset_info=cls.asset_info[['start_date', 'end_date']], calendar=cls.calendar, ) table = cls.writer.write( cls.temp_dir.getpath('testdata.bcolz'), cls.calendar, cls.all_asset_ids, ) cls.pipeline_loader = USEquityPricingLoader( BcolzDailyBarReader(table), NullAdjustmentReader(), ) except: cls.temp_dir.cleanup() raise