| """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 | |