1import json, random, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report, count_params
  9
 10SEEDS = tuple(range(8))
 11# Shared union: every idea learning rate is also evaluated by baseline sweep.
 12GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 13EPOCHS = 12
 14NTRAIN, NTEST = 400, 400
 15
 16
 17def seed_all(seed):
 18    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 19    if torch.cuda.is_available():
 20        torch.cuda.manual_seed_all(seed)
 21
 22
 23class QuotientFibreGRU(nn.Module):
 24    """64-state triangular recurrent model: autonomous quotient and fibre."""
 25    def __init__(self, out_dim=1, qdim=32, fdim=42):
 26        super().__init__()
 27        self.qdim, self.fdim = qdim, fdim
 28        self.qcell = nn.GRUCell(3, qdim)
 29        self.fcell = nn.GRUCell(3 + qdim, fdim)
 30        self.head = nn.Linear(qdim + fdim, out_dim)
 31
 32    def forward(self, x, return_states=False):
 33        seq = x.view(x.shape[0], -1, 3)
 34        z = x.new_zeros((x.shape[0], self.qdim))
 35        y = x.new_zeros((x.shape[0], self.fdim))
 36        zs, ys = [], []
 37        for t in range(seq.shape[1]):
 38            u = seq[:, t]
 39            z = self.qcell(u, z)
 40            # Conditional fibre transition uses current quotient, never y to update z.
 41            y = self.fcell(torch.cat([u, z], dim=-1), y)
 42            zs.append(z); ys.append(y)
 43        out = self.head(torch.cat([z, y], dim=-1))
 44        if return_states:
 45            return out, torch.stack(zs, 1), torch.stack(ys, 1)
 46        return out
 47
 48
 49def train_one(kind, seed, lr):
 50    seed_all(seed)
 51    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 52    if kind == "baseline":
 53        model = make_model("rnn_small", ds["xtr"].shape[1:], ds["ytr"].shape[-1] if ds["ytr"].ndim > 1 else 1)
 54    else:
 55        model = QuotientFibreGRU(out_dim=1)
 56    _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 57    return float(metric)
 58
 59
 60def model_for_signature(kind, seed, lr):
 61    seed_all(seed)
 62    ds = get_dataset("dynamics", seed, n_train=NTRAIN, n_test=NTEST)
 63    model = (make_model("rnn_small", ds["xtr"].shape[1:], 1)
 64             if kind == "baseline" else QuotientFibreGRU(1))
 65    model, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
 66    return model, ds, metric
 67
 68
 69def fitted_rate(values, start=1):
 70    a = np.maximum(np.asarray(values, dtype=float), 1e-8)
 71    if len(a) <= start + 2: return float("nan")
 72    slope = np.polyfit(np.arange(start, len(a)), np.log(a[start:]), 1)[0]
 73    return float(np.exp(slope))
 74
 75
 76def mechanism_signature(best_lr):
 77    """Re-test max-rate and slow-branch invariance on trained benchmark models."""
 78    model, ds, _ = model_for_signature("idea", 0, best_lr)
 79    model.eval()
 80    dev = next(model.parameters()).device
 81    x = ds["xte"][:32].clone().to(dev)
 82    # Empirical perturbation of initial hidden states, with identical observed inputs.
 83    with torch.no_grad():
 84        _, z1, y1 = model(x, True)
 85        # Since the canonical model starts at zero, rerun with perturbed branch initial
 86        # states through the same trained cells to measure each branch's decay.
 87        seq = x.view(x.shape[0], -1, 3)
 88        z = torch.zeros((len(x), model.qdim), device=dev); y = torch.zeros((len(x), model.fdim), device=dev)
 89        z2 = torch.zeros_like(z); y2 = torch.zeros_like(y)
 90        z2[:, 0] = 1.0; y2[:, 0] = 1.0
 91        qnorm, fnorm, fullnorm = [], [], []
 92        for t in range(seq.shape[1]):
 93            u = seq[:, t]
 94            z = model.qcell(u, z); y = model.fcell(torch.cat([u, z], -1), y)
 95            z2 = model.qcell(u, z2); y2 = model.fcell(torch.cat([u, z2], -1), y2)
 96            qnorm.append((z2-z).norm(dim=1).mean().item())
 97            fnorm.append((y2-y).norm(dim=1).mean().item())
 98            fullnorm.append(torch.sqrt((z2-z).pow(2).sum(1)+(y2-y).pow(2).sum(1)).mean().item())
 99    aq, af, observed = fitted_rate(qnorm), fitted_rate(fnorm), fitted_rate(fullnorm)
100    pred = max(aq, af)
101    rel = abs(observed-pred) / max(abs(pred), 1e-8)
102    return {"predicted_max_rate": pred, "observed_full_rate": observed,
103            "observed_quotient_rate": aq, "observed_fibre_rate": af,
104            "relative_error": rel, "confirmed": bool(np.isfinite(rel) and rel <= .20),
105            "measurement": "trained quotient-fibre GRU hidden-state perturbation decay"}
106
107
108def main():
109    # Baseline sweep on four seeds, then full eight-seed reevaluation by harness.
110    base = sweep_baseline(lambda cfg: (lambda s: train_one("baseline", s, cfg["lr"])), GRID)
111    best_lr = float(base["best_cfg"]["lr"])
112    idea_cfgs = [best_lr] + [float(g["lr"]) for g in GRID if float(g["lr"]) != best_lr]
113    idea_runs = []
114    for lr in idea_cfgs:
115        r = evaluate(lambda s, lr=lr: train_one("idea", s, lr), seeds=SEEDS)
116        idea_runs.append({"lr": lr, "result": r})
117    idea_best = min(idea_runs, key=lambda q: q["result"]["mean"])
118    report = make_report("dynamics", "rnn_small", base, idea_best["result"], {
119        "predicted_vs_observed": mechanism_signature(idea_best["lr"]),
120        "idea_lr_sweep": idea_runs,
121        "parameter_counts": {"baseline": count_params(make_model("rnn_small", (24,), 1)), "idea": count_params(QuotientFibreGRU(1))},
122        "track_choice": "dynamics structurally matches stability/control and multi-step dynamical memory"
123    })
124    Path("bench_report.json").write_text(json.dumps(report, indent=2))
125    print(json.dumps(report, indent=2))
126
127if __name__ == "__main__":
128    main()