TST: update adjustment tests - add gaps between sids

TST: add a seed for permuting
This commit is contained in:
Maya Tydykov
2016-10-21 16:53:56 -04:00
parent 4ea60c2843
commit 086ea6be6b
+59 -45
View File
@@ -825,7 +825,7 @@ class WithEstimateWindows(WithEstimates):
SID_FIELD_NAME: 0, SID_FIELD_NAME: 0,
}) })
sid_1_timeline = pd.DataFrame({ sid_10_timeline = pd.DataFrame({
TS_FIELD_NAME: [pd.Timestamp('2015-01-09'), TS_FIELD_NAME: [pd.Timestamp('2015-01-09'),
pd.Timestamp('2015-01-12'), pd.Timestamp('2015-01-12'),
pd.Timestamp('2015-01-09'), pd.Timestamp('2015-01-09'),
@@ -836,12 +836,12 @@ class WithEstimateWindows(WithEstimates):
'estimate': [110., 111.] + [310., 311.], 'estimate': [110., 111.] + [310., 311.],
FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [3] * 2, FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [3] * 2,
FISCAL_YEAR_FIELD_NAME: 2015, FISCAL_YEAR_FIELD_NAME: 2015,
SID_FIELD_NAME: 1 SID_FIELD_NAME: 10
}) })
# Extra sid to make sure we have correct overwrites when sid quarter # Extra sid to make sure we have correct overwrites when sid quarter
# boundaries collide. # boundaries collide.
sid_3_timeline = pd.DataFrame({ sid_20_timeline = pd.DataFrame({
TS_FIELD_NAME: [pd.Timestamp('2015-01-05'), TS_FIELD_NAME: [pd.Timestamp('2015-01-05'),
pd.Timestamp('2015-01-07'), pd.Timestamp('2015-01-07'),
pd.Timestamp('2015-01-05'), pd.Timestamp('2015-01-05'),
@@ -854,9 +854,21 @@ class WithEstimateWindows(WithEstimates):
'estimate': [120., 121.] + [220., 221.], 'estimate': [120., 121.] + [220., 221.],
FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [2] * 2, FISCAL_QUARTER_FIELD_NAME: [1] * 2 + [2] * 2,
FISCAL_YEAR_FIELD_NAME: 2015, FISCAL_YEAR_FIELD_NAME: 2015,
SID_FIELD_NAME: 2 SID_FIELD_NAME: 20
}) })
return pd.concat([sid_0_timeline, sid_1_timeline, sid_3_timeline]) concatted = pd.concat([sid_0_timeline,
sid_10_timeline,
sid_20_timeline]).reset_index()
np.random.seed(0)
return concatted.reindex(np.random.permutation(concatted.index))
@classmethod
def get_sids(cls):
sids = sorted(cls.events[SID_FIELD_NAME].unique())
# Add extra sids between sids in our data. We want to test that we
# apply adjustments to the correct sids.
return [sid for i in range(len(sids) - 1)
for sid in range(sids[i], sids[i+1])] + [sids[-1]]
@classmethod @classmethod
def make_expected_timelines(cls): def make_expected_timelines(cls):
@@ -888,6 +900,8 @@ class WithEstimateWindows(WithEstimates):
df.index = df.index.rename('knowledge_date') df.index = df.index.rename('knowledge_date')
df['at_date'] = end_date.tz_localize('utc') df['at_date'] = end_date.tz_localize('utc')
df = df.set_index(['at_date', df.index.tz_localize('utc')]).ffill() df = df.set_index(['at_date', df.index.tz_localize('utc')]).ffill()
new_sids = set(cls.get_sids()) - set(df.columns)
df = df.reindex(columns=df.columns.union(new_sids))
return df return df
@parameterized.expand(window_test_cases) @parameterized.expand(window_test_cases)
@@ -943,44 +957,44 @@ class PreviousEstimateWindows(WithEstimateWindows, ZiplineTestCase):
oneq_previous = pd.concat([ oneq_previous = pd.concat([
cls.create_expected_df( cls.create_expected_df(
[(0, np.NaN, cls.window_test_start_date), [(0, np.NaN, cls.window_test_start_date),
(1, np.NaN, cls.window_test_start_date), (10, np.NaN, cls.window_test_start_date),
(2, np.NaN, cls.window_test_start_date)], (20, np.NaN, cls.window_test_start_date)],
pd.Timestamp('2015-01-09') pd.Timestamp('2015-01-09')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 101, pd.Timestamp('2015-01-10')), [(0, 101, pd.Timestamp('2015-01-10')),
(1, 111, pd.Timestamp('2015-01-12')), (10, 111, pd.Timestamp('2015-01-12')),
(2, 121, pd.Timestamp('2015-01-10'))], (20, 121, pd.Timestamp('2015-01-10'))],
pd.Timestamp('2015-01-12') pd.Timestamp('2015-01-12')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 101, pd.Timestamp('2015-01-10')), [(0, 101, pd.Timestamp('2015-01-10')),
(1, 111, pd.Timestamp('2015-01-12')), (10, 111, pd.Timestamp('2015-01-12')),
(2, 121, pd.Timestamp('2015-01-10'))], (20, 121, pd.Timestamp('2015-01-10'))],
pd.Timestamp('2015-01-13') pd.Timestamp('2015-01-13')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 101, pd.Timestamp('2015-01-10')), [(0, 101, pd.Timestamp('2015-01-10')),
(1, 111, pd.Timestamp('2015-01-12')), (10, 111, pd.Timestamp('2015-01-12')),
(2, 121, pd.Timestamp('2015-01-10'))], (20, 121, pd.Timestamp('2015-01-10'))],
pd.Timestamp('2015-01-14') pd.Timestamp('2015-01-14')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 101, pd.Timestamp('2015-01-10')), [(0, 101, pd.Timestamp('2015-01-10')),
(1, 311, pd.Timestamp('2015-01-15')), (10, 311, pd.Timestamp('2015-01-15')),
(2, 121, pd.Timestamp('2015-01-10'))], (20, 121, pd.Timestamp('2015-01-10'))],
pd.Timestamp('2015-01-15') pd.Timestamp('2015-01-15')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 101, pd.Timestamp('2015-01-10')), [(0, 101, pd.Timestamp('2015-01-10')),
(1, 311, pd.Timestamp('2015-01-15')), (10, 311, pd.Timestamp('2015-01-15')),
(2, 121, pd.Timestamp('2015-01-10'))], (20, 121, pd.Timestamp('2015-01-10'))],
pd.Timestamp('2015-01-16') pd.Timestamp('2015-01-16')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 201, pd.Timestamp('2015-01-17')), [(0, 201, pd.Timestamp('2015-01-17')),
(1, 311, pd.Timestamp('2015-01-15')), (10, 311, pd.Timestamp('2015-01-15')),
(2, 221, pd.Timestamp('2015-01-17'))], (20, 221, pd.Timestamp('2015-01-17'))],
pd.Timestamp('2015-01-20') pd.Timestamp('2015-01-20')
), ),
]) ])
@@ -988,14 +1002,14 @@ class PreviousEstimateWindows(WithEstimateWindows, ZiplineTestCase):
twoq_previous = pd.concat( twoq_previous = pd.concat(
[cls.create_expected_df( [cls.create_expected_df(
[(0, np.NaN, cls.window_test_start_date), [(0, np.NaN, cls.window_test_start_date),
(1, np.NaN, cls.window_test_start_date), (10, np.NaN, cls.window_test_start_date),
(2, np.NaN, cls.window_test_start_date)], (20, np.NaN, cls.window_test_start_date)],
end_date end_date
) for end_date in pd.date_range('2015-01-09', '2015-01-19')] + ) for end_date in pd.date_range('2015-01-09', '2015-01-19')] +
[cls.create_expected_df( [cls.create_expected_df(
[(0, 101, pd.Timestamp('2015-01-20')), [(0, 101, pd.Timestamp('2015-01-20')),
(1, np.NaN, cls.window_test_start_date), (10, np.NaN, cls.window_test_start_date),
(2, 121, pd.Timestamp('2015-01-20'))], (20, 121, pd.Timestamp('2015-01-20'))],
pd.Timestamp('2015-01-20') pd.Timestamp('2015-01-20')
)] )]
) )
@@ -1016,49 +1030,49 @@ class NextEstimateWindows(WithEstimateWindows, ZiplineTestCase):
cls.create_expected_df( cls.create_expected_df(
[(0, 100, cls.window_test_start_date), [(0, 100, cls.window_test_start_date),
(0, 101, pd.Timestamp('2015-01-07')), (0, 101, pd.Timestamp('2015-01-07')),
(1, 110, pd.Timestamp('2015-01-09')), (10, 110, pd.Timestamp('2015-01-09')),
(2, 120, cls.window_test_start_date), (20, 120, cls.window_test_start_date),
(2, 121, pd.Timestamp('2015-01-07'))], (20, 121, pd.Timestamp('2015-01-07'))],
pd.Timestamp('2015-01-09') pd.Timestamp('2015-01-09')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 200, cls.window_test_start_date), [(0, 200, cls.window_test_start_date),
(1, 110, pd.Timestamp('2015-01-09')), (10, 110, pd.Timestamp('2015-01-09')),
(1, 111, pd.Timestamp('2015-01-12')), (10, 111, pd.Timestamp('2015-01-12')),
(2, 220, cls.window_test_start_date)], (20, 220, cls.window_test_start_date)],
pd.Timestamp('2015-01-12') pd.Timestamp('2015-01-12')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 200, cls.window_test_start_date), [(0, 200, cls.window_test_start_date),
(1, 310, pd.Timestamp('2015-01-09')), (10, 310, pd.Timestamp('2015-01-09')),
(2, 220, cls.window_test_start_date)], (20, 220, cls.window_test_start_date)],
pd.Timestamp('2015-01-13') pd.Timestamp('2015-01-13')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 200, cls.window_test_start_date), [(0, 200, cls.window_test_start_date),
(1, 310, pd.Timestamp('2015-01-09')), (10, 310, pd.Timestamp('2015-01-09')),
(2, 220, cls.window_test_start_date)], (20, 220, cls.window_test_start_date)],
pd.Timestamp('2015-01-14') pd.Timestamp('2015-01-14')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 200, cls.window_test_start_date), [(0, 200, cls.window_test_start_date),
(1, 310, pd.Timestamp('2015-01-09')), (10, 310, pd.Timestamp('2015-01-09')),
(1, 311, pd.Timestamp('2015-01-15')), (10, 311, pd.Timestamp('2015-01-15')),
(2, 220, cls.window_test_start_date)], (20, 220, cls.window_test_start_date)],
pd.Timestamp('2015-01-15') pd.Timestamp('2015-01-15')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 200, cls.window_test_start_date), [(0, 200, cls.window_test_start_date),
(1, np.NaN, cls.window_test_start_date), (10, np.NaN, cls.window_test_start_date),
(2, 220, cls.window_test_start_date)], (20, 220, cls.window_test_start_date)],
pd.Timestamp('2015-01-16') pd.Timestamp('2015-01-16')
), ),
cls.create_expected_df( cls.create_expected_df(
[(0, 200, cls.window_test_start_date), [(0, 200, cls.window_test_start_date),
(0, 201, pd.Timestamp('2015-01-17')), (0, 201, pd.Timestamp('2015-01-17')),
(1, np.NaN, cls.window_test_start_date), (10, np.NaN, cls.window_test_start_date),
(2, 220, cls.window_test_start_date), (20, 220, cls.window_test_start_date),
(2, 221, pd.Timestamp('2015-01-17'))], (20, 221, pd.Timestamp('2015-01-17'))],
pd.Timestamp('2015-01-20') pd.Timestamp('2015-01-20')
), ),
]) ])
@@ -1066,14 +1080,14 @@ class NextEstimateWindows(WithEstimateWindows, ZiplineTestCase):
twoq_next = pd.concat( twoq_next = pd.concat(
[cls.create_expected_df( [cls.create_expected_df(
[(0, 200, pd.Timestamp(cls.window_test_start_date)), [(0, 200, pd.Timestamp(cls.window_test_start_date)),
(1, np.NaN, pd.Timestamp(cls.window_test_start_date)), (10, np.NaN, pd.Timestamp(cls.window_test_start_date)),
(2, 220, pd.Timestamp(cls.window_test_start_date))], (20, 220, pd.Timestamp(cls.window_test_start_date))],
pd.Timestamp('2015-01-09') pd.Timestamp('2015-01-09')
)] + )] +
[cls.create_expected_df( [cls.create_expected_df(
[(0, np.NaN, pd.Timestamp(cls.window_test_start_date)), [(0, np.NaN, pd.Timestamp(cls.window_test_start_date)),
(1, np.NaN, pd.Timestamp(cls.window_test_start_date)), (10, np.NaN, pd.Timestamp(cls.window_test_start_date)),
(2, np.NaN, pd.Timestamp(cls.window_test_start_date))], (20, np.NaN, pd.Timestamp(cls.window_test_start_date))],
end_date end_date
) for end_date in pd.date_range('2015-01-12', '2015-01-20')] ) for end_date in pd.date_range('2015-01-12', '2015-01-20')]
) )