1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6import torch.nn.functional as F
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_report, evaluate, sweep_baseline
  9
 10TRACK = 'geometry_support_regression'
 11SEEDS = tuple(range(8))
 12GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]
 13R2 = 7.5
 14
 15class MatchedMLP(nn.Module):
 16    def __init__(self, seed, idea=False):
 17        super().__init__()
 18        torch.manual_seed(seed)
 19        self.idea = idea
 20        self.features = nn.Sequential(nn.Linear(2, 64), nn.ReLU(), nn.Linear(64, 5), nn.Tanh())
 21        self.head = nn.Linear(5, 1)
 22    def encode(self, x):
 23        return self.features(x)
 24    def forward(self, x):
 25        return self.head(self.encode(x))
 26
 27def compact_score(z):
 28    d = z.shape[1]
 29    a = (R2 - d - 2.0) / 2.0
 30    mu = z.detach().mean(0)
 31    q = z.detach() - mu
 32    cov = q.T @ q / max(1, z.shape[0] - 1) + 1e-3 * torch.eye(d, device=z.device)
 33    L = torch.linalg.cholesky(cov)
 34    v = torch.linalg.solve_triangular(L, (z - mu).T, upper=False).T
 35    r2 = (v * v).sum(1)
 36    logc = (torch.lgamma(torch.tensor(d / 2 + 1 + a, device=z.device))
 37            - d / 2 * math.log(math.pi) - d / 2 * math.log(R2)
 38            - torch.lgamma(torch.tensor(1 + a, device=z.device)))
 39    score = -logc - a * torch.log(torch.clamp(1 - r2 / R2, min=1e-6)) + F.softplus(r2 - R2)
 40    return score, r2
 41
 42def baseline_train(cfg, seed):
 43    np.random.seed(seed); torch.manual_seed(seed)
 44    ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
 45    net = MatchedMLP(seed, idea=False)
 46    net, metric, _ = __import__('bench').train_model(net, ds, epochs=25, lr=cfg['lr'], batch=128, log=lambda *_: None)
 47    return float(metric)
 48
 49def idea_train(cfg, seed, capture=False):
 50    np.random.seed(seed); torch.manual_seed(seed)
 51    ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
 52    try:
 53        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 54    except Exception:
 55        device = torch.device('cpu')
 56    try:
 57        net = MatchedMLP(seed, idea=True).to(device)
 58        x = ds['xtr'].to(device); y = ds['ytr'].to(device)
 59        xt = ds['xte'].to(device); yt = ds['yte'].to(device)
 60        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
 61        for _ in range(25):
 62            net.train(); opt.zero_grad()
 63            z = net.encode(x); out = net.head(z)
 64            loss = F.mse_loss(out, y) + 0.002 * compact_score(z)[0].mean()
 65            loss.backward(); opt.step()
 66        net.eval()
 67        with torch.no_grad():
 68            zte = net.encode(xt); pred = net.head(zte); metric = F.mse_loss(pred, yt).item()
 69            _, r2 = compact_score(zte)
 70            sig = {'mean_test_radius2': float(r2.mean().cpu()),
 71                   'test_inside_support_fraction': float((r2 < R2).float().mean().cpu())}
 72        return (float(metric), sig) if capture else float(metric)
 73    except RuntimeError:
 74        torch.backends.cudnn.enabled = False
 75        net = MatchedMLP(seed, idea=True)
 76        x, y, xt, yt = ds['xtr'], ds['ytr'], ds['xte'], ds['yte']
 77        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
 78        for _ in range(25):
 79            opt.zero_grad(); z = net.encode(x); out = net.head(z)
 80            (F.mse_loss(out, y) + 0.002 * compact_score(z)[0].mean()).backward(); opt.step()
 81        with torch.no_grad():
 82            zte = net.encode(xt); pred = net.head(zte); _, r2 = compact_score(zte)
 83        metric = F.mse_loss(pred, yt).item()
 84        sig = {'mean_test_radius2': float(r2.mean()), 'test_inside_support_fraction': float((r2 < R2).float().mean())}
 85        return (float(metric), sig) if capture else float(metric)
 86
 87def main():
 88    base = sweep_baseline(lambda cfg: (lambda seed: baseline_train(cfg, seed)), GRID, seeds=(0,1,2,3))
 89    best = base['best_cfg']
 90    idea = evaluate(lambda seed: idea_train(best, seed), seeds=SEEDS)
 91    sigrows = [idea_train(best, s, capture=True)[1] for s in SEEDS]
 92    observed = float(np.mean([r['mean_test_radius2'] for r in sigrows]))
 93    extra = {'prediction': 'calibrated compact density predicts nominal E[r2]=d=5',
 94             'predicted_mean_radius2': 5.0,
 95             'observed_mean_radius2_on_trained_models': observed,
 96             'absolute_error': abs(observed - 5.0),
 97             'mean_inside_support_fraction': float(np.mean([r['test_inside_support_fraction'] for r in sigrows])),
 98             'confirmed': bool(abs(observed - 5.0) < 0.5)}
 99    report = make_report(TRACK, 'matched_mlp_embedding_regressor', base, idea, extra)
100    report['idea']['per_seed_signatures'] = sigrows
101    Path('bench_report.json').write_text(json.dumps(report, indent=2))
102    print(json.dumps(report, indent=2))
103
104if __name__ == '__main__': main()