1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  5from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report, permutation_pvalue
  6
  7SEEDS = tuple(range(8))
  8# Shared union: every idea lr is also evaluated by baseline.
  9GRID = [
 10    {'lr': 0.0015, 'epochs': 12},
 11    {'lr': 0.0030, 'epochs': 12},
 12    {'lr': 0.0060, 'epochs': 12},
 13]
 14K = 20
 15BATCH_SELECT = 10
 16NTRAIN = 400
 17NTEST = 400
 18
 19
 20def epsilon_scenario(N, k, beta):
 21    # Cheap finite-sample-style sanity proxy: binomial tail/union bound.
 22    # It is not substituted for the paper theorem in the NN experiment.
 23    return min(1.0, (k + math.log(1.0 / beta)) / N)
 24
 25
 26def train_on_indices(ds, cfg, indices=None):
 27    if indices is None:
 28        sub = ds
 29    else:
 30        sub = dict(ds)
 31        ix = np.asarray(indices, dtype=np.int64)
 32        sub['xtr'] = ds['xtr'][ix]
 33        sub['ytr'] = ds['ytr'][ix]
 34    net = make_model('mlp_tiny', tuple(sub['xtr'].shape[1:]), sub['out_dim'])
 35    return train_model(net, sub, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=0.0, log=lambda *_: None)
 36
 37
 38def greedy_compress(ds, cfg, k=K, batch_size=BATCH_SELECT):
 39    n = len(ds['xtr'])
 40    chosen = []
 41    remaining = np.arange(n)
 42    while len(chosen) < k:
 43        # Retrain on the current active set before scoring all scenarios.
 44        model, _, _ = train_on_indices(ds, cfg, chosen if chosen else remaining[:1])
 45        model.eval()
 46        import torch
 47        dev = next(model.parameters()).device
 48        with torch.no_grad():
 49            pred = model(ds['xtr'].to(dev)).reshape(-1).cpu().numpy()
 50        y = ds['ytr'].reshape(-1).cpu().numpy()
 51        residual = np.abs(y - pred)
 52        mask = np.ones(n, dtype=bool); mask[np.asarray(chosen, dtype=int)] = False
 53        cand = remaining[mask[remaining]]
 54        take = min(batch_size, k-len(chosen))
 55        # Select the most informative/high-loss scenarios, a practical active-set rule.
 56        picked = cand[np.argsort(residual[cand])[-take:]]
 57        chosen.extend([int(x) for x in picked])
 58    model, metric, hist = train_on_indices(ds, cfg, chosen)
 59    return float(metric), chosen, model
 60
 61
 62def run_baseline(cfg):
 63    def one(seed):
 64        ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
 65        _, metric, _ = train_on_indices(ds, cfg)
 66        return metric
 67    return one
 68
 69
 70def run_idea(cfg):
 71    def one(seed):
 72        ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
 73        metric, _, _ = greedy_compress(ds, cfg)
 74        return metric
 75    return one
 76
 77
 78def behavior_signature(cfg, seed=0):
 79    ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
 80    # Retain models to measure predicted-vs-observed behavior, not an analytic identity.
 81    base_net, base_metric, _ = train_on_indices(ds, cfg)
 82    idea_m, chosen, idea_net = greedy_compress(ds, cfg)
 83    import torch
 84    with torch.no_grad():
 85        db = next(base_net.parameters()).device
 86        di = next(idea_net.parameters()).device
 87        pb = base_net(ds['xtr'].to(db)).reshape(-1).cpu().numpy()
 88        pi = idea_net(ds['xtr'].to(di)).reshape(-1).cpu().numpy()
 89    y = ds['ytr'].reshape(-1).cpu().numpy()
 90    rb, ri = np.abs(y-pb), np.abs(y-pi)
 91    # Stage-1 prediction: active compression should cover the high-residual tail.
 92    q = float(np.quantile(ri, 1-K/NTRAIN))
 93    selected_tail = float(np.mean(ri[np.asarray(chosen)] >= q))
 94    baseline_tail = float(np.mean(np.sort(rb)[-K:] >= np.quantile(rb, 1-K/NTRAIN)))
 95    return {'prediction': 'greedy selected scenarios cover the high-loss tail',
 96            'observed_selected_tail_fraction': selected_tail,
 97            'observed_baseline_top_tail_fraction': baseline_tail,
 98            'selected_count': len(chosen), 'idea_train_mse': float(np.mean(ri**2)),
 99            'baseline_train_mse': float(np.mean(rb**2)),
100            'confirmed': bool(selected_tail > 0.55)}
101
102
103def main():
104    # Core numerical check first: increasing N contracts the finite-sample proxy.
105    check = [epsilon_scenario(n, 2, 1e-5) for n in (100, 200, 400, 800)]
106    assert all(check[i] > check[i+1] for i in range(3))
107    baseline = sweep_baseline(run_baseline, GRID, seeds=SEEDS)
108    best = baseline['best_cfg']
109    # Three idea settings are exactly the shared three-point grid; report best.
110    idea_trials = []
111    for cfg in GRID:
112        r = evaluate(run_idea(cfg), seeds=SEEDS)
113        idea_trials.append({'cfg': cfg, **r})
114    idea_best = min(idea_trials, key=lambda x: x['mean'])
115    idea_res = {k: idea_best[k] for k in ('mean','std','per_seed','n')}
116    diffs = [a-b for a,b in zip(idea_res['per_seed'], baseline['full']['per_seed'])]
117    cmp = make_report('tabular', 'mlp_tiny', baseline, idea_res,
118                      extra=behavior_signature(idea_best['cfg']))
119    cmp['idea']['sweep'] = idea_trials
120    cmp['core_math_check'] = {'proxy_epsilon_N_100_200_400_800': check,
121                              'monotone_decrease': True,
122                              'N': 400, 'k': K, 'beta': 1e-5}
123    cmp['comparison']['paired_delta_mean'] = float(np.mean(diffs))
124    cmp['comparison']['permutation_pvalue'] = float(permutation_pvalue(diffs, n_perm=20000))
125    cmp['bench_report'] = {'track_justification': 'tabular is the harness-matched track for optimizer/training-dynamics/calibration interventions',
126                           'custom_track': None,
127                           'protocol': '8 paired seeds; baseline sweep and idea sweep share the exact lr/epoch union',
128                           'budget_note': 'small 400-sample, 12-epoch MLP; compression uses active-set retraining and final retraining'}
129    Path('bench_report.json').write_text(json.dumps(cmp, indent=2))
130    print(json.dumps(cmp, indent=2))
131
132if __name__ == '__main__': main()