1import os, sys, json, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import get_dataset, sweep_baseline, evaluate, make_report
  8
  9SEEDS = tuple(range(8))
 10LRS = [1e-3, 3e-3, 1e-2]  # union is used for both methods
 11EPOCHS = 12
 12BATCH = 128
 13SIGMA = 0.20
 14TARGET = 0.92
 15
 16
 17def seed_all(seed):
 18    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 19    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 20
 21
 22def device():
 23    # Probe allocation and cuDNN with the exact recurrent primitive; shared GPU
 24    # failures are handled by the required CPU fallback.
 25    if torch.cuda.is_available():
 26        try:
 27            torch.zeros(1, device='cuda')
 28            probe = nn.GRU(3, 4, batch_first=True).cuda()
 29            probe(torch.zeros(2, 2, 3, device='cuda'))
 30            return 'cuda'
 31        except Exception:
 32            try: torch.cuda.empty_cache()
 33            except Exception: pass
 34    return 'cpu'
 35
 36
 37class RNNSmall(nn.Module):
 38    # Same 32-unit GRU-style recurrent architecture for both systems.
 39    def __init__(self, out_dim=1):
 40        super().__init__()
 41        self.rnn = nn.GRU(3, 32, batch_first=True)
 42        self.head = nn.Linear(32, out_dim)
 43
 44    def forward(self, x, return_hidden=False):
 45        seq = x.view(x.shape[0], -1, 3)
 46        z, h = self.rnn(seq)
 47        if return_hidden:
 48            return self.head(h[-1]), h[-1], z
 49        return self.head(h[-1])
 50
 51
 52def rho_proxy(net):
 53    # For scalar multiplicative uncertainty M_t=(1+sigma xi)M0,
 54    # lifted spectral radius is approximately (1+sigma^2)*rho(M0)^2.
 55    # Use differentiable power iteration on recurrent hidden-hidden blocks.
 56    W = net.rnn.weight_hh_l0
 57    v = torch.ones(W.shape[1], device=W.device) / np.sqrt(W.shape[1])
 58    for _ in range(12):
 59        v = W.T @ (W @ v)
 60        v = v / (v.norm() + 1e-8)
 61    s2 = (W @ v).pow(2).sum()
 62    return (1.0 + SIGMA**2) * s2
 63
 64
 65def train_one(seed, lr, penalty):
 66    seed_all(seed)
 67    ds = get_dataset('dynamics', seed=seed, n_train=400, n_test=400)
 68    dev = device()
 69    net = RNNSmall().to(dev)
 70    x, y = ds['xtr'].to(dev), ds['ytr'].to(dev)
 71    opt = torch.optim.Adam(net.parameters(), lr=lr)
 72    lossf = nn.MSELoss()
 73    for _ in range(EPOCHS):
 74        net.train(); perm = torch.randperm(len(x), device=dev)
 75        for i in range(0, len(x), BATCH):
 76            q = perm[i:i+BATCH]
 77            loss = lossf(net(x[q]), y[q])
 78            if penalty:
 79                loss = loss + penalty * torch.relu(rho_proxy(net) - TARGET).pow(2)
 80            opt.zero_grad(); loss.backward(); opt.step()
 81    net.eval()
 82    with torch.no_grad():
 83        return float(lossf(net(ds['xte'].to(dev)), ds['yte'].to(dev)).cpu())
 84
 85
 86def train_and_signature(seed, lr, penalty):
 87    seed_all(seed)
 88    ds = get_dataset('dynamics', seed=seed, n_train=400, n_test=400)
 89    dev = device(); net = RNNSmall().to(dev)
 90    x, y = ds['xtr'].to(dev), ds['ytr'].to(dev)
 91    opt = torch.optim.Adam(net.parameters(), lr=lr); lossf = nn.MSELoss()
 92    for _ in range(EPOCHS):
 93        perm = torch.randperm(len(x), device=dev)
 94        for i in range(0, len(x), BATCH):
 95            q=perm[i:i+BATCH]; loss=lossf(net(x[q]),y[q])
 96            if penalty: loss = loss + penalty*torch.relu(rho_proxy(net)-TARGET).pow(2)
 97            opt.zero_grad(); loss.backward(); opt.step()
 98    with torch.no_grad():
 99        metric=float(lossf(net(ds['xte'].to(dev)),ds['yte'].to(dev)).cpu())
100        W=net.rnn.weight_hh_l0.detach(); v=torch.ones(32,device=dev)/np.sqrt(32)
101        for _ in range(30): v=W.T@(W@v); v=v/(v.norm()+1e-8)
102        pred=float((1+SIGMA**2)*(W@v).pow(2).sum().cpu())
103        # Re-test the trained recurrent system under multiplicative recurrent noise.
104        h=torch.randn(512,32,device=dev); zero=torch.zeros(512,1,3,device=dev)
105        vals=[float(h.pow(2).mean().cpu())]
106        for _ in range(20):
107            # GRU zero-input rollout, with recurrent weights perturbed as specified.
108            old=net.rnn.weight_hh_l0.data.clone()
109            net.rnn.weight_hh_l0.data = old*(1+SIGMA*torch.randn_like(old))
110            _,hh=net.rnn(zero,h.unsqueeze(0)); h=hh[-1]
111            net.rnn.weight_hh_l0.data = old
112            vals.append(float(h.pow(2).mean().cpu()))
113        ratios=np.asarray(vals[1:])/np.maximum(np.asarray(vals[:-1]),1e-12)
114        obs=float(np.median(ratios[-8:]))
115    return metric, pred, obs
116
117
118def main():
119    # Cheap exact math sanity check: scalar K=(a^2+sigma^2), boundary at 1.
120    a=.8; sig=.6; math_check=float(a*a+sig*sig)
121    grid=[{'lr':lr, 'penalty':0.0} for lr in LRS]
122    base=sweep_baseline(lambda c: lambda s: train_one(s,c['lr'],False), grid, seeds=(0,1,2,3))
123    # Mandatory parity: baseline was evaluated at every idea lr; final baseline is best sweep config.
124    best_lr=float(base['best_cfg']['lr'])
125    idea_grid=[best_lr, 1e-3 if best_lr != 1e-3 else 3e-3, 1e-2 if best_lr != 1e-2 else 3e-3]
126    idea_cfgs=[{'lr':float(lr),'penalty':p} for lr,p in zip(idea_grid,[1.0,3.0,10.0])]
127    # Keep idea sweep size 3; all its lrs are in baseline union.
128    best_idea_cfg=min(idea_cfgs, key=lambda c: np.mean([train_one(s,c['lr'],True) for s in (0,1,2,3)]))
129    base_full=evaluate(lambda s: train_one(s,best_lr,False), SEEDS)
130    idea_full=evaluate(lambda s: train_one(s,best_idea_cfg['lr'],True), SEEDS)
131    sig=[train_and_signature(s,best_idea_cfg['lr'],True) for s in SEEDS]
132    pred=float(np.mean([r[1] for r in sig])); obs=float(np.mean([r[2] for r in sig]))
133    signature={'noise_sigma':SIGMA,'predicted_lifted_rho_mean':pred,
134               'observed_hidden_second_moment_ratio_mean':obs,
135               'relative_error':abs(obs-pred)/max(abs(pred),1e-8),
136               'confirmed': bool(abs(obs-pred)/max(abs(pred),1e-8) < .20),
137               'source':'trained idea GRU models on dynamics test rollout'}
138    rep=make_report('dynamics','rnn_small',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base_full},idea_full,
139                    {'mechanism_signature':signature,'math_sanity':{'a':a,'sigma':sig,'K_scalar':math_check},
140                     'idea_sweep':{'configs':idea_cfgs,'best_cfg':best_idea_cfg}})
141    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
142    print(json.dumps(rep,indent=2))
143
144if __name__=='__main__': main()