1import json, sys
  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, train_model, evaluate, sweep_baseline, make_report
  8
  9SEED0 = 2915
 10EPOCHS = 18
 11NTRAIN = 800
 12NTEST = 300
 13BATCH = 128
 14LRS = [1e-3, 3e-3, 1e-2]
 15
 16class MatchedGRU(nn.Module):
 17    """The benchmark rnn_small recurrent architecture, exposed as GRUCell.
 18    Observer differs only by correction between recurrent updates."""
 19    def __init__(self, observer=False, gain=0.0, hidden=64):
 20        super().__init__()
 21        self.cell = nn.GRUCell(3, hidden)
 22        self.head = nn.Linear(hidden, 1)
 23        self.observer = observer
 24        self.gain = float(gain)
 25
 26    def forward(self, x):
 27        z = x.view(x.shape[0], -1, 3)
 28        h = torch.zeros(z.shape[0], self.cell.hidden_size, device=x.device, dtype=x.dtype)
 29        for k in range(z.shape[1]):
 30            # theta is the measured feature; C selects the first latent coordinate.
 31            if self.observer:
 32                innovation = z[:, k, 0] - h[:, 0]
 33                h = h + self.gain * innovation.unsqueeze(1) * torch.eye(1, self.cell.hidden_size, device=x.device, dtype=x.dtype)
 34            h = self.cell(z[:, k], h)
 35        return self.head(h)
 36
 37
 38def make_ds(seed):
 39    return get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
 40
 41def train_one(seed, lr, idea, gain):
 42    torch.manual_seed(seed); np.random.seed(seed)
 43    d = make_ds(seed)
 44    model = MatchedGRU(observer=idea, gain=gain)
 45    _, metric, _ = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 46    return float(metric)
 47
 48def fn(idea, gain, lr):
 49    return lambda seed: train_one(seed, lr, idea, gain)
 50
 51def signature(seed, lr, gain):
 52    """Measure contraction on hidden states of a trained observer model.
 53    Prediction is the continuous certificate upper bound using empirical local
 54    Jacobian one-sided/QIB constants and P=I; observation is actual squared
 55    perturbation ratio after one trained recurrent update."""
 56    torch.manual_seed(seed); np.random.seed(seed)
 57    d = make_ds(seed)
 58    model = MatchedGRU(observer=True, gain=gain)
 59    model, _, _ = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
 60    model.eval(); device = next(model.parameters()).device; x = d['xte'][:128].to(device)
 61    with torch.no_grad():
 62        z = x.view(x.shape[0], -1, 3)
 63        h = torch.zeros(x.shape[0], 64, device=device)
 64        # collect actual hidden states and current observations
 65        for k in range(4):
 66            innovation=z[:,k,0]-h[:,0]
 67            h=h+gain*innovation[:,None]*torch.eye(1,64,device=device)
 68            h=model.cell(z[:,k],h)
 69        base=h[:32].clone(); eps=1e-3
 70        perturb=torch.randn_like(base); perturb=perturb/(perturb.norm(dim=1,keepdim=True)+1e-9)*eps
 71        obs=z[:32,4,0]
 72        def step(q):
 73            qq=q+gain*(obs-q[:,0])[:,None]*torch.eye(1,64,device=device)
 74            return model.cell(z[:32,4],qq)
 75        hn=step(base); hp=step(base+perturb)
 76        ratios=((hp-hn)**2).sum(1)/(perturb**2).sum(1)
 77        observed=float(ratios.mean())
 78        # Empirical local slope bound from finite differences of the trained map.
 79        q=base[:8].detach(); vals=[]; fd=1e-3
 80        for i in range(8):
 81            qi=q[i:i+1]; oi=obs[i:i+1]; zi=z[i:i+1,4]
 82            def local(v):
 83                vv=v+gain*(oi-v[:,0])[:,None]*torch.eye(1,64,device=device)
 84                return model.cell(zi,vv)[0]
 85            cols=[]
 86            for j in range(64):
 87                dv=torch.zeros_like(qi); dv[0,j]=fd
 88                cols.append(((local(qi+dv)-local(qi-dv))/(2*fd)).detach())
 89            vals.append(torch.stack(cols,dim=1))
 90        J=torch.stack(vals)
 91        spectral=float(torch.linalg.matrix_norm(J,ord=2,dim=(1,2)).max())
 92        # empirical local slope bound for the correction+cell map via autograd JVP
 93    predicted=spectral*spectral
 94    return {'gain':gain,'lr':lr,'predicted_V_ratio_bound':predicted,
 95            'observed_mean_V_ratio':observed,'observed_max_V_ratio':float(ratios.max()),
 96            'certificate_contraction_predicted':bool(predicted<1),
 97            'observed_contraction':bool(observed<1),'confirmed':bool(predicted<1 and observed<1)}
 98
 99def main():
100    # Baseline sweep deliberately covers every LR used by the idea.
101    grid=[{'lr':lr,'gain':0.0} for lr in LRS]
102    base=sweep_baseline(lambda c: fn(False, 0.0, c['lr']), grid)
103    # Three observer settings: best baseline LR plus two nearby method settings.
104    idea_grid=[(base['best_cfg']['lr'], 0.25),(base['best_cfg']['lr'],0.5),(base['best_cfg']['lr'],1.0)]
105    # Also ensure all idea lrs have baseline evaluations (already true above).
106    trials=[]
107    for lr,g in idea_grid:
108        r=evaluate(fn(True,g,lr))
109        trials.append({'cfg':{'lr':lr,'gain':g},'result':r})
110    best=min(trials,key=lambda t:t['result']['mean'])
111    idea=best['result']; bestcfg=best['cfg']
112    sig=signature(0,bestcfg['lr'],bestcfg['gain'])
113    report=make_report('dynamics','rnn_small',base,idea,{
114        'description':'trained GRUCell hidden-state observer correction on pendulum windows',
115        'prediction':'local trained recurrent map should contract latent perturbations in V=||e||^2',
116        'trained_model_measurement':sig,
117        'idea_sweep':trials})
118    report['idea_best_cfg']=bestcfg
119    report['protocol_notes']={'track_match':'dynamics: controlled damped pendulum', 'epochs':EPOCHS,
120                              'n_train':NTRAIN,'n_test':NTEST,'lr_union':LRS,
121                              'paired_seeds':8,'baseline_sweep_seeds':4}
122    Path('bench_report.json').write_text(json.dumps(report,indent=2))
123    print(json.dumps(report,indent=2))
124
125if __name__=='__main__': main()