1from __future__ import annotations
  2import json, sys
  3from pathlib import Path
  4import numpy as np
  5import torch
  6from torch import nn
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8import bench
  9
 10SEEDS = tuple(range(8))
 11DEPTHS = (1, 2, 3, 4)
 12LRS = (0.0015, 0.003, 0.006)
 13EPOCHS = 12
 14BATCH = 128
 15WIDTH = 64
 16
 17class VariableMLP(nn.Module):
 18    def __init__(self, input_dim, depth, width=WIDTH):
 19        super().__init__()
 20        layers, d = [], input_dim
 21        for _ in range(depth):
 22            layers += [nn.Linear(d, width), nn.ReLU()]
 23            d = width
 24        layers.append(nn.Linear(d, 1))
 25        self.net = nn.Sequential(*layers)
 26    def forward(self, x):
 27        return self.net(x.reshape(x.shape[0], -1))
 28
 29def make_model(ds, depth):
 30    return VariableMLP(int(np.prod(tuple(ds['xtr'].shape[1:]))), int(depth))
 31
 32def train_metric(ds, depth, lr, return_model=False):
 33    model = make_model(ds, depth)
 34    net, metric, _ = bench.train_model(model, ds, epochs=EPOCHS, lr=float(lr),
 35                                       batch=BATCH, weight_decay=0.0, log=lambda *_: None)
 36    if metric is None:
 37        raise RuntimeError('bench training failed')
 38    return (float(metric), net) if return_model else float(metric)
 39
 40def eval_cfg(cfg, seeds=SEEDS):
 41    vals = []
 42    for seed in seeds:
 43        torch.manual_seed(10000 + seed)
 44        np.random.seed(20000 + seed)
 45        ds = bench.get_dataset('tabular', seed, n_train=400, n_test=400)
 46        vals.append(train_metric(ds, cfg['depth'], cfg['lr']))
 47    return {'per_seed': vals, 'mean': float(np.mean(vals)),
 48            'std': float(np.std(vals, ddof=1))}
 49
 50def baseline_train(cfg):
 51    return lambda seed: train_metric(bench.get_dataset('tabular', seed, 400, 400), cfg['depth'], cfg['lr'])
 52
 53def calibration(seed=9001):
 54    ds = bench.get_dataset('tabular', seed, n_train=400, n_test=400)
 55    # Calibration discrepancy is validation MSE to a fixed high-precision teacher.
 56    torch.manual_seed(30000 + seed)
 57    teacher, _ = train_metric(ds, 4, 0.003, return_model=True)
 58    del teacher
 59    rows = []
 60    for d in DEPTHS:
 61        torch.manual_seed(31000 + d)
 62        err = train_metric(ds, d, 0.003)
 63        rows.append((d, err))
 64    # Fit E(d)=floor+C_syn/d; residual margin gives conservative [L,U].
 65    X = np.array([[1.0, 1.0/d] for d, _ in rows])
 66    y = np.array([e for _, e in rows])
 67    coef = np.linalg.lstsq(X, y, rcond=None)[0]
 68    resid = y - X @ coef
 69    margin = max(0.002, 2.0 * float(np.std(resid)))
 70    return {'floor_L': float(coef[0]-margin), 'floor_U': float(coef[0]+margin),
 71            'C_syn': float(max(0.0, coef[1])), 'rows': rows,
 72            'residual_std': float(np.std(resid))}
 73
 74def main():
 75    # Search-space parity: every idea lr/depth pair is included in baseline grid.
 76    grid = [{'depth': d, 'lr': lr} for d in DEPTHS for lr in LRS]
 77    base = bench.sweep_baseline(baseline_train, grid, seeds=SEEDS[:4])
 78    base['full'] = eval_cfg(base['best_cfg'], SEEDS)
 79    cal = calibration()
 80    eps = float(base['full']['mean'])
 81    candidates = []
 82    for d in DEPTHS:
 83        radius = cal['C_syn'] / d + cal['floor_U'] - cal['floor_L']
 84        upper = cal['floor_U'] + radius
 85        lower = cal['floor_L'] - radius
 86        candidates.append((upper, d, lower))
 87    feasible = [x for x in candidates if x[0] <= eps]
 88    chosen_depth = min((x[1] for x in feasible), default=4)
 89    # Same-size idea sweep at the baseline best lr and two nearby union-grid lrs.
 90    idea_grid = [{'depth': chosen_depth, 'lr': lr} for lr in LRS]
 91    idea_trials = [{'cfg': c, 'result': eval_cfg(c, SEEDS[:4])} for c in idea_grid]
 92    best = min(idea_trials, key=lambda z: z['result']['mean'])['cfg']
 93    idea = eval_cfg(best, SEEDS)
 94    comparison = bench.compare_results(base['full'], idea)
 95    # NN-scale mechanism signature: observed fitted depth law from trained models.
 96    observed = []
 97    for d in DEPTHS:
 98        r = eval_cfg({'depth': d, 'lr': best['lr']}, SEEDS[:2])
 99        observed.append((d, r['mean']))
100    slope = float(np.polyfit(np.log([d for d,_ in observed]), np.log(np.maximum([e for _,e in observed],1e-12)), 1)[0])
101    signature = {'predicted_radius_depth_exponent': -1.0,
102                 'observed_test_error_depth_slope': slope,
103                 'predicted_vs_observed_tolerance': 0.35,
104                 'confirmed': bool(abs(slope + 1.0) <= 0.35),
105                 'observed_depth_errors': observed,
106                 'note': 'Measured on independently trained benchmark models.'}
107    report = bench.make_report('tabular', 'variable_mlp', base, idea,
108        {'certificate_calibration': cal, 'candidate_certificates': candidates,
109         'chosen_depth': chosen_depth, 'idea_grid': idea_trials,
110         'mechanism_signature': signature})
111    report['mechanism_signature'] = signature
112    report['bench_report'] = {'track': 'tabular', 'model': 'variable_mlp',
113                              'calibration': cal, 'chosen_depth': chosen_depth}
114    Path('bench_report.json').write_text(json.dumps(report, indent=2))
115    print(json.dumps(report, indent=2))
116
117if __name__ == '__main__':
118    main()