1import os, sys, json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset, make_model, train_model, sweep_baseline, evaluate, make_report
  7
  8SEEDS=tuple(range(8))
  9GRID=[{'lr':1e-3,'penalty':0.0},{'lr':3e-3,'penalty':0.0},{'lr':6e-3,'penalty':0.0},
 10      {'lr':1e-3,'penalty':0.01},{'lr':3e-3,'penalty':0.01},{'lr':6e-3,'penalty':0.01},
 11      {'lr':1e-3,'penalty':0.05},{'lr':3e-3,'penalty':0.05},{'lr':6e-3,'penalty':0.05}]
 12EPOCHS=14
 13sig_rows=[]
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 18
 19def baseline_fn(cfg):
 20    def run(seed):
 21        seed_all(seed)
 22        ds=get_dataset('dynamics',seed,n_train=800,n_test=300)
 23        model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 24        net, metric, _=train_model(model,ds,epochs=EPOCHS,lr=cfg['lr'],batch=128)
 25        # Behavior signature measured on this trained baseline model.
 26        dev=next(net.parameters()).device
 27        with torch.no_grad():
 28            xt=ds['xte'].to(dev)
 29            pred=net(xt); theta=xt[:,21:22]
 30            viol=(pred.abs()>0.94*theta.abs()).float().mean().item()
 31            ratio=(pred.abs()/(theta.abs()+1e-3)).mean().item()
 32        sig_rows.append({'kind':'baseline','seed':seed,'lr':cfg['lr'],
 33                         'mse':float(metric),'contraction_violation_rate':viol,
 34                         'mean_pred_abs_over_theta':ratio})
 35        return float(metric)
 36    return run
 37
 38def idea_train(seed, cfg):
 39    seed_all(seed)
 40    ds=get_dataset('dynamics',seed,n_train=800,n_test=300)
 41    model=make_model('rnn_small',ds['input_shape'],ds['out_dim'])
 42    for device in (['cuda'] if torch.cuda.is_available() else []) + ['cpu']:
 43        try:
 44            net=model.to(device); x=ds['xtr'].to(device); y=ds['ytr'].to(device)
 45            opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
 46            rho=0.94
 47            for _ in range(EPOCHS):
 48                net.train(); perm=torch.randperm(len(x),device=device)
 49                for i in range(0,len(x),128):
 50                    ix=perm[i:i+128]; pred=net(x[ix]); target=y[ix]
 51                    theta=x[ix,21:22]
 52                    lyap=torch.relu(pred.abs()-rho*theta.abs())**2
 53                    loss=((pred-target)**2).mean()+cfg['penalty']*lyap.mean()
 54                    opt.zero_grad(); loss.backward(); opt.step()
 55            net.eval()
 56            with torch.no_grad():
 57                pred=net(ds['xte'].to(device)); target=ds['yte'].to(device)
 58                mse=((pred-target)**2).mean().item()
 59                theta=ds['xte'][:,21:22].to(device)
 60                ratio=(pred.abs()/(theta.abs()+1e-3)).mean().item()
 61                violation=(pred.abs()>rho*theta.abs()).float().mean().item()
 62            sig_rows.append({'kind':'idea','seed':seed,'lr':cfg['lr'],'penalty':cfg['penalty'],
 63                             'mse':mse,'mean_pred_abs_over_theta':ratio,
 64                             'contraction_violation_rate':violation})
 65            return float(mse)
 66        except RuntimeError:
 67            if device=='cuda': torch.cuda.empty_cache(); continue
 68            raise
 69    raise RuntimeError('training failed')
 70
 71def idea_fn(cfg): return lambda seed: idea_train(seed,cfg)
 72
 73def main():
 74    base=sweep_baseline(baseline_fn,GRID)
 75    blr=base['best_cfg']['lr']
 76    idea_grid=[{'lr':blr,'penalty':0.01},{'lr':blr,'penalty':0.05},
 77               {'lr':(blr/3 if blr>1.1e-3 else 3e-3),'penalty':0.01}]
 78    tried=[]; best=None
 79    for cfg in idea_grid:
 80        r=evaluate(idea_fn(cfg),SEEDS); tried.append({'cfg':cfg,'mean':r['mean'],'full':r})
 81        if best is None or r['mean']<best['result']['mean']:
 82            best={'cfg':cfg,'result':r}
 83    chosen=best['cfg']
 84    ir=[r for r in sig_rows if r.get('kind')=='idea' and r.get('penalty')==chosen['penalty'] and r.get('lr')==chosen['lr']]
 85    br=[r for r in sig_rows if r.get('kind')=='baseline' and r.get('lr')==blr]
 86    iv=float(np.mean([r['contraction_violation_rate'] for r in ir]))
 87    bv=float(np.mean([r['contraction_violation_rate'] for r in br]))
 88    signature={'prediction':'Lyapunov penalty lowers observed contraction-violation rate',
 89               'predicted_direction':'lower','baseline_violation_rate':bv,
 90               'idea_violation_rate':iv,
 91               'baseline_mean_abs_prediction_over_theta':float(np.mean([r['mean_pred_abs_over_theta'] for r in br])),
 92               'idea_mean_abs_prediction_over_theta':float(np.mean([r['mean_pred_abs_over_theta'] for r in ir])),
 93               'rho':0.94,'confirmed':bool(iv < bv)}
 94    report=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},best['result'],
 95                       {'mechanism_signature':signature,'idea_sweep':tried,
 96                        'protocol_notes':'Dynamics is the structurally matched control/stability track; identical rnn_small systems and datasets, only training loss differs.'})
 97    report['idea_selected_cfg']=chosen
 98    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
 99    print(json.dumps(report,indent=2))
100if __name__=='__main__': main()