File size: 1,086 Bytes
87d0bc2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """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)
|