def testTransformWithDifferentGroup(self):
test_df = pd.DataFrame({'code': [1, 2, 3, 4, 5, 6, 7],
'b': [4, 5, 6, 7, 6, 5, 4],
'c': [9, 8, 7, 6, 5, 4, 3]},
index=[1, 1, 1, 1, 2, 2, 2],
dtype=float)
expression = SecurityMovingAverage(2, 'b')
calculated = expression.transform(test_df, name='new_factor', category_field='code')
expected = [4., 5., 6., 7., 6., 5., 4.]
np.testing.assert_array_almost_equal(calculated['new_factor'], expected)
def testTransformWithoutCategoryAndDuplicatedIndex(self):
sample_data = pd.DataFrame(
data={'code': [1, 2, 1, 2, 1, 2],
'open': [2.0, 1.0, 1.5, 3.0, 2.4, 3.5],
'close': [1.7, 1.6, 0.9, 3.8, 1.6, 2.1]},
index=[1, 1, 2, 2, 3, 3],
dtype=float
)
expression = SecurityMovingAverage(2, 'close')
calculated = expression.transform(sample_data, name='new_factor')
expected = sample_data.rolling(window=2).mean()
np.testing.assert_array_almost_equal(calculated['new_factor'].values[1:],
expected['close'].values[1:])
