File size: 519 Bytes
3c5cae5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 | """Shared dataset splitting so every script uses the same held-out events."""
import numpy as np
def split_by_event(df, test_frac=0.15, seed=42):
"""Return (train_mask, test_mask, test_events) grouped by whole events."""
rng = np.random.default_rng(seed)
events = np.array(df["event"].unique())
rng.shuffle(events)
n_test = max(1, int(len(events) * test_frac))
test_events = set(events[:n_test].tolist())
mask = df["event"].isin(test_events).to_numpy()
return ~mask, mask, test_events
|