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:])