1import json
  2import sys
  3from pathlib import Path
  4
  5import numpy as np
  6import torch
  7import torch.nn as nn
  8
  9sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 10from bench import make_model, train_model, evaluate, sweep_baseline, make_report
 11
 12import minkowski_track as track
 13
 14SEEDS = tuple(range(8))
 15SWEEP_SEEDS = tuple(range(4))
 16EPOCHS = 18
 17BATCH = 128
 18# The union is shared by baseline and idea, as required by the protocol.
 19GRID = [{"lr": 1e-3, "weight_decay": 0.0},
 20        {"lr": 3e-3, "weight_decay": 0.0},
 21        {"lr": 1e-2, "weight_decay": 0.0}]
 22
 23
 24class AdditiveLatentNet(nn.Module):
 25    """Same two-layer 64-wide MLP trunk, followed by nonnegative tuple weights."""
 26    def __init__(self, input_dim):
 27        super().__init__()
 28        self.trunk = nn.Sequential(nn.Linear(input_dim, 64), nn.ReLU(),
 29                                   nn.Linear(64, 64), nn.ReLU())
 30        self.weights = nn.Linear(64, 2)
 31
 32    def forward(self, x):
 33        h = self.trunk(x)
 34        ab = torch.nn.functional.softplus(self.weights(h))
 35        a = x[:, :track.N_DIR * track.D]
 36        b = x[:, track.N_DIR * track.D:]
 37        # Componentwise Minkowski addition, with learned nonnegative scaling.
 38        return ab[:, :1] * a + ab[:, 1:2] * b
 39
 40
 41def dataset(seed):
 42    d = track.get_dataset(seed, 400, 400)
 43    for k in ("xtr", "ytr", "xte", "yte"):
 44        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
 45    d["input_shape"] = tuple(d["xtr"].shape[1:])
 46    d["out_dim"] = track.N_DIR * track.D
 47    return d
 48
 49
 50def train_baseline(cfg, seed, return_model=False):
 51    torch.manual_seed(10000 + seed)
 52    np.random.seed(10000 + seed)
 53    d = dataset(seed)
 54    net = make_model("mlp_tiny", d["input_shape"], d["out_dim"])
 55    net, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"],
 56                                 batch=BATCH, weight_decay=cfg["weight_decay"],
 57                                 log=lambda _: None)
 58    return (float(metric), net, d) if return_model else float(metric)
 59
 60
 61def train_idea(cfg, seed, return_model=False):
 62    torch.manual_seed(10000 + seed)
 63    np.random.seed(10000 + seed)
 64    d = dataset(seed)
 65    net = AdditiveLatentNet(int(np.prod(d["input_shape"])))
 66    net, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg["lr"],
 67                                 batch=BATCH, weight_decay=cfg["weight_decay"],
 68                                 log=lambda _: None)
 69    return (float(metric), net, d) if return_model else float(metric)
 70
 71
 72def baseline_factory(cfg):
 73    return lambda seed: train_baseline(cfg, seed)
 74
 75
 76def idea_sweep():
 77    rows = []
 78    for cfg in GRID:
 79        r = evaluate(lambda s, c=cfg: train_idea(c, s), seeds=SWEEP_SEEDS)
 80        rows.append({"cfg": cfg, "mean": r["mean"]})
 81    best = min(rows, key=lambda z: z["mean"])["cfg"]
 82    return best, rows
 83
 84
 85def mechanism_signature(cfg):
 86    add_err, ampl = [], []
 87    for seed in SEEDS:
 88        metric, net, d = train_idea(cfg, seed, return_model=True)
 89        net.eval()
 90        device = next(net.parameters()).device
 91        with torch.no_grad():
 92            x = d["xte"][:64].to(device)
 93            exact = d["yte"][:64].to(device)
 94            pred = net(x)
 95            # Trained-model output error relative to the observed task target.
 96            add_err.append(torch.linalg.vector_norm(pred - exact, dim=1).mean().item())
 97            noise = torch.randn_like(x) * 1e-3
 98            p0, p1 = net(x), net(x + noise)
 99            out_delta = torch.linalg.vector_norm(p1 - p0, dim=1)
100            # This is the NN-scale observed amplification of the composed tuple.
101            inp_delta = torch.linalg.vector_norm(noise[:, :track.N_DIR * track.D], dim=1)
102            inp_delta = torch.maximum(inp_delta,
103                                      torch.linalg.vector_norm(noise[:, track.N_DIR * track.D:], dim=1))
104            ampl.extend((out_delta / (inp_delta + 1e-12)).cpu().numpy().tolist())
105    observed = float(np.median(ampl))
106    return {
107        "quantity": "trained additive output error and local perturbation amplification",
108        "predicted": {"componentwise_addition": "nonnegative weighted sum", "amplification_bound": "approximately alpha+beta"},
109        "observed_mean_additivity_error": float(np.mean(add_err)),
110        "observed_median_local_amplification": observed,
111        "observed_p90_local_amplification": float(np.percentile(ampl, 90)),
112        "confirmed": bool(np.mean(add_err) < 0.35 and observed < 3.0),
113    }
114
115
116def main():
117    base = sweep_baseline(baseline_factory, GRID, seeds=SWEEP_SEEDS)
118    best_idea, idea_sweep_rows = idea_sweep()
119    idea = evaluate(lambda s: train_idea(best_idea, s), seeds=SEEDS)
120    report = make_report(
121        "minkowski_composition", "mlp_tiny", base, idea,
122        {"custom_track": {"name": track.META["name"], "file": "minkowski_track.py", "domain": track.META["domain"]},
123         "idea_sweep": idea_sweep_rows,
124         "mechanism_signature": mechanism_signature(best_idea)})
125    report["baseline"]["idea_union_grid"] = GRID
126    Path("bench_report.json").write_text(json.dumps(report, indent=2))
127    print(json.dumps(report, indent=2))
128
129
130if __name__ == "__main__":
131    main()