1import sys, json, math, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
  7
  8TRACK = 'resolution_hilbert_regression'
  9MODEL = 'mlp_tiny'
 10EPOCHS = 12
 11LRS = (0.001, 0.003, 0.01)
 12RESOLUTIONS = (16, 32, 64)
 13KAPPAS = (0.5, 1.0, 1.5)
 14
 15
 16def math_check():
 17    ms = np.array([4, 8, 16, 32, 64])
 18    grid = np.linspace(0, 1, 20001)
 19    hs = np.array([np.max(np.min(np.abs(grid[:, None] - np.linspace(0, 1, m)[None, :]), axis=1)) for m in ms])
 20    theory = 1.0 / (2.0 * (ms - 1))
 21    slope = float(np.polyfit(np.log(ms), np.log(hs), 1)[0])
 22    return {'m': ms.tolist(), 'fill_distance': hs.tolist(), 'theory': theory.tolist(),
 23            'loglog_slope': slope, 'max_theory_error': float(np.max(np.abs(hs-theory))),
 24            'confirmed': bool(abs(slope + 1) < .15 and np.max(np.abs(hs-theory)) < 1e-8)}
 25
 26
 27def load(seed, n_pairs, m):
 28    d = get_dataset(TRACK, int(seed), int(n_pairs), 200)
 29    # Resolution is the retained number of function/Hilbert coordinates.
 30    for key in ('xtr', 'xte'):
 31        d[key] = d[key][:, :int(m)].contiguous()
 32    return d
 33
 34
 35def train_one(seed, n_pairs, m, lr, return_model=False):
 36    d = load(seed, n_pairs, m)
 37    torch.manual_seed(10000 + int(seed))
 38    np.random.seed(10000 + int(seed))
 39    net = make_model(MODEL, tuple(d['xtr'].shape[1:]), 1)
 40    t0 = time.perf_counter()
 41    trained, metric, hist = train_model(net, d, epochs=EPOCHS, lr=float(lr), batch=128, log=lambda *_: None)
 42    rec = {'metric': float(metric), 'seconds': float(time.perf_counter()-t0),
 43           'n_pairs': int(n_pairs), 'm': int(m), 'lr': float(lr),
 44           'train_last': float(hist[-1]) if hist else None}
 45    if return_model and trained is not None:
 46        trained.eval()
 47        dev = next(trained.parameters()).device
 48        with torch.no_grad():
 49            pred = trained(d['xte'].to(dev)).detach().cpu().numpy().reshape(-1)
 50        truth = d['yte'].numpy().reshape(-1)
 51        rec.update({'pred_std': float(np.std(pred)), 'truth_std': float(np.std(truth)),
 52                    'pred_truth_corr': float(np.corrcoef(pred, truth)[0, 1])})
 53    return rec
 54
 55
 56def main():
 57    check = math_check()
 58    assert check['confirmed']
 59    # All budgets used by the idea are also evaluated for the baseline.
 60    budgets = sorted(set(min(400, int(math.ceil(m ** k))) for m in RESOLUTIONS for k in KAPPAS))
 61    base_grid = [{'lr': lr, 'n_pairs': n, 'm': 64, 'epochs': EPOCHS} for lr in LRS for n in budgets]
 62    def base_factory(cfg):
 63        return lambda seed: train_one(seed, cfg['n_pairs'], cfg['m'], cfg['lr'])['metric']
 64    base = sweep_baseline(base_factory, base_grid)
 65    best = base['best_cfg']
 66
 67    # Idea sweep has the same learning-rate union and the a-priori kappa values.
 68    idea_grid = []
 69    for lr in LRS:
 70        for k in KAPPAS:
 71            n = min(400, int(math.ceil(64 ** k)))
 72            vals = [train_one(s, n, 64, lr)['metric'] for s in range(4)]
 73            idea_grid.append({'cfg': {'lr': lr, 'kappa': k, 'n_pairs': n, 'm': 64, 'epochs': EPOCHS},
 74                              'mean': float(np.mean(vals))})
 75    chosen = min(idea_grid, key=lambda z: z['mean'])['cfg']
 76
 77    bs, ins, details = [], [], []
 78    for seed in range(8):
 79        b = train_one(seed, best['n_pairs'], 64, best['lr'])
 80        i = train_one(seed, chosen['n_pairs'], 64, chosen['lr'])
 81        bs.append(b['metric']); ins.append(i['metric'])
 82        details.append({'seed': seed, 'baseline': b, 'idea': i})
 83    base['full']['paired_details'] = details
 84    idea = {'config': chosen, 'per_seed': ins, 'mean': float(np.mean(ins)),
 85            'std': float(np.std(ins)), 'n': 8, 'details': details}
 86
 87    # Behavior-based signature: predictions from trained models across resolutions.
 88    observed = []
 89    for m in RESOLUTIONS:
 90        n = min(400, int(math.ceil(m ** 1.0)))
 91        r = train_one(0, n, m, chosen['lr'], return_model=True)
 92        observed.append({'m': m, 'n_pairs': n, 'test_mse': r['metric'],
 93                         'pred_std': r['pred_std'], 'truth_std': r['truth_std'],
 94                         'pred_truth_corr': r['pred_truth_corr']})
 95    sig = {'trained_model_observations': observed,
 96           'prediction_error_at_larger_resolution_lower': bool(observed[-1]['test_mse'] < observed[0]['test_mse']),
 97           'confirmed': bool(observed[-1]['test_mse'] < observed[0]['test_mse'])}
 98    report = make_report(TRACK, MODEL,
 99                         {'sweep': base['sweep'], 'best_cfg': best, 'full': {'per_seed': bs, 'mean': float(np.mean(bs)), 'std': float(np.std(bs)), 'n': 8}},
100                         idea,
101                         {'math_check': check, 'idea_sweep': idea_grid,
102                          'mechanism_signature': sig})
103    Path('bench_report.json').write_text(json.dumps(report, indent=2))
104    print(json.dumps(report, indent=2))
105
106if __name__ == '__main__':
107    main()