1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEEDS = list(range(8))
  8DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
  9
 10
 11def required_rollouts(r, delta=0.05):
 12    r = min(max(float(r), 1e-12), 1 - 1e-12)
 13    return int(math.ceil(math.log(delta) / math.log1p(-r)))
 14
 15
 16def beta_lower(events, total, alpha=0.05):
 17    if total <= 0:
 18        return 0.0
 19    if events == 0:
 20        return 1.0 - alpha ** (1.0 / (total + 1.0))
 21    z = 1.644853626951
 22    p = events / total
 23    d = 1.0 + z * z / total
 24    return max(0.0, ((p + z*z/(2*total)) - z*math.sqrt(p*(1-p)/total + z*z/(4*total*total))) / d)
 25
 26
 27def true_transition(x, u):
 28    th, om = x[:, 0], x[:, 1]
 29    next_om = om + 0.12*u - 0.08*torch.sin(th) - 0.025*om
 30    next_th = th + 0.10*next_om
 31    event = next_th.abs() > 1.15
 32    next_om = torch.where(event, -0.65*next_om, next_om)
 33    next_th = torch.clamp(next_th, -1.5, 1.5)
 34    return torch.stack((next_th, next_om), 1), event
 35
 36
 37class TinyRNN(nn.Module):
 38    def __init__(self, hidden=24):
 39        super().__init__()
 40        self.rnn = nn.GRU(3, hidden, batch_first=True)
 41        self.head = nn.Linear(hidden, 2)
 42
 43    def forward(self, z):
 44        return self.head(self.rnn(z)[0][:, -1])
 45
 46
 47def make_data(seed, n):
 48    g = torch.Generator().manual_seed(seed)
 49    x = torch.empty(n, 1, 2).uniform_(-1.5, 1.5, generator=g)
 50    u = torch.empty(n, 1, 1).uniform_(-1.0, 1.0, generator=g)
 51    y, event = true_transition(x[:, 0], u[:, 0, 0])
 52    return torch.cat((x, u), 2), y, event
 53
 54
 55def train_model(seed, lr, idea, epochs=35):
 56    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
 57    x, y, _ = make_data(seed, 400)
 58    model = TinyRNN().to(DEVICE)
 59    opt = torch.optim.Adam(model.parameters(), lr=lr)
 60    x, y = x.to(DEVICE), y.to(DEVICE)
 61    for _ in range(epochs):
 62        opt.zero_grad()
 63        pred = model(x)
 64        loss = (pred-y).pow(2).mean()
 65        if idea:
 66            # Risk-calibrated transition loss: oversample high-risk boundary states
 67            # and weight their transition error, a train-time proxy for directed probes.
 68            boundary = x[:, 0, 0].abs() > 0.95
 69            if boundary.any():
 70                loss = 0.65*loss + 0.35*(pred[boundary]-y[boundary]).pow(2).mean()
 71        loss.backward(); opt.step()
 72    return model
 73
 74
 75def evaluate(model, seed):
 76    x, y, event = make_data(seed+1000, 500)
 77    with torch.no_grad():
 78        pred = model(x.to(DEVICE)).cpu()
 79    mse = float((pred-y).pow(2).mean())
 80    # Boundary stress distribution is the planner-directed probe distribution.
 81    g = torch.Generator().manual_seed(seed+2000)
 82    xs = torch.empty(300, 1, 2).uniform_(-1.5, 1.5, generator=g)
 83    xs[:, 0, 0] = torch.where(torch.rand(300, generator=g) < .7,
 84                              torch.empty(300).uniform_(1.0, 1.35, generator=g), xs[:,0,0])
 85    us = torch.empty(300, 1, 1).uniform_(-1, 1, generator=g)
 86    z = torch.cat((xs, us), 2)
 87    truth, ev = true_transition(xs[:,0], us[:,0,0])
 88    with torch.no_grad(): pp = model(z.to(DEVICE)).cpu()
 89    predicted_event = (pp[:,0].abs() > 1.15) | ((pp[:,1]-xs[:,0,1]).abs() > .35)
 90    tp = int((predicted_event & ev).sum()); fn = int((~predicted_event & ev).sum())
 91    return {"mse": mse, "event_rate": float(ev.float().mean()), "event_detect_rate": float(tp/max(1,tp+fn)), "probe_error": float((pp-truth).pow(2).mean())}
 92
 93
 94def permutation_p(deltas, reps=20000):
 95    rng=np.random.default_rng(991); obs=abs(float(np.mean(deltas))); n=len(deltas)
 96    signs=rng.choice([-1,1],size=(reps,n)); null=np.abs((signs*np.asarray(deltas)).mean(1))
 97    return float((1+np.sum(null>=obs))/(reps+1))
 98
 99
100def run():
101    # Union search space is evaluated for both systems: baseline lr sweep and two nearby values.
102    lrs=[0.001,0.003,0.006]
103    all_results={"baseline":{},"idea":{}}
104    for lr in lrs:
105        all_results["baseline"][str(lr)] = []
106        all_results["idea"][str(lr)] = []
107        for s in SEEDS:
108            all_results["baseline"][str(lr)].append(evaluate(train_model(s,lr,False),s))
109            all_results["idea"][str(lr)].append(evaluate(train_model(s,lr,True),s))
110    def mean(key, lr, side): return float(np.mean([r[key] for r in all_results[side][str(lr)]]))
111    best_b=min(lrs,key=lambda l:mean("mse",l,"baseline")); best_i=min(lrs,key=lambda l:mean("mse",l,"idea"))
112    b=np.array([r["mse"] for r in all_results["baseline"][str(best_b)]])
113    i=np.array([r["mse"] for r in all_results["idea"][str(best_i)]])
114    delta=i-b
115    # Mechanism signature is measured from trained models, not an analytical identity.
116    eb=np.array([r["event_detect_rate"] for r in all_results["baseline"][str(best_b)]])
117    ei=np.array([r["event_detect_rate"] for r in all_results["idea"][str(best_i)]])
118    observed_miss=float(np.mean(1-ei)); predicted_miss=float(np.mean(1-eb))
119    signature={"predicted_probe_miss_reduction": predicted_miss-observed_miss, "observed_baseline_mse":float(b.mean()), "observed_idea_mse":float(i.mean()), "required_rollouts_at_r_0.02":required_rollouts(.02), "confirmed": bool((ei.mean()>eb.mean()) and (observed_miss<predicted_miss))}
120    report={"track":"dynamics","device":DEVICE,"seeds":SEEDS,"baseline_sweep":{str(l):mean("mse",l,"baseline") for l in lrs},"idea_sweep":{str(l):mean("mse",l,"idea") for l in lrs},"best_baseline_lr":best_b,"best_idea_lr":best_i,"baseline_per_seed":all_results["baseline"][str(best_b)],"idea_per_seed":all_results["idea"][str(best_i)],"paired_delta_mean":float(delta.mean()),"permutation_p_value":permutation_p(delta),"mechanism_signature":signature,"harness_status":"unavailable: advertised /home/maxwelhelp/all/math2nn/bench was absent and could not be imported","custom_track":None}
121    Path("bench_report.json").write_text(json.dumps(report,indent=2))
122    print(json.dumps(report,indent=2))
123
124if __name__ == "__main__":
125    try:
126        run()
127    except Exception as exc:
128        if DEVICE == "cuda":
129            print("CUDA failure; retrying on CPU:", repr(exc))
130            DEVICE = "cpu"
131            run()
132        else:
133            raise