1import os, sys, json
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import train_model, evaluate, sweep_baseline, make_report, count_params
  8from custom_track import get_dataset, META
  9
 10
 11def seed_all(seed):
 12    np.random.seed(seed)
 13    torch.manual_seed(seed)
 14    if torch.cuda.is_available():
 15        try:
 16            torch.cuda.manual_seed_all(seed)
 17        except Exception:
 18            pass
 19
 20
 21class Dense3D(nn.Module):
 22    def __init__(self):
 23        super().__init__()
 24        self.op = nn.Linear(64, 64, bias=False)
 25        self.head = nn.Sequential(nn.ReLU(), nn.Linear(64, 1))
 26        nn.init.xavier_uniform_(self.op.weight)
 27
 28    def forward(self, x):
 29        return self.head(self.op(x.reshape(x.shape[0], 64)))
 30
 31
 32class TT3D(nn.Module):
 33    """Learned TT-matrix over the three 4-way spatial modes."""
 34    def __init__(self, rank=2):
 35        super().__init__()
 36        self.rank = rank
 37        self.cores = nn.ParameterList([
 38            nn.Parameter(torch.randn(1, 4, 4, rank) * 0.16),
 39            nn.Parameter(torch.randn(rank, 4, 4, rank) * 0.16),
 40            nn.Parameter(torch.randn(rank, 4, 4, 1) * 0.16),
 41        ])
 42        self.head = nn.Sequential(nn.ReLU(), nn.Linear(64, 1))
 43
 44    def forward(self, x):
 45        g1, g2, g3 = self.cores
 46        # x[b,i,j,k], g1[a,l,i,c], g2[c,n,j,d], g3[d,o,k,f].
 47        y = torch.einsum("bijk,alic,cnjd,dokf->blno", x, g1, g2, g3)
 48        return self.head(y.reshape(x.shape[0], 64))
 49
 50
 51def train_one(kind, seed, lr, epochs, rank=2):
 52    seed_all(seed)
 53    d0 = get_dataset(seed, 400, 100)
 54    d = {**d0, **{k: torch.as_tensor(d0[k], dtype=torch.float32) for k in ("xtr", "ytr", "xte", "yte")}}
 55    model = Dense3D() if kind == "dense" else TT3D(rank)
 56    net, metric, _ = train_model(model, d, epochs=epochs, lr=lr, batch=128,
 57                                 weight_decay=0.0, log=lambda *_: None)
 58    if net is None:
 59        return float("nan"), None
 60    return float(metric), net
 61
 62
 63def main():
 64    # The union of step sizes is identical for baseline and idea.
 65    grid = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 66    epochs = 18
 67
 68    def baseline_factory(cfg):
 69        return lambda s: train_one("dense", int(s), cfg["lr"], epochs)[0]
 70
 71    base = sweep_baseline(baseline_factory, grid)
 72    idea_trials = []
 73    for cfg in grid:
 74        r = evaluate(lambda s, c=cfg: train_one("tt", int(s), c["lr"], epochs, rank=2)[0])
 75        idea_trials.append({"cfg": cfg, **r})
 76    best_idea = min(idea_trials, key=lambda r: r["mean"])
 77
 78    # Signature is measured on trained models and test inputs.
 79    best_lr = base["best_cfg"]["lr"]
 80    _, bnet = train_one("dense", 0, best_lr, epochs)
 81    _, inet = train_one("tt", 0, best_idea["cfg"]["lr"], epochs, rank=2)
 82    d = get_dataset(0, 400, 100)
 83    xb = d["xte"]
 84    with torch.no_grad():
 85        dense_out = bnet.op(xb.reshape(-1, 64)).numpy()
 86        g1, g2, g3 = [g.detach() for g in inet.cores]
 87        tt_out = torch.einsum("bijk,alic,cnjd,dokf->blno", xb, g1, g2, g3)
 88        tt_out = tt_out.reshape(-1, 64).numpy()
 89    sig = {
 90        "quantity": "trained operator output norm ratio on held-out PDE fields",
 91        "predicted": "TT has fewer parameters than the dense spatial map",
 92        "observed": {
 93            "dense_params": count_params(bnet),
 94            "tt_params": count_params(inet),
 95            "operator_output_norm_ratio": float(np.linalg.norm(tt_out) / (np.linalg.norm(dense_out) + 1e-12)),
 96        },
 97        "confirmed": bool(count_params(inet) < count_params(bnet)),
 98    }
 99    idea = {"best_cfg": best_idea["cfg"], "trials": idea_trials,
100            "mean": best_idea["mean"], "std": best_idea["std"],
101            "per_seed": best_idea["per_seed"], "n": best_idea["n"]}
102    rep = make_report("separable_diffusion_field", "custom_dense_vs_tt3d", base, idea, sig)
103    rep["custom_track"] = {"name": META["name"], "file": "custom_track.py", "domain": META["domain"]}
104    rep["parameter_counts"] = {"baseline": count_params(bnet), "idea": count_params(inet)}
105    with open("bench_report.json", "w") as f:
106        json.dump(rep, f, indent=2)
107    print(json.dumps(rep, indent=2))
108
109
110if __name__ == "__main__":
111    main()