def test_split_apply_check_end_size(mixed_size_scramble_data):
result = split_apply(mixed_size_scramble_data, ["A"], shuffle)
assert result.shape[0] == mixed_size_scramble_data.shape[0]
def test_split_apply_shuffle_some_small(mixed_size_scramble_data):
result = split_apply(mixed_size_scramble_data, ["A"], shuffle)
assert result["B"].isna().any()
def test_split_apply_shuffle_breaks_corr(split_apply_frm):
result = split_apply(split_apply_frm, ["ID"], shuffle)
cov_matrix = np.cov(result["Value"], split_apply_frm["Value"])
assert cov_matrix[0, 0] != cov_matrix[0, 1]
def test_split_apply_shuffle(split_apply_frm):
raw_means = split_apply_frm.groupby("ID").mean()
result = split_apply(split_apply_frm, ["ID"], shuffle)
shuffle_means = result.groupby("ID").mean()
assert (shuffle_means == raw_means).values.all()
def test_split_apply_two_levels(split_apply_frm, split_apply_increment_frm):
"""Check multi level split_apply performs as expected"""
result = split_apply(split_apply_frm, ["ID", "ID2"], split_function)
assert (result == split_apply_increment_frm).values.all()
def test_split_apply2(split_apply_frm, split_apply_increment_frm):
"""Check split_apply performs as expected"""
result = split_apply(split_apply_frm, "ID2", split_function)
assert (result == split_apply_increment_frm).values.all()