| """Shared utilities: trajectory bounds, nested splits, Wilson CI, seeding.""" | |
| import math | |
| import numpy as np | |
| def get_traj_bounds(terminals): | |
| """Return list of (start, end_exclusive) for each trajectory.""" | |
| ends = np.where(np.asarray(terminals) > 0.5)[0] | |
| bounds, start = [], 0 | |
| for e in ends: | |
| bounds.append((int(start), int(e) + 1)) | |
| start = e + 1 | |
| return bounds | |
| def nested_traj_order(n_traj, seed=0): | |
| """Fixed permutation of trajectory ids so that prefixes are NESTED across demo sizes.""" | |
| return np.random.default_rng(seed).permutation(n_traj) | |
| def wilson_ci(k, n, z=1.96): | |
| """95% Wilson score interval for a binomial proportion.""" | |
| if n == 0: | |
| return 0.0, 0.0, 0.0 | |
| p = k / n | |
| denom = 1.0 + z * z / n | |
| center = (p + z * z / (2 * n)) / denom | |
| half = (z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))) / denom | |
| return p, max(0.0, center - half), min(1.0, center + half) | |
| def seed_everything(seed): | |
| import torch | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed_all(seed) | |