1import os, sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6import torch
  7import torch.nn as nn
  8from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report
  9
 10ROOT = Path(__file__).resolve().parent
 11SEEDS = tuple(range(8))
 12# Common learning-rate union is used on both sides.
 13LR_GRID = [0.0015, 0.0030, 0.0045]
 14EPOCHS = 18
 15BATCH = 64
 16WEIGHT_DECAY = 0.0
 17PERIOD = 8
 18AMP_GRID = [0.15, 0.30, 0.45]
 19
 20
 21def seed_all(seed):
 22    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 23    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 24
 25
 26def device():
 27    return 'cuda' if torch.cuda.is_available() else 'cpu'
 28
 29
 30def baseline_one(cfg, seed, retain=False):
 31    seed_all(seed)
 32    ds = get_dataset('dynamics', seed, n_train=400, n_test=100)
 33    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 34    if not retain:
 35        _, metric, _ = train_model(
 36            model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
 37            weight_decay=WEIGHT_DECAY, log=lambda *a, **k: None)
 38        if metric is None:
 39            raise RuntimeError('baseline training failed')
 40        return float(metric)
 41    return model, ds
 42
 43
 44def cyclic_one(cfg, seed, return_model=False):
 45    """Adam with phase-periodic step size and first-moment coefficient."""
 46    seed_all(seed)
 47    ds = get_dataset('dynamics', seed, n_train=400, n_test=100)
 48    dev = device()
 49    try:
 50        model = make_model('rnn_small', ds['input_shape'], ds['out_dim']).to(dev)
 51        x, y = ds['xtr'].to(dev), ds['ytr'].to(dev)
 52        xe, ye = ds['xte'].to(dev), ds['yte'].to(dev)
 53        lossf = nn.MSELoss()
 54        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], betas=(0.9, 0.999), weight_decay=WEIGHT_DECAY)
 55        n = x.shape[0]; step = 0
 56        for ep in range(EPOCHS):
 57            g = torch.Generator(device='cpu'); g.manual_seed(seed + 1009 * ep)
 58            order = torch.randperm(n, generator=g).to(dev)
 59            for st in range(0, n, BATCH):
 60                phase = 2 * math.pi * (step % PERIOD) / PERIOD
 61                eta = cfg['lr'] * (1.0 + cfg['amp'] * math.sin(phase))
 62                beta1 = min(0.98, max(0.50, 0.82 + 0.10 * math.cos(phase)))
 63                for group in opt.param_groups:
 64                    group['lr'] = eta; group['betas'] = (beta1, 0.999)
 65                ix = order[st:st + BATCH]
 66                opt.zero_grad(set_to_none=True)
 67                loss = lossf(model(x[ix]), y[ix]); loss.backward()
 68                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
 69                opt.step(); step += 1
 70        with torch.no_grad(): metric = float(lossf(model(xe), ye).detach().cpu())
 71        return (model, ds, step) if return_model else metric
 72    except Exception:
 73        if dev == 'cuda':
 74            torch.cuda.empty_cache()
 75            return cyclic_one_cpu(cfg, seed, return_model)
 76        raise
 77
 78
 79def cyclic_one_cpu(cfg, seed, return_model=False):
 80    seed_all(seed)
 81    ds = get_dataset('dynamics', seed, n_train=400, n_test=100)
 82    model = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
 83    x, y, xe, ye = ds['xtr'], ds['ytr'], ds['xte'], ds['yte']
 84    lossf = nn.MSELoss(); opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], betas=(.9, .999), weight_decay=WEIGHT_DECAY)
 85    step = 0
 86    for ep in range(EPOCHS):
 87        gen = torch.Generator(); gen.manual_seed(seed + 1009 * ep)
 88        order = torch.randperm(len(x), generator=gen)
 89        for st in range(0, len(x), BATCH):
 90            phase = 2 * math.pi * (step % PERIOD) / PERIOD
 91            for q in opt.param_groups:
 92                q['lr'] = cfg['lr'] * (1 + cfg['amp'] * math.sin(phase))
 93                q['betas'] = (min(.98, max(.5, .82 + .10 * math.cos(phase))), .999)
 94            opt.zero_grad(set_to_none=True)
 95            loss = lossf(model(x[order[st:st + BATCH]]), y[order[st:st + BATCH]])
 96            loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step(); step += 1
 97    with torch.no_grad(): metric = float(lossf(model(xe), ye))
 98    return (model, ds, step) if return_model else metric
 99
100
101def mechanism_signature(cfg, seed=0):
102    """Measure one-period local contraction using a trained rnn_small model."""
103    model, ds, steps = cyclic_one(cfg, seed, return_model=True)
104    model.eval(); params = [p for p in model.parameters() if p.requires_grad]
105    def clone_state(): return [p.detach().clone() for p in params]
106    base = clone_state(); eps = 1e-5
107    with torch.no_grad(): params[0].add_(eps)
108    pert = clone_state()
109    lossf = nn.MSELoss(); probe_dev = next(model.parameters()).device
110    x, y = ds['xtr'].to(probe_dev), ds['ytr'].to(probe_dev)
111    def run(vec):
112        with torch.no_grad():
113            for p, v in zip(params, vec): p.copy_(v)
114        local = torch.optim.SGD(params, lr=cfg['lr'])
115        for k in range(PERIOD):
116            local.zero_grad(); lossf(model(x), y).backward(); local.step()
117        return clone_state()
118    out0, out1 = run(base), run(pert)
119    d0 = math.sqrt(sum(float(((a - b) ** 2).sum()) for a, b in zip(pert, base)))
120    d1 = math.sqrt(sum(float(((a - b) ** 2).sum()) for a, b in zip(out1, out0)))
121    observed = d1 / max(d0, 1e-30)
122    mus = [.82 + .10 * math.cos(2 * math.pi * k / PERIOD) for k in range(PERIOD)]
123    predicted = float(np.prod(mus))
124    return {'period': PERIOD, 'predicted_transverse_multiplier': predicted,
125            'observed_trained_model_multiplier': observed,
126            'ratio_observed_to_predicted': observed / max(abs(predicted), 1e-30),
127            'confirmed': bool(abs(math.log(max(observed, 1e-30)) - math.log(max(abs(predicted), 1e-30))) < 1.0),
128            'seed': seed, 'probe': 'full_batch_gradient_on_trained_rnn'}
129
130
131def main():
132    grid = [{'lr': lr, 'amp': amp} for lr in LR_GRID for amp in AMP_GRID]
133    base = sweep_baseline(lambda c: (lambda s: baseline_one(c, s)), grid, seeds=(0, 1, 2, 3))
134    idea_runs = []
135    for c in grid:
136        idea_runs.append({'cfg': c, 'result': evaluate(lambda s, c=c: cyclic_one(c, s), seeds=SEEDS)})
137    best = min(idea_runs, key=lambda z: z['result']['mean'])
138    sig = mechanism_signature(best['cfg'], 0)
139    report = make_report('dynamics', 'rnn_small', base, best['result'], {
140        **sig, 'idea_grid': idea_runs,
141        'protocol_note': '8 paired seeds; baseline and idea share lr/amp union; 400/100 samples'
142    })
143    (ROOT / 'bench_report.json').write_text(json.dumps(report, indent=2))
144    print(json.dumps(report, indent=2))
145
146
147if __name__ == '__main__': main()