def test_impute_transform_with_groupby(): data = pd.DataFrame({ 'x': [1, 2, 4, 1, 3, 4], 'y': [1, 2, 4, 2, 4, 5], 'cat': list('AAABBB') }) transform = alt.ImputeTransform(impute='y', key='x', method='max', groupby=['cat']) imputed = apply(data, transform) assert_equal(imputed.x.values, np.tile(range(1, 5), 2)) assert_equal(imputed.y.values, [1, 2, 4, 4, 2, 5, 4, 5])
def test_impute_transform_with_groupby(): data = pd.DataFrame({ "x": [1, 2, 4, 1, 3, 4], "y": [1, 2, 4, 2, 4, 5], "cat": list("AAABBB") }) transform = alt.ImputeTransform(impute="y", key="x", method="max", groupby=["cat"]) imputed = apply(data, transform) assert_equal(imputed.x.values, np.tile(range(1, 5), 2)) assert_equal(imputed.y.values, [1, 2, 4, 4, 2, 5, 4, 5])
def test_impute_transform_no_groupby(method): data = pd.DataFrame({'x': [1, 2], 'y': [2, 3]}) transform = alt.ImputeTransform(impute='y', key='x', keyvals={ 'start': 0, 'stop': 5 }, value=0, method=method) if method == 'value': value = 0 else: value = data.y.agg(method) imputed = apply(data, transform) assert_equal(imputed.x.values, range(5)) assert_equal(imputed.y[[1, 2]].values, data.y.values) assert_equal(imputed.y[[0, 3, 4]].values, value)
def test_impute_transform_no_groupby(method): data = pd.DataFrame({"x": [1, 2], "y": [2, 3]}) transform = alt.ImputeTransform(impute="y", key="x", keyvals={ "start": 0, "stop": 5 }, value=0, method=method) if method == "value": value = 0 else: value = data.y.agg(method) imputed = apply(data, transform) assert_equal(imputed.x.values, range(5)) assert_equal(imputed.y[[1, 2]].values, data.y.values) assert_equal(imputed.y[[0, 3, 4]].values, value)