1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, make_report
  9from bench.protocol import evaluate, sweep_baseline
 10
 11SEEDS = tuple(range(8))
 12EPOCHS = 8
 13BATCH = 128
 14# Union parity: every idea step size is also evaluated by Adam.
 15LRS = [1e-3, 3e-3, 6e-3]
 16WDS = [0.0, 1e-4]
 17
 18
 19def seed_all(seed):
 20    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 21    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 22
 23
 24def dev():
 25    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 26
 27
 28def batches(n, seed, device):
 29    g = torch.Generator(device=device).manual_seed(seed)
 30    ix = torch.randperm(n, generator=g, device=device)
 31    for i in range(0, n, BATCH):
 32        yield ix[i:i+BATCH]
 33
 34
 35def train_adam(seed, lr, wd, capture=False):
 36    seed_all(seed)
 37    ds = get_dataset('tabular', seed, 400, 400)
 38    device = dev()
 39    try:
 40        net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device)
 41        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 42        opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)
 43        history = []
 44        for ep in range(EPOCHS):
 45            net.train()
 46            for ix in batches(len(x), seed + 1009 * ep, device):
 47                loss = nn.functional.mse_loss(net(x[ix]), y[ix])
 48                opt.zero_grad(set_to_none=True)
 49                loss.backward(); opt.step()
 50            history.append(float(loss.detach().cpu()))
 51        net.eval()
 52        with torch.no_grad():
 53            metric = float(nn.functional.mse_loss(net(ds['xte'].to(device)), ds['yte'].to(device)).cpu())
 54        return metric, {'history': history, 'energy': []}
 55    except RuntimeError:
 56        if device.type == 'cuda':
 57            torch.cuda.empty_cache()
 58            return train_adam_cpu(seed, lr, wd, capture)
 59        raise
 60
 61
 62def train_adam_cpu(seed, lr, wd, capture=False):
 63    seed_all(seed)
 64    ds = get_dataset('tabular', seed, 400, 400)
 65    net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim'])
 66    x, y = ds['xtr'], ds['ytr']
 67    opt = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=wd)
 68    history = []
 69    for ep in range(EPOCHS):
 70        for ix in batches(len(x), seed + 1009 * ep, torch.device('cpu')):
 71            loss = nn.functional.mse_loss(net(x[ix]), y[ix])
 72            opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 73        history.append(float(loss.detach()))
 74    with torch.no_grad(): metric = float(nn.functional.mse_loss(net(ds['xte']), ds['yte']))
 75    return metric, {'history': history, 'energy': []}
 76
 77
 78def flat_params(params):
 79    return torch.cat([p.detach().reshape(-1) for p in params])
 80
 81
 82def train_leapfrog(seed, h, mass=1.0, capture=False):
 83    seed_all(seed)
 84    ds = get_dataset('tabular', seed, 400, 400)
 85    device = dev()
 86    try:
 87        net = make_model('mlp_tiny', ds['input_shape'], ds['out_dim']).to(device)
 88        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
 89        params = list(net.parameters())
 90        mom = [torch.zeros_like(p) for p in params]
 91        energies = []; history = []
 92        # Two gradient evaluations per leapfrog step, matching Adam's budget by epochs.
 93        for ep in range(EPOCHS):
 94            net.train()
 95            for bi, ix in enumerate(batches(len(x), seed + 1009 * ep, device)):
 96                net.zero_grad(set_to_none=True)
 97                loss = nn.functional.mse_loss(net(x[ix]), y[ix]); loss.backward()
 98                g1 = [p.grad.detach().clone() for p in params]
 99                with torch.no_grad():
100                    for p, q, g in zip(params, mom, g1):
101                        q.sub_(0.5 * h * g); p.add_(h * q / mass)
102                net.zero_grad(set_to_none=True)
103                loss2 = nn.functional.mse_loss(net(x[ix]), y[ix]); loss2.backward()
104                g2 = [p.grad.detach().clone() for p in params]
105                with torch.no_grad():
106                    for q, g in zip(mom, g2): q.sub_(0.5 * h * g)
107                if capture and (bi == 0 or bi == len(list(batches(len(x), seed + 1009 * ep, device))) - 1):
108                    with torch.no_grad():
109                        kinetic = 0.5 * sum(float((q*q).sum().cpu()) / mass for q in mom)
110                        potential = float(nn.functional.mse_loss(net(x), y).detach().cpu())
111                        energies.append(potential + kinetic)
112            history.append(float(loss2.detach().cpu()))
113        net.eval()
114        with torch.no_grad(): metric = float(nn.functional.mse_loss(net(ds['xte'].to(device)), ds['yte'].to(device)).cpu())
115        return metric, {'history': history, 'energy': energies}
116    except RuntimeError:
117        if device.type == 'cuda':
118            torch.cuda.empty_cache()
119            old = globals()['dev']; globals()['dev'] = lambda: torch.device('cpu')
120            try: return train_leapfrog(seed, h, mass, capture)
121            finally: globals()['dev'] = old
122        raise
123
124
125def main():
126    # Baseline sweep includes all step sizes and both central Adam weight-decay values.
127    grid = [{'lr': lr, 'weight_decay': wd} for lr in LRS for wd in WDS]
128    base = sweep_baseline(lambda c: lambda s: train_adam(s, c['lr'], c['weight_decay'])[0], grid)
129    # Idea sweep: best baseline h plus two nearby settings; all are in baseline union.
130    idea_grid = [{'h': h, 'mass': 1.0} for h in LRS]
131    idea_cfg_results = []
132    for cfg in idea_grid:
133        r = evaluate(lambda s, c=cfg: train_leapfrog(s, c['h'], c['mass'])[0])
134        idea_cfg_results.append({'cfg': cfg, 'result': r})
135    best = min(idea_cfg_results, key=lambda z: z['result']['mean'])
136    idea = best['result']; cfg = best['cfg']
137
138    # Re-test the mechanism on trained systems: energy oscillation versus h.
139    hs = [0.001, 0.003, 0.006]
140    ranges = []
141    for h in hs:
142        vals = []
143        for s in range(4):
144            z = train_leapfrog(s, h, 1.0, capture=True)[1]['energy']
145            vals.append(max(z) - min(z) if z else float('nan'))
146        ranges.append(float(np.nanmean(vals)))
147    slope = float(np.polyfit(np.log(hs), np.log(np.maximum(ranges, 1e-30)), 1)[0])
148    sig = {'prediction': 'leapfrog trained-model energy oscillation scales approximately as h^2',
149           'step_sizes': hs, 'observed_energy_ranges': ranges,
150           'loglog_slope': slope, 'confirmed': bool(1.5 < slope < 2.5)}
151    report = make_report('tabular', 'mlp_tiny', base, idea, {
152        'track_choice': 'optimizer intervention structurally matches the tabular optimizer track',
153        'baseline_grid': grid, 'idea_grid': idea_grid, 'idea_best_cfg': cfg,
154        'mechanism_signature': sig})
155    report['idea_sweep'] = idea_cfg_results
156    Path('bench_report.json').write_text(json.dumps(report, indent=2))
157    print(json.dumps(report, indent=2))
158
159
160if __name__ == '__main__': main()