Six scans from four patients. Split the rows, and then split the patients, and look at who ends up where.
rows = [('ann', 1), ('ann', 2), ('bob', 3),
('bob', 4), ('cid', 5), ('dot', 6)]
# Splitting rows: ann can easily land on both sides
# train: ('ann', 1) ('bob', 3) ('cid', 5)
# test: ('ann', 2) ('bob', 4) ('dot', 6)
# -> the model has seen ann and bob before
# Splitting groups: shuffle ['ann', 'bob', 'cid', 'dot'], then take rows
# test groups {'bob', 'dot'}
# train: ann's two rows and cid's one
# test: bob's two rows and dot's one
# -> no patient is in both halves
- The unit of independence is not always the row
Every evaluation assumes the test set is data the model has not learned from, and 'has not learned from' means the underlying thing, not the particular row. Two scans of one patient share almost everything — anatomy, scanner, positioning — so a model that memorises the patient scores well without learning anything about the disease. Ask what the model could cheat by recognising; that is your group.
- Shuffle and split the groups, then collect their rows
The shape is the same as train-test-split with one level of indirection: the list being divided is the distinct groups, and the rows follow whichever half their group went to. That is also why the fraction now means a fraction of groups — you cannot promise an exact share of rows when groups have different sizes.
- The halves will not be the sizes you asked for
Splitting four patients in half gives two each, and if one has fifteen scans and another has three, the row counts are nowhere near equal. That is the price, and it is worth paying: an imbalanced but honest split beats a perfectly balanced one that lies. On very uneven group sizes it is worth checking the split you got rather than assuming.
The thing being shuffled and divided is wrong. It should not be the rows — work out what the list of independent units actually is, and split that instead.
Collect the distinct group names in first-seen order, shuffle those, and take the first round(len(groups) * test_fraction) of them as the test groups. Then each row goes to whichever half its group is in.
dict.fromkeys gives the distinct groups in first-seen order, which keeps the shuffle reproducible — a set would make the order depend on hashing. Build the two halves with complementary comprehensions over the original rows: `if group in test_groups` and `if group not in test_groups`, which preserves the row order and cannot lose a row.
import random def split_by_group(rows, test_fraction=0.2, seed=0): if not 0 <= test_fraction <= 1: raise ValueError('test_fraction must be between 0 and 1') groups = list(dict.____(group for group, _ in rows)) random.Random(seed).shuffle(____) size = round(len(____) * test_fraction) held_back = ____(groups[:size]) return ( [row for row in rows if row[0] ____ held_back], [row for row in rows if row[0] in held_back], )
import random
def split_by_group(rows, test_fraction=0.2, seed=0):
if not 0 <= test_fraction <= 1:
raise ValueError(
f"test_fraction must be between 0 and 1, got {test_fraction}"
)
# The groups are what gets split. fromkeys keeps first-seen order, so
# the shuffle depends on the seed and not on how names happen to hash.
groups = list(dict.fromkeys(group for group, _ in rows))
random.Random(seed).shuffle(groups)
# A fraction of the groups, not of the rows.
size = round(len(groups) * test_fraction)
held_back = set(groups[:size])
# Complementary conditions over the original rows: order preserved,
# every row in exactly one half, no group in both.
return (
[row for row in rows if row[0] not in held_back],
[row for row in rows if row[0] in held_back],
)
set(...) would give the group names in an order that depends on hashing, and for strings that varies between runs unless PYTHONHASHSEED is fixed — so shuffling it would produce a different split from the same seed. dict.fromkeys gives a deterministic first-seen order, and the seed is then the only thing deciding the split. The set is built afterwards, where only membership matters.
`in held_back` and `not in held_back` cover every row exactly once by construction, and iterating the original rows keeps their order — so the coverage property is guaranteed rather than checked. It walks the rows twice, which is irrelevant next to the guarantee, and the membership test is against a set so each pass is linear.
Half the patients is not half the scans, and on real data the test half can easily be a third or two thirds of the rows. The starter’s split was beautifully balanced and measured the wrong thing. When a fix makes the numbers look worse and the reasoning look better, the numbers were the problem.
sklearn has GroupShuffleSplit and GroupKFold, which take a groups array alongside X and y and do exactly this — use them, and use StratifiedGroupKFold when the target is also imbalanced. The hard part was never the code: it is identifying the group, and no library can do that for you. Ask what the model could score well by recognising. Patient, user, product, document, session, site, device, and — the one people miss most — time, where the group is the day and the correct split is chronological rather than random. If several of those apply at once, you have a genuinely difficult evaluation design problem, and the honest move is to write down which one you split on and why.
- Take the group as a key function rather than assuming it is the first element, so the same split works on dictionaries or objects.
- Report the row counts of both halves so a caller can see how uneven the split came out, and decide at what point they should re-seed.
- Build the k-fold version — every group in exactly one test fold — by reusing the k-fold-split arithmetic on the list of groups.
- Swap dict.fromkeys for a set, run the split in two separate processes, and confirm for yourself that the seed stops deciding the answer — string hashing is randomised per process unless PYTHONHASHSEED is fixed, which is why this one cannot be caught by a test inside a single run.