1import sys, os, json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
  8from bench.protocol import evaluate
  9
 10SEEDS = tuple(range(8))
 11SWEEP = tuple(range(4))
 12EPOCHS = 15
 13NTR, NTE = 1000, 400
 14
 15# Nominal cheap structured solver for the actuated pendulum. It holds each
 16# observed control over the same four microsteps used by the bench generator.
 17def background(x, g=9.81, damp=0.25):
 18    z = x.reshape(-1, 8, 3)
 19    th, om = z[:, 0, 0], z[:, 0, 1]
 20    states = []
 21    for k in range(8):
 22        u = z[:, k, 2]
 23        states.append(torch.stack((th, om, u), 1))
 24        for _ in range(4):
 25            om = om + (-g/10 * torch.sin(th) - damp*om + 2*u) * (.05/4)
 26            th = th + om * (.05/4)
 27    return torch.stack(states, 1), th.unsqueeze(1)
 28
 29class ResidualRNN(nn.Module):
 30    """Same rnn_small GRU+head, with analytic background and defect gate."""
 31    def __init__(self, gate_mid=0.08, gate_slope=3.0):
 32        super().__init__()
 33        self.core = make_model('rnn_small', (24,), 1)
 34        self.gate_mid, self.gate_slope = gate_mid, gate_slope
 35    def forward(self, x):
 36        bgseq, bgfinal = background(x)
 37        raw = x.reshape(-1, 8, 3)
 38        residual = raw.clone()
 39        residual[:, :, :2] = raw[:, :, :2] - bgseq[:, :, :2]
 40        # Discrete background defect: observed acceleration minus nominal one.
 41        dt = .05
 42        obs_acc = (raw[:, 1:, 1] - raw[:, :-1, 1]) / dt
 43        nom_acc = -9.81/10*torch.sin(raw[:, :-1, 0]) - .25*raw[:, :-1, 1] + 2*raw[:, :-1, 2]
 44        defect = torch.sqrt(torch.mean((obs_acc-nom_acc)**2, dim=1) + 1e-8)
 45        gate = torch.sigmoid(self.gate_slope*(torch.log(defect+1e-6)-math.log(self.gate_mid))).unsqueeze(1)
 46        closure = self.core(residual.reshape(x.shape[0], -1))
 47        return bgfinal + gate * closure
 48
 49def seed_all(s):
 50    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 51    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 52
 53def ds_for(seed):
 54    return get_dataset('dynamics', seed, n_train=NTR, n_test=NTE)
 55
 56def baseline_fn(cfg):
 57    def run(seed):
 58        seed_all(seed); d=ds_for(seed); net=make_model('rnn_small', d['input_shape'], d['ytr'].shape[-1])
 59        _, metric, _ = train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,weight_decay=cfg.get('weight_decay',0.0),log=lambda *_:None)
 60        return metric
 61    return run
 62
 63def idea_fn(cfg):
 64    def run(seed):
 65        seed_all(seed); d=ds_for(seed); net=ResidualRNN(cfg.get('gate_mid', 0.08), cfg.get('gate_slope', 3.0))
 66        _, metric, _ = train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,weight_decay=cfg.get('weight_decay',0.0),log=lambda *_:None)
 67        return metric
 68    return run
 69
 70def signature(cfg):
 71    seed_all(0); d=ds_for(0); net=ResidualRNN(cfg['gate_mid'],cfg['gate_slope'])
 72    net,_,_=train_model(net,d,epochs=EPOCHS,lr=cfg['lr'],batch=128,log=lambda *_:None)
 73    dev=next(net.parameters()).device
 74    x=d['xte'].to(dev)
 75    with torch.no_grad():
 76        bgseq,bgf=background(x); raw=x.reshape(-1,8,3)
 77        acc=(raw[:,1:,1]-raw[:,:-1,1])/.05
 78        nom=-.981*torch.sin(raw[:,:-1,0])-.25*raw[:,:-1,1]+2*raw[:,:-1,2]
 79        defect=torch.sqrt(torch.mean((acc-nom)**2,1)+1e-8)
 80        res=raw.clone()
 81        res[:,:,:2]=raw[:,:,:2]-bgseq[:,:,:2]
 82        closure=net.core(res.reshape(x.shape[0],-1)).abs().mean(1)
 83        bgerr=((bgf-d['yte'].to(dev))**2).sqrt().squeeze(1)
 84        order=torch.argsort(defect); lo=order[:len(order)//3]; hi=order[-len(order)//3:]
 85        corr=float(np.corrcoef(defect.cpu(),closure.cpu())[0,1])
 86        return {'trained_model':True,'defect_low_closure':float(closure[lo].mean()),'defect_high_closure':float(closure[hi].mean()),'defect_closure_corr':corr,'background_rmse':float(bgerr.mean()),'prediction':'closure should rise with defect and fall as background becomes accurate','confirmed':bool(closure[hi].mean()>closure[lo].mean() and corr>0)}
 87
 88def main():
 89    grid=[{'lr':1e-3,'gate_mid':0.08,'gate_slope':3.0},{'lr':3e-3,'gate_mid':0.08,'gate_slope':3.0},{'lr':1e-2,'gate_mid':0.08,'gate_slope':3.0}]
 90    base=sweep_baseline(baseline_fn,grid,seeds=SWEEP)
 91    # Search-space parity: identical three learning rates on both sides.
 92    idea_trials=[]
 93    for cfg in grid:
 94        r=evaluate(idea_fn(cfg),seeds=SEEDS); idea_trials.append({'cfg':cfg,'result':r})
 95    best=min(idea_trials,key=lambda q:q['result']['mean'])
 96    rep=make_report('dynamics','rnn_small',base,best['result'],extra=signature(best['cfg']))
 97    rep['idea_sweep']=idea_trials
 98    rep['protocol']={'paired_seeds':list(SEEDS),'sweep_seeds':list(SWEEP),'epochs':EPOCHS,'n_train':NTR,'n_test':NTE,'same_architecture':True,'baseline_grid':grid,'idea_grid':grid,'intervention':'analytic pendulum background, residual coordinates, defect sigmoid gate'}
 99    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
100    print(json.dumps(rep,indent=2))
101if __name__=='__main__': main()