def test_PurgedKFold_3_with_purge_without_embargo():
df = pd.DataFrame({'A': [0, 1, 2, 3, 4, 5]})
pkf = fl.PurgedKFold(n_splits=3, n_overlaps=1)
idxs = list(map(lambda x: (list(x[0]), list(x[1])), pkf.split(df)))
assert idxs == [([3, 4, 5], [0, 1]), ([0, 5], [2, 3]), ([0, 1, 2], [4, 5])]
def test_PurgedKFold_3_no_purge_no_embargo():
df = pd.DataFrame({'A': [0, 1, 2]})
pkf = fl.PurgedKFold(n_splits=3)
idxs = list(map(lambda x: (list(x[0]), list(x[1])), pkf.split(df)))
assert idxs == [([1, 2], [0]), ([0, 2], [1]), ([0, 1], [2])]
def test_PurgedKFold_3_with_purge_embargo():
df = pd.DataFrame({'A': [0, 1, 2, 3, 4, 5, 6]})
pkf = fl.PurgedKFold(n_splits=3, n_overlaps=1, pct_embargo=0.16)
idxs = list(map(lambda x: (list(x[0]), list(x[1])), pkf.split(df)))
assert idxs == [([4, 5, 6], [0, 1]), ([0, 1, 6], [3]),
([0, 1, 2, 3], [5, 6])]
def test_PurgedKFold_2_special_index():
df = pd.DataFrame({'A': [0, 1, 2, 3]})
df.index = [0, 10, 100, 1000]
pkf = fl.PurgedKFold(n_splits=2)
idxs = list(map(lambda x: (list(x[0]), list(x[1])), pkf.split(df)))
assert idxs == [([100, 1000], [0, 10]), ([0, 10], [100, 1000])]
def test_PurgedKFold_2_with_embargo_without_purge():
df = pd.DataFrame({'A': [0, 1, 2, 3]})
pkf = fl.PurgedKFold(n_splits=2, n_overlaps=0, pct_embargo=0.25)
idxs = list(map(lambda x: (list(x[0]), list(x[1])), pkf.split(df)))
assert idxs == [([2, 3], [0]), ([0, 1], [2, 3])]