1import sys, json, random, math
  2import numpy as np
  3import torch
  4import torch.nn.functional as F
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
  8
  9SEED0 = 1215
 10EPOCHS = 12
 11BATCH = 128
 12THREADS = 4
 13
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available():
 18        torch.cuda.manual_seed_all(seed)
 19
 20
 21def js_batch(p, q):
 22    m = (p + q) / 2
 23    return 0.5 * ((p * (p.clamp_min(1e-8).log() - m.clamp_min(1e-8).log())).sum(-1) +
 24                  (q * (q.clamp_min(1e-8).log() - m.clamp_min(1e-8).log())).sum(-1))
 25
 26
 27def math_check():
 28    p = torch.tensor([[.9, .1], [.9, .1], [.2, .8]], dtype=torch.float64)
 29    z = float(js_batch(p[:1], p[1:2])[0])
 30    unequal = float(js_batch(p[:1], p[2:3])[0])
 31    sets = [{0, 1}, {0}, {0, 2}]
 32    empty = [{0}, {1}]
 33    return {
 34        "js_equal": z, "js_unequal": unequal,
 35        "js_symmetric_error": abs(float(js_batch(p[:1], p[2:3])[0] - js_batch(p[2:3], p[:1])[0])),
 36        "intersection_feasible": bool(set.intersection(*map(set, sets))),
 37        "intersection_empty": not bool(set.intersection(*map(set, empty))),
 38        "passed": z < 1e-12 and unequal > 0 and not bool(set.intersection(*map(set, empty))),
 39    }
 40
 41
 42def grouped_consistency(pred, x):
 43    # Dynamics observations are aliased by a coarse theta projection. The
 44    # scalar forecast induces a distribution over three forcing-mask choices.
 45    theta_obs = x[:, -3]
 46    keys = torch.round(theta_obs / 0.35).to(torch.int64)
 47    centers = pred.new_tensor([-0.35, 0.0, 0.35])
 48    q = F.softmax(-((pred - centers[None, :]) / 0.30) ** 2, dim=1)
 49    vals = []
 50    for k in torch.unique(keys):
 51        ix = keys == k
 52        if int(ix.sum()) > 1:
 53            bar = q[ix].mean(0, keepdim=True)
 54            vals.append(js_batch(q[ix], bar.expand_as(q[ix])).mean())
 55    return torch.stack(vals).mean() if vals else pred.new_zeros(())
 56
 57
 58def robust_safety(pred, x):
 59    # Candidate forcing corrections; require one shared correction to keep all
 60    # histories in each observation class inside the safe angular region.
 61    theta_obs = x[:, -3]
 62    keys = torch.round(theta_obs / 0.35).to(torch.int64)
 63    deltas = pred.new_tensor([-0.20, 0.0, 0.20])
 64    margins = torch.sigmoid((1.20 - (pred[:, None] + deltas[None, :]).abs()) / 0.12)
 65    losses = []
 66    for k in torch.unique(keys):
 67        ix = keys == k
 68        if int(ix.sum()) > 1:
 69            worst = margins[ix].min(0).values
 70            best = worst.max()
 71            losses.append(F.relu(pred.new_tensor(0.75) - best))
 72    return torch.stack(losses).mean() if losses else pred.new_zeros(())
 73
 74
 75def train_one(seed, lr, weight_decay=0.0, fc_weight=0.0, safety_weight=0.0, return_sig=False):
 76    seed_all(seed)
 77    ds = get_dataset("dynamics", seed, n_train=400, n_test=200)
 78    model = make_model("rnn_small", ds["input_shape"], ds["out_dim"])
 79    device = "cuda" if torch.cuda.is_available() else "cpu"
 80    try:
 81        model.to(device)
 82        opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
 83        x, y = ds["xtr"].to(device), ds["ytr"].to(device)
 84        model.train()
 85        for _ in range(EPOCHS):
 86            order = torch.randperm(len(x), device=device)
 87            for start in range(0, len(x), BATCH):
 88                ix = order[start:start+BATCH]
 89                pred = model(x[ix])
 90                task = F.mse_loss(pred, y[ix])
 91                loss = task
 92                if fc_weight:
 93                    loss = loss + fc_weight * grouped_consistency(pred, x[ix])
 94                if safety_weight:
 95                    loss = loss + safety_weight * robust_safety(pred, x[ix])
 96                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 97        model.eval()
 98        with torch.no_grad():
 99            pred = model(ds["xte"].to(device))
100            metric = float(F.mse_loss(pred, ds["yte"].to(device)).cpu())
101            # Signature measured on the trained model, not an analytic toy.
102            kappa = float(grouped_consistency(pred, ds["xte"].to(device)).cpu())
103            safety = float(robust_safety(pred, ds["xte"].to(device)).cpu())
104        if return_sig:
105            return metric, {"observed_kappa": kappa, "observed_safety_hinge": safety}
106        return metric
107    except RuntimeError:
108        if device != "cpu":
109            torch.cuda.empty_cache()
110            return train_one(seed, lr, weight_decay, fc_weight, safety_weight, return_sig)
111        raise
112
113
114def main():
115    check = math_check()
116    assert check["passed"]
117    # Baseline decisive knob (weight decay) and learning rate are both swept.
118    grid = [{"lr": lr, "weight_decay": wd}
119            for lr in (1e-2, 3e-2, 1e-1) for wd in (0.0, 1e-4)]
120    base = sweep_baseline(
121        lambda c: lambda s: train_one(s, c["lr"], c["weight_decay"]), grid)
122    best_lr = base["best_cfg"]["lr"]
123    idea_grid = [{"lr": lr, "weight_decay": 0.0, "fc_weight": 0.08, "safety_weight": 0.08}
124                 for lr in (1e-2, best_lr, 1e-1)]
125    idea_trials = []
126    for cfg in idea_grid:
127        r = evaluate(lambda s, c=cfg: train_one(s, c["lr"], c["weight_decay"], c["fc_weight"], c["safety_weight"]), (0,1,2,3))
128        idea_trials.append({"cfg": cfg, "mean": r["mean"]})
129    best_idea_cfg = min(idea_trials, key=lambda z: z["mean"])["cfg"]
130    idea = evaluate(lambda s: train_one(s, best_idea_cfg["lr"], best_idea_cfg["weight_decay"], best_idea_cfg["fc_weight"], best_idea_cfg["safety_weight"]))
131    _, sigs = zip(*(train_one(s, best_idea_cfg["lr"], best_idea_cfg["weight_decay"], best_idea_cfg["fc_weight"], best_idea_cfg["safety_weight"], True) for s in range(8)))
132    base_sig = [train_one(s, base["best_cfg"]["lr"], base["best_cfg"]["weight_decay"], return_sig=True)[1] for s in range(8)]
133    signature = {
134        "prediction": "FC should lower observational kappa on trained recurrent models",
135        "baseline_kappa_mean": float(np.mean([z["observed_kappa"] for z in base_sig])),
136        "idea_kappa_mean": float(np.mean([z["observed_kappa"] for z in sigs])),
137        "baseline_safety_hinge_mean": float(np.mean([z["observed_safety_hinge"] for z in base_sig])),
138        "idea_safety_hinge_mean": float(np.mean([z["observed_safety_hinge"] for z in sigs])),
139        "confirmed": float(np.mean([z["observed_kappa"] for z in sigs])) < float(np.mean([z["observed_kappa"] for z in base_sig]))
140    }
141    report = make_report("dynamics", "rnn_small", base, idea, {
142        "math_check": check, "idea_sweep": idea_trials,
143        "best_idea_cfg": best_idea_cfg, "signature": signature
144    })
145    with open("bench_report.json", "w") as f: json.dump(report, f, indent=2)
146    print(json.dumps(report, indent=2))
147
148if __name__ == "__main__": main()