1import json
  2import random
  3import sys
  4import numpy as np
  5import torch
  6from torch import nn
  7
  8sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
  9from bench import train_model, evaluate, sweep_baseline, make_report
 10from tree_polynomial_track import get_dataset, META
 11
 12SEEDS = tuple(range(8))
 13# The union is used for both methods, satisfying learning-rate parity.
 14LRS = [1e-3, 3e-3, 1e-2]
 15EPOCHS = 24
 16NTR, NTE = 400, 200
 17YS = np.array([-3., -2., -1., -.5, .5, 1., 2., 3.], dtype=np.float64)
 18
 19
 20def seed_all(seed):
 21    random.seed(seed)
 22    np.random.seed(seed)
 23    torch.manual_seed(seed)
 24    if torch.cuda.is_available():
 25        try:
 26            torch.cuda.manual_seed_all(seed)
 27        except Exception:
 28            pass
 29
 30
 31def polynomial_values(tree_x):
 32    """Evaluate P(root; 1,y) for each encoded fixed-size rooted tree.
 33
 34    Parent ids are stored as parent_id/5 and edge lengths in column 1.
 35    This is the exact product recursion, evaluated independently at eight y's.
 36    """
 37    n = len(tree_x)
 38    out = np.ones((n, len(YS)), dtype=np.float64)
 39    for i in range(n):
 40        parent = np.rint(tree_x[i, :, 0] * 5).astype(int)
 41        z = np.ones((6, len(YS)), dtype=np.float64)
 42        # Every non-root node is a leaf or has already been evaluated when used.
 43        for v in (5, 4, 3, 2, 1):
 44            children = np.flatnonzero(parent == v)
 45            if len(children) == 0:
 46                z[v] = 1.0
 47            else:
 48                z[v] = np.prod(tree_x[i, children, 1, None] * YS[None, :] + z[children], axis=0)
 49        children = np.flatnonzero(parent == 0)
 50        z[0] = 1.0 if len(children) == 0 else np.prod(
 51            tree_x[i, children, 1, None] * YS[None, :] + z[children], axis=0)
 52        out[i] = z[0]
 53    return out
 54
 55
 56def representation(ds, kind):
 57    q = dict(ds)
 58    if kind == "baseline":
 59        # Standard sum-style scalar message passing: root receives the sum of
 60        # edge lengths and terminal messages, which is blind to this topology.
 61        a = q["xtr"]
 62        b = q["xte"]
 63        q["xtr"] = a[:, :, 1].sum(1, keepdims=True)
 64        q["xte"] = b[:, :, 1].sum(1, keepdims=True)
 65    else:
 66        # Finite compressed branch-length polynomial fingerprint.
 67        q["xtr"] = np.log(np.abs(polynomial_values(q["xtr"])) + 1e-10).astype(np.float32)
 68        q["xte"] = np.log(np.abs(polynomial_values(q["xte"])) + 1e-10).astype(np.float32)
 69    q["xtr"] = torch.as_tensor(q["xtr"], dtype=torch.float32)
 70    q["xte"] = torch.as_tensor(q["xte"], dtype=torch.float32)
 71    q["ytr"] = torch.as_tensor(q["ytr"], dtype=torch.long)
 72    q["yte"] = torch.as_tensor(q["yte"], dtype=torch.long)
 73    q["input_shape"] = tuple(q["xtr"].shape[1:])
 74    q["out_dim"] = 2
 75    return q
 76
 77
 78def run(kind, lr, seed):
 79    seed_all(seed)
 80    ds = representation(get_dataset(seed, NTR, NTE), kind)
 81    model = nn.Sequential(nn.Flatten(), nn.Linear(int(np.prod(ds["input_shape"])), 32),
 82                          nn.ReLU(), nn.Linear(32, 2))
 83    _, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=float(lr), batch=128,
 84                               log=lambda *args, **kwargs: None)
 85    return float(metric)
 86
 87
 88def base_factory(cfg):
 89    return lambda seed: run("baseline", cfg["lr"], seed)
 90
 91
 92def idea_factory(cfg):
 93    return lambda seed: run("idea", cfg["lr"], seed)
 94
 95
 96def mechanism_signature():
 97    """Retest the claimed topology sensitivity through trained model outputs."""
 98    seed_all(9001)
 99    ds0 = get_dataset(9001, 200, 200)
100    db = representation(ds0, "baseline")
101    di = representation(ds0, "idea")
102    # Same branch-length multiset, one example of each topology.
103    xb = torch.cat([db["xte"][0:1], db["xte"][1:2]])
104    xi = torch.cat([di["xte"][0:1], di["xte"][1:2]])
105    def fit(ds):
106        model = nn.Sequential(nn.Flatten(), nn.Linear(int(np.prod(ds["input_shape"])), 32),
107                              nn.ReLU(), nn.Linear(32, 2))
108        net, _, _ = train_model(model, ds, epochs=EPOCHS, lr=3e-3, batch=128,
109                                log=lambda *a, **k: None)
110        return net
111    mb, mi = fit(db), fit(di)
112    devb = next(mb.parameters()).device
113    devi = next(mi.parameters()).device
114    with torch.no_grad():
115        pb = torch.softmax(mb(xb.to(devb)), 1).cpu()
116        pi = torch.softmax(mi(xi.to(devi)), 1).cpu()
117    observed_baseline = float(torch.linalg.vector_norm(pb[0] - pb[1]))
118    observed_idea = float(torch.linalg.vector_norm(pi[0] - pi[1]))
119    # Prediction: polynomial representation should yield larger trained-output
120    # separation than the sum representation on this structural pair.
121    return {"prediction": "trained fingerprint outputs separate topology pair more than sum baseline",
122            "observed_baseline_output_distance": observed_baseline,
123            "observed_idea_output_distance": observed_idea,
124            "predicted_ratio_lower_bound": 1.0,
125            "observed_ratio": observed_idea / max(observed_baseline, 1e-12),
126            "confirmed": bool(observed_idea > observed_baseline)}
127
128
129def main():
130    grid = [{"lr": lr} for lr in LRS]
131    base = sweep_baseline(base_factory, grid, seeds=(0, 1, 2, 3))
132    trials = [{"cfg": c, "result": evaluate(idea_factory(c), SEEDS)} for c in grid]
133    best = min(trials, key=lambda z: z["result"]["mean"])
134    report = make_report("rooted_tree_branch_classification", "mlp_tiny", base,
135                         best["result"], {
136                             "custom_track": {"name": META["name"], "file": "tree_polynomial_track.py", "domain": META["domain"]},
137                             "idea_config": best["cfg"], "idea_sweep": trials,
138                             "mechanism_signature": mechanism_signature()})
139    with open("bench_report.json", "w") as f:
140        json.dump(report, f, indent=2)
141    print(json.dumps(report, indent=2))
142
143
144if __name__ == "__main__":
145    main()