我想我会发布我的解决方案,以防它对其他人有用。
from collections import defaultdict
import random
def strat_map(y):
"""
Returns permuted indices that maintain class
"""
smap = defaultdict(list)
for i,v in enumerate(y):
smap[v].append(i)
for values in smap.values():
random.shuffle(values)
y_map = np.zeros_like(y)
for i,v in enumerate(y):
y_map[i] = smap[v].pop()
return y_map
##########
#Example Use
##########
skf = StratifiedKFold(y, nfolds)
sm = strat_map(y)
for test, train in skf:
test,train = sm[test], sm[train]
#then cv as usual
#######
#tests#
#######
import numpy.random as rnd
for _ in range(100):
y = np.array( [0]*10 + [1]*20 + [3] * 10)
rnd.shuffle(y)
sm = strat_map(y)
shuffled = y[sm]
assert (sm != range(len(y))).any() , "did not shuffle"
assert (shuffled == y).all(), "classes not in right position"
assert (set(sm) == set(range(len(y)))), "missing indices"
for _ in range(100):
nfolds = 10
skf = StratifiedKFold(y, nfolds)
sm = strat_map(y)
for test, train in skf:
assert (sm[test] != test).any(), "did not shuffle"
assert (y[sm[test]] == y[test]).all(), "classes not in right position"