3. Stratified Train-Validation-Test Split for Imbalanced Injury Classes

Why stratification matters for rare injury classes

Occupant injury datasets are strongly imbalanced: fatal or severe injuries (often coded as MAIS 3+ or KABCO "K/A") typically represent only a few percent of all records. A plain random split can easily place too few severe cases in the validation or test partition, making recall and AUC estimates unstable and non-comparable across experiments. Stratified splitting fixes the class proportions in every partition, so the rare class is represented consistently and the evaluation floor is reproducible.

This case builds a single, reproducible three-way split (train / validation / test) that preserves the severe-injury proportion in each partition, then verifies the balance with a compact audit table. The split is written to disk as an index file so every later Random Forest and XGBoost experiment reuses exactly the same partitions.

Inputs and assumptions

• Input table: a cleaned occupant-level dataframe df with one row per occupant and a binary label column severe_injury (1 = severe/fatal, 0 = otherwise).
• Row identifier: a stable key column occupant_id used to persist the split.
• Target proportions: 70% train, 15% validation, 15% test.
• Random seed: fixed at SEED = 42 for reproducibility.
• Environment: Python 3.10+, pandas, numpy, scikit-learn. No version-specific UI steps are assumed.

Two-stage stratified split

Scikit-learn's train_test_split supports stratification on a single call, so a three-way split is built as two sequential stratified splits: first isolate the test set, then split the remainder into train and validation. Both calls pass stratify= on the label so the rare class is preserved at each stage.

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

SEED = 42
TEST_SIZE = 0.15
VAL_SIZE = 0.15  # fraction of the full dataset

# df: cleaned occupant table with columns ['occupant_id', 'severe_injury', ...]
y = df['severe_injury'].astype(int)

# Stage 1: hold out the test set
train_val_idx, test_idx = train_test_split(
    df.index,
    test_size=TEST_SIZE,
    stratify=y,
    random_state=SEED,
)

# Stage 2: split the remainder into train and validation
# val fraction is expressed relative to the remaining rows
val_relative = VAL_SIZE / (1.0 - TEST_SIZE)
train_idx, val_idx = train_test_split(
    train_val_idx,
    test_size=val_relative,
    stratify=y.loc[train_val_idx],
    random_state=SEED,
)

split = pd.Series(index=df.index, dtype='object')
split.loc[train_idx] = 'train'
split.loc[val_idx]   = 'val'
split.loc[test_idx]  = 'test'
df['split'] = split

Verifying class balance in each partition

Stratification is only trustworthy once it is measured. The audit below reports row counts, severe-injury counts, and the severe-injury rate per partition, plus the absolute deviation from the overall rate. A well-behaved split keeps the deviation within roughly one percentage point for a rare class of a few percent.

overall_rate = y.mean()

audit = (
    df.groupby('split')['severe_injury']
      .agg(rows='size', severe='sum')
      .assign(rate=lambda d: d['severe'] / d['rows'])
      .assign(deviation=lambda d: (d['rate'] - overall_rate).abs())
      .loc[['train', 'val', 'test']]
)
print(audit.round(4))

For a teaching example with 20,000 occupants and a 3.0% severe-injury rate (600 severe cases), a correctly stratified split produces approximately the following — these are illustrative expected values, not measured results:

In practice the counts will not land on exact integers, so expect small deviations (typically under 0.3 pp for a 3% class). If any partition's severe count is zero, or the deviation exceeds about 1 pp, treat the split as failed and investigate before proceeding.

Persisting the split for downstream experiments

Every later model must train and evaluate on the same partitions, otherwise AUC and recall are not comparable. Persist the assignment as a CSV keyed by occupant_id and load it in each experiment rather than re-splitting.

split_map = df[['occupant_id', 'split']].copy()
split_map.to_csv('splits/occupant_split_seed42.csv', index=False)

# In a later experiment:
split_map = pd.read_csv('splits/occupant_split_seed42.csv')
df = df.merge(split_map, on='occupant_id', how='left', validate='one_to_one')
assert df['split'].notna().all(), 'Unassigned rows after merge'

Reading the split structure

The diagram below shows how the two-stage stratified split partitions the occupant table and how the severe-injury rate is preserved at each stage. It also marks the verification checkpoint and the persisted artifact that downstream experiments consume.

Failure symptoms and troubleshooting

• Zero severe cases in val or test: the rare class is too small for the chosen partition size. Either increase the dataset, reduce the number of partitions, or use a larger test/val fraction. Do not "fix" this by re-splitting until it looks good — that leaks test information into the design.
• Length mismatch on the second split: stratify was passed the full label vector instead of y.loc[train_val_idx].
• Deviation above ~1 pp: check that severe_injury is integer-coded (0/1) and not a float with NaNs; stratification silently degrades if the label has missing values.
• Different results across runs: confirm random_state=SEED is set on both calls and that the input row order is stable before splitting.