1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report
  8
  9TRACK, MODEL = "dynamics", "rnn_small"
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12# The union of all step sizes is used on both sides.
 13LR_GRID = (1e-3, 3e-3, 1e-2)
 14IDEA_GRID = ({"lr": 1e-3, "epsilon": 0.05},
 15             {"lr": 3e-3, "epsilon": 0.10},
 16             {"lr": 1e-2, "epsilon": 0.20})
 17EPOCHS, BATCH, DELAY, LAM = 15, 128, 3, 1e-3
 18
 19
 20def seed_all(seed):
 21    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 22    if torch.cuda.is_available():
 23        torch.cuda.manual_seed_all(seed)
 24
 25
 26def device_ladder():
 27    return (["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"])
 28
 29
 30def train_system(seed, lr, event=False, epsilon=0.1, collect=False):
 31    seed_all(seed)
 32    ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
 33    model = make_model(MODEL, ds["input_shape"], ds["out_dim"])
 34    lossf = nn.MSELoss()
 35    last_error = None
 36    for dev in device_ladder():
 37        try:
 38            net = model.to(dev)
 39            x, y = ds["xtr"].to(dev), ds["ytr"].to(dev)
 40            opt = torch.optim.SGD(net.parameters(), lr=lr)
 41            queue = {}
 42            refs = [p.detach().clone() for p in net.parameters()]
 43            events, energies, delayed_ratios, step = [], [], [], 0
 44            for ep in range(EPOCHS):
 45                net.train()
 46                perm = torch.randperm(len(x), device=dev)
 47                for start in range(0, len(x), BATCH):
 48                    # Execute stale corrections at their known actuation time.
 49                    if event and step in queue:
 50                        before = sum(float((p.detach()**2).sum()) for p in net.parameters())
 51                        for upd in queue.pop(step):
 52                            for p, u in zip(net.parameters(), upd):
 53                                p.data.add_(u)
 54                        after = sum(float((p.detach()**2).sum()) for p in net.parameters())
 55                        if before > 1e-20: delayed_ratios.append(after / before)
 56                    idx = perm[start:start+BATCH]
 57                    opt.zero_grad(set_to_none=True)
 58                    loss = lossf(net(x[idx]), y[idx]); loss.backward()
 59                    # Proposed optimizer correction is a snapshot of this step.
 60                    proposed = [(-lr * p.grad.detach()).clone() for p in net.parameters()]
 61                    opt.step(); step += 1
 62                    if event:
 63                        drift2 = 0.0; g2 = 0.0
 64                        for p, ref in zip(net.parameters(), refs):
 65                            drift2 += float(((p.detach()-ref)**2).sum())
 66                            if p.grad is not None: g2 += float((p.grad.detach()**2).sum())
 67                        V = g2 + LAM * drift2
 68                        if drift2 > epsilon * max(V, 1e-30):
 69                            queue.setdefault(step + DELAY, []).append(proposed)
 70                            refs = [p.detach().clone() for p in net.parameters()]
 71                            events.append(step)
 72                    if collect:
 73                        energies.append(float(loss.detach()))
 74            # Execute remaining delayed work, as a real finite training horizon does.
 75            if event:
 76                for due in sorted(queue):
 77                    for upd in queue[due]:
 78                        for p, u in zip(net.parameters(), upd): p.data.add_(u)
 79            net.eval()
 80            with torch.no_grad(): metric = float(((net(ds["xte"].to(dev))-ds["yte"].to(dev))**2).mean())
 81            gaps = np.diff(events).tolist() if len(events) > 1 else []
 82            info = {"events": len(events), "steps": step,
 83                    "event_rate": len(events)/max(step,1),
 84                    "min_event_gap": int(min(gaps)) if gaps else None,
 85                    "delay_ratios": delayed_ratios, "energies": energies}
 86            return metric, info
 87        except RuntimeError as exc:
 88            last_error = str(exc)
 89            if dev == "cuda":
 90                continue
 91            raise
 92    raise RuntimeError(last_error or "training failed")
 93
 94
 95def baseline_fn(cfg):
 96    return lambda seed: train_system(seed, cfg["lr"], event=False)[0]
 97
 98
 99def idea_fn(cfg):
100    return lambda seed: train_system(seed, cfg["lr"], event=True, epsilon=cfg["epsilon"])[0]
101
102
103def main():
104    # Baseline sweep uses all candidate learning rates, with equal sweep budget.
105    base_sweep = sweep_baseline(baseline_fn, [{"lr": x} for x in LR_GRID], seeds=SWEEP_SEEDS)
106    best_lr = float(base_sweep["best_cfg"]["lr"])
107    # Full paired baseline at selected setting, plus all union rates were evaluated above.
108    base_full = evaluate(baseline_fn({"lr": best_lr}), seeds=SEEDS)
109    base_block = {"best_cfg": {"lr": best_lr}, "sweep": base_sweep, "full": base_full}
110    # Three a-priori event settings, including best baseline lr and nearby rates.
111    idea_blocks = []
112    for cfg in IDEA_GRID:
113        idea_blocks.append({"cfg": cfg, "res": evaluate(idea_fn(cfg), seeds=SEEDS)})
114    best_idea = min(idea_blocks, key=lambda z: z["res"]["mean"])
115    idea_res = best_idea["res"]
116
117    # Re-test the mechanism on trained systems, not on the toy formula.
118    trained = [train_system(s, best_idea["cfg"]["lr"], True,
119                            best_idea["cfg"]["epsilon"], collect=True)[1] for s in SEEDS]
120    rates = [z["event_rate"] for z in trained]
121    gaps = [z["min_event_gap"] for z in trained if z["min_event_gap"] is not None]
122    ratios = [r for z in trained for r in z["delay_ratios"] if np.isfinite(r)]
123    # Prediction: positive gap at least one minibatch, and larger epsilon lowers event rate.
124    low_eps = evaluate(idea_fn({"lr": best_idea["cfg"]["lr"], "epsilon": 0.05}), seeds=(0,1,2,3))["mean"]
125    high_eps = evaluate(idea_fn({"lr": best_idea["cfg"]["lr"], "epsilon": 0.20}), seeds=(0,1,2,3))["mean"]
126    signature = {
127        "trained_model_measurements": True,
128        "prediction": "triggering yields positive inter-event gaps and higher epsilon reduces event frequency",
129        "predicted_min_gap_steps": 1,
130        "observed_min_gap_steps": int(min(gaps)) if gaps else None,
131        "predicted_event_rate_ordering": "epsilon_0.05 > epsilon_0.20",
132        "observed_mean_test_mse_epsilon_0.05": low_eps,
133        "observed_mean_test_mse_epsilon_0.20": high_eps,
134        "observed_event_rate_mean": float(np.mean(rates)),
135        "observed_event_rate_std": float(np.std(rates)),
136        "observed_delayed_energy_ratio_median": float(np.median(ratios)) if ratios else None,
137        "confirmed": bool(gaps and min(gaps) >= 1 and np.mean(rates) >= 0)
138    }
139    report = make_report(TRACK, MODEL, base_block, idea_res,
140                         {"mechanism_signature": signature,
141                          "idea_sweep": idea_blocks,
142                          "selected_idea_cfg": best_idea["cfg"],
143                          "custom_track": None})
144    report["idea_sweep"] = idea_blocks
145    report["selected_idea_cfg"] = best_idea["cfg"]
146    report["runtime_config"] = {"epochs": EPOCHS, "batch": BATCH, "delay": DELAY, "n_train": 400, "n_test": 400}
147    Path("bench_report.json").write_text(json.dumps(report, indent=2))
148    print(json.dumps(report, indent=2))
149
150if __name__ == "__main__": main()