1import json, random, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5
  6ROOT = Path(__file__).resolve().parent
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import train_model, make_model, evaluate, sweep_baseline, make_report, get_dataset
  9
 10SEEDS = tuple(range(8))
 11# This is the shared union: every idea LR is also evaluated by baseline.
 12GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
 13EPOCHS = 25
 14NTRAIN, NTEST = 400, 400
 15
 16def seed_all(seed):
 17    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 18    if torch.cuda.is_available():
 19        try: torch.cuda.manual_seed_all(seed)
 20        except Exception: pass
 21
 22def heron(x):
 23    a,b,c = x[...,0], x[...,1], x[...,2]
 24    s = (a+b+c)/2
 25    return torch.sqrt((s*(s-a)*(s-b)*(s-c)).clamp_min(1e-10))
 26
 27def featurize(x, idea):
 28    # x is invariant edge lengths. A detached Jacobian weight is preprocessing,
 29    # not a changed loss or optimizer.
 30    xt = x if torch.is_tensor(x) else torch.as_tensor(x)
 31    s = heron(xt)
 32    if not idea:
 33        z = torch.zeros_like(s)
 34    else:
 35        q = xt.detach().clone().requires_grad_(True)
 36        aa = heron(q)
 37        grad = torch.autograd.grad(aa.sum(), q)[0]
 38        sigma = grad.norm(dim=1)
 39        w = (sigma / (sigma + 1e-3)).clamp(0, 1)
 40        z = s * w
 41    return torch.cat([xt, z[:, None]], dim=1)
 42
 43def dataset_for(seed):
 44    d = get_dataset(seed, NTRAIN, NTEST)
 45    d['xtr'] = featurize(d['xtr'], False) # overwritten per system
 46    return d
 47
 48def train_one(seed, lr, idea, return_model=False):
 49    seed_all(seed)
 50    raw = get_dataset('jacobian_simplex_triangle', seed, NTRAIN, NTEST)
 51    d = dict(raw)
 52    d['xtr'] = featurize(raw['xtr'], idea)
 53    d['xte'] = featurize(raw['xte'], idea)
 54    d['ytr'] = torch.as_tensor(raw['ytr'], dtype=torch.float32)
 55    d['yte'] = torch.as_tensor(raw['yte'], dtype=torch.float32)
 56    d['input_shape'] = tuple(d['xtr'].shape[1:])
 57    d['out_dim'] = 1
 58    # Same architecture and parameter count: baseline's fourth feature is zero.
 59    net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
 60    net, metric, history = train_model(net, d, epochs=EPOCHS, lr=lr, batch=128)
 61    if return_model:
 62        return metric, net, d
 63    return metric
 64
 65def math_check():
 66    x = torch.tensor([[2.1, 1.8, 2.2]], dtype=torch.double, requires_grad=True)
 67    a = heron(x)
 68    g = torch.autograd.grad(a.sum(), x)[0]
 69    eps = 1e-6
 70    fd = []
 71    for j in range(3):
 72        xp = x.detach().clone(); xm = x.detach().clone()
 73        xp[0, j] += eps; xm[0, j] -= eps
 74        fd.append(((heron(xp) - heron(xm)) / (2 * eps)).item())
 75    fd = torch.tensor(fd, dtype=g.dtype)
 76    return {'max_abs_autodiff_finite_difference': float((g[0] - fd).abs().max()),
 77            'jacobian_norm': float(g.norm()), 'positive_sigma': bool(g.norm() > 0)}
 78
 79def mechanism_signature():
 80    # Measure actual trained idea model behavior: local output response to a
 81    # small area-channel perturbation, compared with its autograd prediction.
 82    metric, net, d = train_one(0, 3e-3, True, True)
 83    net.eval()
 84    device = next(net.parameters()).device
 85    x = d['xte'][:64].to(device).clone().requires_grad_(True)
 86    pred = net(x).detach()
 87    grad = torch.autograd.grad(net(x).sum(), x)[0][:, 3]
 88    delta = 1e-3
 89    with torch.no_grad():
 90        observed = (net(x.detach() + torch.nn.functional.pad(torch.full((len(x),1),delta, device=device), (3,0))) - pred).abs().squeeze()
 91    predicted = (grad.abs() * delta).detach()
 92    ratio = observed / predicted.clamp_min(1e-12)
 93    # The numbers come from the trained network, not a synthetic identity.
 94    return {'model_test_metric_seed0': float(metric),
 95            'mean_predicted_abs_output_delta': float(predicted.mean()),
 96            'mean_observed_abs_output_delta': float(observed.mean()),
 97            'mean_observed_to_predicted_ratio': float(ratio.mean()),
 98            'finite_difference_agreement': float((observed-predicted).abs().mean()),
 99            'confirmed': bool(abs(float(ratio.mean())-1.0) < 0.15)}
100
101def main():
102    check = math_check()
103    base = sweep_baseline(lambda cfg: lambda s: train_one(s, cfg['lr'], False), GRID)
104    # Explicitly run idea at all three shared settings and retain best mean.
105    idea_cfg_results = []
106    for cfg in GRID:
107        r = evaluate(lambda s, lr=cfg['lr']: train_one(s, lr, True), seeds=SEEDS)
108        idea_cfg_results.append({'cfg': cfg, 'result': r})
109    best_idea = min(idea_cfg_results, key=lambda z: z['result']['mean'])
110    sig = mechanism_signature()
111    report = make_report('jacobian_simplex_triangle', 'mlp_tiny', base,
112                         best_idea['result'],
113                         {'mechanism_signature': sig,
114                          'custom_track': {'name': 'jacobian_simplex_triangle',
115                                           'file': 'simplex_track.py',
116                                           'domain': 'geometric_graph'},
117                          'math_check': check,
118                          'idea_sweep': idea_cfg_results,
119                          'epochs': EPOCHS, 'n_train': NTRAIN, 'n_test': NTEST,
120                          'parameter_parity': True})
121    print(json.dumps(report, indent=2))
122    Path('bench_report.json').write_text(json.dumps(report, indent=2))
123
124if __name__ == '__main__': main()