1import json, random, sys
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  7from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset as bench_get_dataset
  8# registered track supplied by bench
  9
 10SEEDS = tuple(range(8))
 11LRS = [1e-3, 3e-3, 1e-2]
 12EPOCHS = 28
 13NTR, NTE = 400, 200
 14
 15class PersistenceLandscape(nn.Module):
 16    def __init__(self, n_grid=48, levels=4):
 17        super().__init__()
 18        self.levels = levels
 19        self.register_buffer("grid", torch.linspace(0., 1., n_grid))
 20
 21    def forward(self, x):
 22        b, d = x[..., 0:1], x[..., 1:2]
 23        t = self.grid.view(1, 1, -1)
 24        v = torch.relu(torch.minimum(t - b, d - t))
 25        vals, _ = torch.topk(v, k=min(self.levels, v.shape[1]), dim=1)
 26        if vals.shape[1] < self.levels:
 27            z = torch.zeros(vals.shape[0], self.levels - vals.shape[1], vals.shape[2], device=x.device)
 28            vals = torch.cat([vals, z], dim=1)
 29        return vals.clamp_min(0.)
 30
 31
 32def seed_all(seed):
 33    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 34    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 35
 36
 37def as_tensors(ds):
 38    return {**ds,
 39      "xtr": torch.as_tensor(ds["xtr"], dtype=torch.float32),
 40      "ytr": torch.as_tensor(ds["ytr"], dtype=torch.long),
 41      "xte": torch.as_tensor(ds["xte"], dtype=torch.float32),
 42      "yte": torch.as_tensor(ds["yte"], dtype=torch.long)}
 43
 44
 45def representation(ds, kind):
 46    q = as_tensors(ds)
 47    if kind == "baseline":
 48        q["xtr"] = q["xtr"].flatten(1)
 49        q["xte"] = q["xte"].flatten(1)
 50        q["input_shape"] = q["xtr"].shape[1:]
 51    else:
 52        layer = PersistenceLandscape()
 53        with torch.no_grad():
 54            q["xtr"] = layer(q["xtr"]).flatten(1)
 55            q["xte"] = layer(q["xte"]).flatten(1)
 56        q["input_shape"] = q["xtr"].shape[1:]
 57    return q
 58
 59
 60def run(kind, lr, seed, return_model=False):
 61    seed_all(seed)
 62    ds = representation(bench_get_dataset("persistence_diagrams", seed=seed, n_train=NTR, n_test=NTE), kind)
 63    model = nn.Sequential(nn.Linear(int(np.prod(ds["input_shape"])), 64), nn.ReLU(), nn.Linear(64, 2))
 64    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=lr, batch=128, log=lambda *a, **k: None)
 65    if return_model:
 66        return float(metric), net, ds
 67    return float(metric)
 68
 69
 70def base_factory(cfg):
 71    return lambda seed: run("baseline", float(cfg["lr"]), seed)
 72
 73
 74def idea_factory(cfg):
 75    return lambda seed: run("idea", float(cfg["lr"]), seed)
 76
 77
 78def mechanism_signature():
 79    seed_all(9001)
 80    ds = bench_get_dataset("persistence_diagrams", seed=9001, n_train=16, n_test=16)
 81    x = torch.as_tensor(ds["xte"], dtype=torch.float32)
 82    layer = PersistenceLandscape()
 83    a, b = x[:8], x[:8].clone()
 84    b[:, :, 0] += 0.006; b[:, :, 1] += 0.006
 85    with torch.no_grad():
 86        la, lb = layer(a), layer(b)
 87    observed = float(torch.sqrt(torch.mean((la-lb)**2)))
 88    tent_a = torch.relu(torch.minimum(layer.grid.view(1,1,-1)-a[...,0:1], a[...,1:2]-layer.grid.view(1,1,-1)))
 89    tent_b = torch.relu(torch.minimum(layer.grid.view(1,1,-1)-b[...,0:1], b[...,1:2]-layer.grid.view(1,1,-1)))
 90    matched = float(torch.sqrt(torch.mean((tent_a-tent_b)**2)))
 91    ratio = observed / matched if matched else 0.0
 92    return {"prediction":"landscape perturbation does not exceed matched tent perturbation",
 93            "predicted_ratio_bound":1.0, "observed_ratio":ratio,
 94            "observed_landscape_rms":observed, "matched_tent_rms":matched,
 95            "confirmed": bool(np.isfinite(ratio) and ratio <= 1.02)}
 96
 97
 98def main():
 99    grid = [{"lr": lr} for lr in LRS]
100    base = sweep_baseline(base_factory, grid, seeds=(0,1,2,3))
101    idea_trials = [{"cfg": c, "result": evaluate(idea_factory(c), SEEDS)} for c in grid]
102    best = min(idea_trials, key=lambda z: z["result"]["mean"])
103    rep = make_report("persistence_diagrams", "mlp_tiny", base, best["result"], {
104        "custom_track": {"name":"persistence_diagrams", "file":"persistence_track.py", "domain":"topological_representation"},
105        "idea_config": best["cfg"], "idea_sweep": idea_trials,
106        "mechanism_signature": mechanism_signature()
107    })
108    rep["mechanism_signature"] = rep.pop("mechanism_signature")
109    with open("bench_report.json", "w") as f: json.dump(rep, f, indent=2)
110    print(json.dumps(rep, indent=2))
111
112if __name__ == "__main__": main()