1import json
  2import math
  3import sys
  4from pathlib import Path
  5
  6import numpy as np
  7import torch
  8from scipy.special import roots_genlaguerre
  9
 10sys.path.insert(0, "/home/maxwelhelp/all/math2nn")
 11from bench import get_dataset, train_model, sweep_baseline, make_report, evaluate
 12
 13SEED = 2977
 14N_FEATURES = 64
 15EPOCHS = 12
 16NTR, NTE = 1200, 300
 17LR_GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 18
 19
 20def quadrature_frequencies(d=10, n_features=N_FEATURES):
 21    """Product quadrature: generalized-Laguerre radial rule x fixed sphere rule.
 22    If r^2/2 ~ Gamma(d/2,1), Laguerre integrates its radial density; directions
 23    are a fixed antipodal quasi-uniform set, giving a deterministic Gaussian
 24    spectral rule with nonnegative equal weights.
 25    """
 26    ndir, nq = 16, n_features // 16
 27    t, wt = roots_genlaguerre(nq, d / 2.0 - 1.0)
 28    wt = wt / math.gamma(d / 2.0)
 29    rng = np.random.RandomState(18431)
 30    u = rng.normal(size=(ndir // 2, d))
 31    u /= np.linalg.norm(u, axis=1, keepdims=True)
 32    dirs = np.concatenate([u, -u], axis=0)
 33    # Pairing removes odd angular bias and is deterministic once constructed.
 34    W = np.concatenate([np.sqrt(2.0 * r) * dirs[j:j+1]
 35                        for r in t for j in range(ndir)], axis=0)
 36    a = np.repeat(wt / ndir, ndir)
 37    return W.astype(np.float32), a.astype(np.float32)
 38
 39
 40def random_frequencies(d=10, n_features=N_FEATURES, seed=0):
 41    return np.random.RandomState(seed).normal(size=(n_features, d)).astype(np.float32), \
 42           np.full(n_features, 1.0 / n_features, dtype=np.float32)
 43
 44
 45class FourierMLP(torch.nn.Module):
 46    def __init__(self, W, weights):
 47        super().__init__()
 48        self.register_buffer("W", torch.as_tensor(W))
 49        self.register_buffer("sqrt_a", torch.sqrt(torch.as_tensor(weights)))
 50        self.net = torch.nn.Sequential(
 51            torch.nn.Linear(2 * len(W), 64), torch.nn.ReLU(),
 52            torch.nn.Linear(64, 64), torch.nn.ReLU(), torch.nn.Linear(64, 1))
 53
 54    def embedding(self, x):
 55        z = x @ self.W.T
 56        f = torch.cat((torch.cos(z), torch.sin(z)), dim=1)
 57        return f * self.sqrt_a.repeat(2).unsqueeze(0)
 58
 59    def forward(self, x):
 60        return self.net(self.embedding(x))
 61
 62
 63def prep(seed):
 64    d = get_dataset("tabular", seed, n_train=NTR, n_test=NTE)
 65    # Standardization is shared preprocessing and prevents frequency scale from
 66    # being dominated by a coordinate's units.
 67    mu, sd = d["xtr"].mean(0, keepdim=True), d["xtr"].std(0, keepdim=True).clamp_min(1e-5)
 68    for k in ("xtr", "xte"):
 69        d[k] = (d[k] - mu) / sd
 70    return d
 71
 72
 73def train_value(kind, cfg, seed, return_model=False):
 74    torch.manual_seed(SEED + 100 * seed + (0 if kind == "baseline" else 10000))
 75    d = prep(seed)
 76    if kind == "baseline":
 77        W, a = random_frequencies(seed=SEED + seed)
 78    else:
 79        W, a = quadrature_frequencies()
 80    model, metric, _ = train_model(FourierMLP(W, a), d, epochs=EPOCHS,
 81                                   lr=cfg["lr"], batch=128, weight_decay=0.0,
 82                                   log=lambda *_: None)
 83    if return_model:
 84        return metric, model, d, W, a
 85    return metric
 86
 87
 88def make_train(kind):
 89    return lambda cfg: (lambda seed: train_value(kind, cfg, seed))
 90
 91
 92def spectral_error(model, d, W, a):
 93    with torch.no_grad():
 94        x = d["xte"][:120]
 95        z = x @ torch.as_tensor(W).T
 96        F = torch.cat((torch.cos(z), torch.sin(z)), 1) * torch.sqrt(torch.as_tensor(a)).repeat(2).unsqueeze(0)
 97        Kh = F @ F.T
 98        dist = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
 99        K = torch.exp(-0.5 * dist)
100        le = torch.linalg.eigvalsh(K).flip(0)[:10]
101        lh = torch.linalg.eigvalsh(Kh).flip(0)[:10]
102        return float(torch.mean(torch.abs(le-lh) / le.clamp_min(1e-7)))
103
104
105def main():
106    # Baseline sweep and full evaluation are performed by the canonical protocol.
107    base = sweep_baseline(make_train("baseline"), LR_GRID)
108    # Same union of learning rates on idea side; select its best on the same
109    # four tuning seeds, then evaluate that setting on all eight paired seeds.
110    idea_trials = []
111    for cfg in LR_GRID:
112        vals = [train_value("idea", cfg, s) for s in range(4)]
113        idea_trials.append({"cfg": cfg, "mean": float(np.mean(vals))})
114    best_cfg = min(idea_trials, key=lambda x: x["mean"])["cfg"]
115    idea = {"best_cfg": best_cfg, "sweep": idea_trials,
116            "full": __import__("bench").protocol.evaluate(make_train("idea")(best_cfg))}
117
118    # Re-test the claimed spectral effect on the actual trained systems.
119    _, bmodel, bd, bW, ba = train_value("baseline", best_cfg, 0, True)
120    _, imodel, idd, iW, ia = train_value("idea", best_cfg, 0, True)
121    bspec = spectral_error(bmodel, bd, bW, ba)
122    ispec = spectral_error(imodel, idd, iW, ia)
123    sig = {"prediction": "deterministic quadrature has lower top-10 RBF Gram eigenvalue error",
124           "predicted_baseline_top10_relative_error": bspec,
125           "predicted_idea_top10_relative_error": ispec,
126           "observed_delta_idea_minus_baseline": ispec - bspec,
127           "confirmed": bool(ispec < bspec)}
128    report = make_report("tabular", "fourier_mlp_shared", base, idea["full"],
129                         {"mechanism_signature": sig,
130                          "notes": "Tabular is the matched built-in track for an input embedding/kernel feature intervention."})
131    Path("bench_report.json").write_text(json.dumps(report, indent=2))
132    print(json.dumps(report, indent=2))
133
134
135if __name__ == "__main__":
136    main()