1import sys, json
  2import numpy as np
  3import torch
  4
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  7
  8TRACK = 'tabular'
  9MODEL = 'mlp_tiny'
 10EPOCHS = 20
 11BATCH = 128
 12NTR, NTE = 400, 400
 13B = 6.0
 14GRID = [
 15    {'lr': 1e-3, 'weight_decay': 0.0},
 16    {'lr': 3e-3, 'weight_decay': 0.0},
 17    {'lr': 1e-2, 'weight_decay': 0.0},
 18]
 19
 20
 21def spectral_direction(x, y, bound=B):
 22    X = np.asarray(x, dtype=np.float64)
 23    yy = np.asarray(y).reshape(-1)
 24    t = np.clip(yy - yy.mean(), -bound, bound)
 25    D = (X.T * t) @ X / len(X)
 26    D = (D + D.T) * 0.5
 27    vals, vecs = np.linalg.eigh(D)
 28    return vecs[:, -1].astype(np.float32), float(vals[-1] - vals[-2]), float(vals[-1]), float(np.var(t))
 29
 30
 31def init_spectral(net, ds):
 32    direction, gap, top, tvar = spectral_direction(ds['xtr'].cpu().numpy(), ds['ytr'].cpu().numpy())
 33    first = net[0]
 34    gen = torch.Generator(device='cpu').manual_seed(9917)
 35    signs = torch.where(torch.randn(first.out_features, generator=gen) >= 0, 1., -1.)
 36    with torch.no_grad():
 37        first.weight.copy_(0.35 * signs[:, None] * torch.from_numpy(direction)[None, :])
 38        first.bias.zero_()
 39    return {'eigengap': gap, 'top_eigenvalue': top, 'weight_variance': tvar}
 40
 41
 42def make_run(cfg, spectral):
 43    def run(seed):
 44        np.random.seed(seed)
 45        torch.manual_seed(seed)
 46        ds = get_dataset(TRACK, seed, n_train=NTR, n_test=NTE)
 47        net = make_model(MODEL, ds['input_shape'], ds['out_dim'])
 48        sig = init_spectral(net, ds) if spectral else None
 49        _, metric, history = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'],
 50                                         batch=BATCH, weight_decay=cfg['weight_decay'])
 51        run.last_signature = sig
 52        run.last_model = net
 53        run.last_ds = ds
 54        run.last_history = history
 55        return float(metric)
 56    return run
 57
 58
 59def behavior_signature(base_cfg, idea_res):
 60    # Retest the stage-1 prediction on trained-model behavior: compare the
 61    # first-layer projection energy along the data-derived spectral direction.
 62    rows = []
 63    for seed, metric in enumerate(idea_res['per_seed']):
 64        ds = get_dataset(TRACK, seed, n_train=NTR, n_test=NTE)
 65        direction, gap, top, tvar = spectral_direction(ds['xtr'].numpy(), ds['ytr'].numpy())
 66        d = torch.from_numpy(direction)
 67        bnet = make_run(base_cfg, False); bnet(seed)
 68        inet = make_run(base_cfg, True); inet(seed)
 69        bw = bnet.last_model[0].weight.detach().cpu()
 70        iw = inet.last_model[0].weight.detach().cpu()
 71        bp = float(((bw @ d) ** 2).mean())
 72        ip = float(((iw @ d) ** 2).mean())
 73        rows.append({'seed': seed, 'spectral_gap': gap, 'baseline_projection_energy': bp,
 74                     'idea_projection_energy': ip, 'idea_metric': float(metric)})
 75    observed_ratio = float(np.mean([r['idea_projection_energy'] / (r['baseline_projection_energy'] + 1e-12) for r in rows]))
 76    predicted = 'idea starts in recovered spectral subspace, so projection energy is higher'
 77    confirmed = bool(observed_ratio > 1.2)
 78    return {'prediction': predicted, 'predicted_vs_observed': {'observed_energy_ratio': observed_ratio,
 79            'threshold': 1.2}, 'confirmed': confirmed, 'trained_model_measurements': rows}
 80
 81
 82def main():
 83    # Baseline sweep includes every idea learning rate; sweep_baseline selects
 84    # using four seeds and re-evaluates the winner on all eight.
 85    base = sweep_baseline(lambda cfg: make_run(cfg, False), GRID)
 86    # Equal-budget idea sweep at all same settings, selecting on the same four
 87    # sweep seeds, then evaluate its best configuration on all eight.
 88    idea_trials = []
 89    best_cfg, best_mean = None, float('inf')
 90    for cfg in GRID:
 91        r = evaluate(make_run(cfg, True), seeds=(0, 1, 2, 3))
 92        idea_trials.append({'cfg': cfg, 'mean': r['mean']})
 93        if r['mean'] < best_mean:
 94            best_mean, best_cfg = r['mean'], cfg
 95    idea_res = evaluate(make_run(best_cfg, True))
 96    idea_res['sweep'] = idea_trials
 97    extra = behavior_signature(base['best_cfg'], idea_res)
 98    report = make_report(TRACK, MODEL, base, idea_res, extra=extra)
 99    report['protocol'] = {'paired_seeds': list(range(8)), 'n_train': NTR, 'n_test': NTE,
100                          'epochs': EPOCHS, 'batch': BATCH, 'common_grid': GRID,
101                          'track_rationale': 'Initialization intervention uses the designated tabular MLP track.'}
102    with open('bench_report.json', 'w') as f:
103        json.dump(report, f, indent=2)
104    print(json.dumps(report, indent=2))
105
106
107if __name__ == '__main__':
108    main()