1import sys, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6from scipy.linalg import expm, eigvals
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, sweep_baseline, evaluate, make_report
  9
 10OUT = Path('bench_report.json')
 11SEEDS = tuple(range(8))
 12# Union parity: baseline and idea both evaluated at every lr in this grid.
 13GRID = [{'lr': 0.001}, {'lr': 0.003}, {'lr': 0.006}]
 14EPOCHS = 18
 15BATCH = 128
 16
 17
 18def set_seed(s):
 19    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 20    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 21
 22
 23def math_check():
 24    A = np.array([[-2., 5.], [0., -.5]])
 25    d = 1.25
 26    T = .4
 27    g0 = np.log(max(abs(eigvals(expm(A*T))))) / T
 28    gd = np.log(max(abs(eigvals(expm((A-d*np.eye(2))*T))))) / T
 29    return {'predicted_damping_shift': d, 'observed_shift': float(g0-gd),
 30            'abs_error': float(abs((g0-gd)-d)), 'passed': bool(abs(g0-gd-d)<1e-10)}
 31
 32
 33def device():
 34    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 35
 36
 37def train(seed, lr, periodic=False, amp=0.15, period=8, return_model=False):
 38    set_seed(seed)
 39    d = get_dataset('dynamics', seed, n_train=400, n_test=400)
 40    dev = device()
 41    try:
 42        model = make_model('rnn_small', d['input_shape'], d['out_dim']).to(dev)
 43        opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
 44        lossfn = nn.MSELoss()
 45        x, y = d['xtr'].to(dev), d['ytr'].to(dev)
 46        model.train()
 47        step = 0
 48        for ep in range(EPOCHS):
 49            # fixed order per seed; identical minibatches between systems
 50            perm = torch.arange(len(x), device=dev)
 51            for start in range(0, len(x), BATCH):
 52                ix = perm[start:start+BATCH]
 53                phase = 2*math.pi*(step % period)/period
 54                factor = 1.0 + amp*math.sin(phase) if periodic else 1.0
 55                for pg in opt.param_groups: pg['lr'] = lr*factor
 56                opt.zero_grad(set_to_none=True)
 57                pred = model(x[ix]); loss = lossfn(pred, y[ix])
 58                loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
 59                opt.step(); step += 1
 60        model.eval()
 61        with torch.no_grad():
 62            pred = model(d['xte'].to(dev)); metric = lossfn(pred, d['yte'].to(dev)).item()
 63        if return_model: return metric, model, d, dev
 64        return metric
 65    except RuntimeError:
 66        # Explicit CPU fallback for constrained/shared CUDA environments.
 67        if dev.type == 'cuda':
 68            torch.cuda.empty_cache()
 69            return train_cpu(seed, lr, periodic, amp, period, return_model)
 70        raise
 71
 72
 73def train_cpu(seed, lr, periodic=False, amp=.15, period=8, return_model=False):
 74    old = torch.cuda.is_available
 75    # same implementation forced onto CPU, avoiding recursive GPU selection
 76    set_seed(seed); d=get_dataset('dynamics',seed,n_train=400,n_test=400)
 77    model=make_model('rnn_small',d['input_shape'],d['out_dim'])
 78    opt=torch.optim.AdamW(model.parameters(),lr=lr,weight_decay=1e-4); lossfn=nn.MSELoss(); x,y=d['xtr'],d['ytr']; step=0
 79    model.train()
 80    for ep in range(EPOCHS):
 81        for start in range(0,len(x),BATCH):
 82            phase=2*math.pi*(step%period)/period; f=1+amp*math.sin(phase) if periodic else 1
 83            opt.param_groups[0]['lr']=lr*f; opt.zero_grad(); loss=lossfn(model(x[start:start+BATCH]),y[start:start+BATCH]); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step(); step+=1
 84    model.eval()
 85    with torch.no_grad(): metric=lossfn(model(d['xte']),d['yte']).item()
 86    return (metric,model,d,torch.device('cpu')) if return_model else metric
 87
 88
 89def signature():
 90    # Measured on trained models: local output Jacobian products over one schedule.
 91    vals=[]
 92    for s in (0,1):
 93        mb, base, d, dev = train(s,.003,False,return_model=True)
 94        mi, idea, _, _ = train(s,.003,True,return_model=True)
 95        x=d['xte'][:1].to(dev).requires_grad_(True)
 96        def local_jac(m):
 97            z=x.detach().clone().requires_grad_(True); out=m(z)[0,0]
 98            g=torch.autograd.grad(out,z,create_graph=False)[0].detach().flatten().cpu().numpy()
 99            # Scalar observable Jacobian norm is a conservative NN-scale tangent proxy.
100            return float(np.linalg.norm(g))
101        vals.append({'seed':s,'baseline_jacobian_norm':local_jac(base),'idea_jacobian_norm':local_jac(idea),'baseline_mse':mb,'idea_mse':mi})
102    b=np.mean([v['baseline_jacobian_norm'] for v in vals]); i=np.mean([v['idea_jacobian_norm'] for v in vals])
103    return {'quantity':'trained-model local input Jacobian norm','predicted':'periodic modulation reduces perturbation growth','baseline_mean':float(b),'idea_mean':float(i),'relative_change':float((i-b)/max(b,1e-12)),'confirmed':bool(i<b)}
104
105
106def main():
107    sanity=math_check()
108    base=sweep_baseline(lambda cfg: lambda seed: train(seed,cfg['lr'],False), GRID)
109    # Idea at all union lrs; select best on the same four-seed selection protocol.
110    tried=[]
111    for cfg in GRID:
112        r=evaluate(lambda seed,cfg=cfg: train(seed,cfg['lr'],True), seeds=(0,1,2,3))
113        tried.append({'cfg':cfg,'mean':r['mean']})
114    best_cfg=min(tried,key=lambda z:z['mean'])['cfg']
115    idea=evaluate(lambda seed: train(seed,best_cfg['lr'],True), seeds=SEEDS)
116    report=make_report('dynamics','rnn_small',base,idea,{'math_sanity':sanity,'trained_jacobian':signature(),'idea_cfg':best_cfg,'idea_sweep':tried})
117    report['selection_note']='Baseline and idea share rnn_small, data, epochs, batch, AdamW weight decay, and union learning-rate grid; only periodic LR modulation differs.'
118    OUT.write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
119if __name__=='__main__': main()