1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
  9from bench.models import rnn_small
 10
 11SEEDS = tuple(range(8))
 12NTRAIN, NTEST, EPOCHS, BATCH = 1200, 400, 12, 128
 13
 14
 15def seed_all(seed):
 16    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 17    if torch.cuda.is_available():
 18        try: torch.cuda.manual_seed_all(seed)
 19        except Exception: pass
 20
 21
 22def baseline_one(cfg, seed):
 23    seed_all(seed)
 24    d = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
 25    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
 26    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'],
 27                               batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None)
 28    return float(metric)
 29
 30
 31class WeakGRU(nn.Module):
 32    """The canonical rnn_small architecture, exposing its latent GRU states."""
 33    def __init__(self, hidden=64):
 34        super().__init__()
 35        self.rnn = nn.GRU(3, hidden, batch_first=True)
 36        self.head = nn.Linear(hidden, 1)
 37
 38    def forward_states(self, x):
 39        return self.rnn(x.view(x.shape[0], -1, 3))[0]
 40
 41    def forward(self, x):
 42        h = self.forward_states(x)
 43        return self.head(h[:, -1])
 44
 45
 46def weak_weights(n=8):
 47    # Uniform quadrature and a Hann test function vanishing at endpoints.
 48    t = torch.arange(n, dtype=torch.float32)
 49    u = t / (n - 1)
 50    psi = torch.sin(np.pi * u) ** 2
 51    dpsi = np.pi * torch.sin(2 * np.pi * u) / (n - 1)
 52    return psi, -dpsi
 53
 54
 55def idea_one(cfg, seed, return_model=False):
 56    seed_all(seed)
 57    d = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
 58    model = WeakGRU(64)
 59    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 60    try:
 61        model = model.to(device)
 62        x, y = d['xtr'].to(device), d['ytr'].to(device)
 63        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 64        psi, wp = weak_weights(x.shape[1] // 3)
 65        psi, wp = psi.to(device), wp.to(device)
 66        # A jointly learned latent generator. This is the only training change.
 67        A = nn.Parameter(torch.zeros(64, 64, device=device))
 68        nn.init.normal_(A, std=0.01)
 69        opt.add_param_group({'params': [A]})
 70        for _ in range(EPOCHS):
 71            model.train()
 72            perm = torch.randperm(len(x), device=device)
 73            for i in range(0, len(x), BATCH):
 74                ix = perm[i:i+BATCH]
 75                states = model.forward_states(x[ix])
 76                pred = model.head(states[:, -1])
 77                task = ((pred - y[ix]) ** 2).mean()
 78                G = (states * psi[None, :, None]).sum(1)
 79                # B_j = - integral psi' z; for one Hann window this is a vector.
 80                Bv = (states * wp[None, :, None]).sum(1)
 81                # A convention: dz/dt = A z, so weak prediction is G A^T.
 82                weak_pred = G @ A.T
 83                weak = ((Bv - weak_pred) ** 2).mean() / (states.detach().var() + 1e-4)
 84                loss = task + cfg['alpha'] * weak
 85                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(list(model.parameters()) + [A], 5.0); opt.step()
 86        model.eval()
 87        with torch.no_grad():
 88            metric = float(((model(d['xte'].to(device)) - d['yte'].to(device)) ** 2).mean())
 89        if return_model: return metric, model, A.detach(), d, device
 90        return metric
 91    except RuntimeError:
 92        # Explicit CPU fallback for a shared or exhausted CUDA slot.
 93        model = WeakGRU(64)
 94        x, y = d['xtr'], d['ytr']; opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 95        psi, wp = weak_weights(x.shape[1] // 3); A = nn.Parameter(torch.randn(64,64)*0.01); opt.add_param_group({'params':[A]})
 96        for _ in range(EPOCHS):
 97            for i in range(0, len(x), BATCH):
 98                states=model.forward_states(x[i:i+BATCH]); pred=model.head(states[:,-1]); task=((pred-y[i:i+BATCH])**2).mean()
 99                G=(states*psi[None,:,None]).sum(1); Bv=(states*wp[None,:,None]).sum(1); weak=((Bv-G@A.T)**2).mean()/(states.detach().var()+1e-4)
100                opt.zero_grad(); (task+cfg['alpha']*weak).backward(); opt.step()
101        with torch.no_grad(): metric=float(((model(d['xte'])-d['yte'])**2).mean())
102        if return_model: return metric, model, A.detach(), d, 'cpu'
103        return metric
104
105
106def signature():
107    cfg={'lr':0.003,'weight_decay':0.0,'alpha':0.03}
108    metric, model, A, d, dev = idea_one(cfg, 0, True)
109    with torch.no_grad():
110        states=model.forward_states(d['xte'].to(dev)); psi,wp=weak_weights(states.shape[1]); psi,wp=psi.to(dev),wp.to(dev)
111        G=(states*psi[None,:,None]).sum(1); Bv=(states*wp[None,:,None]).sum(1); residual=Bv-G@A.T
112        observed=float(residual.pow(2).mean().sqrt()); predicted=float((Bv.detach().var()+1e-8).sqrt())
113    return {'model_test_mse':metric, 'weak_residual_rms_observed':observed,
114            'weak_scale_predicted_from_observed_B':predicted,
115            'relative_residual':observed/(predicted+1e-12),
116            'confirmed': bool(observed < predicted)}
117
118
119def main():
120    # Union parity: every idea learning rate is present in baseline sweep.
121    grid=[{'lr':lr,'weight_decay':wd} for lr in (0.001,0.003,0.01) for wd in (0.0,1e-4)]
122    base=sweep_baseline(lambda c: lambda s: baseline_one(c,s), grid)
123    best=base['best_cfg']
124    idea_grid=[{'lr':best['lr'],'weight_decay':best['weight_decay'],'alpha':a} for a in (0.01,0.03,0.1)]
125    idea_runs=[]
126    for c in idea_grid:
127        r=evaluate(lambda s, c=c: idea_one(c,s), SEEDS)
128        idea_runs.append((c,r))
129    idea_cfg, idea=min(idea_runs, key=lambda cr: cr[1]['mean'])
130    rep=make_report('dynamics','rnn_small',base,idea,{'weak_koopman':signature(), 'selected_cfg':idea_cfg,
131        'baseline_grid':grid, 'idea_grid':idea_grid, 'track_justification':'Actuated pendulum rollout has explicit dynamical stability/control structure.'})
132    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
133    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
134    print(json.dumps(rep,indent=2))
135
136if __name__=='__main__': main()