1import sys, json, random, math
  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
  9
 10SEEDS = tuple(range(8))
 11SWEEP_SEEDS = (0, 1, 2, 3)
 12EPOCHS = 8
 13BATCH = 128
 14ALPHA = 0.08
 15MARGIN = 0.01
 16PENALTY_WEIGHT = 0.08
 17
 18# The intervention is a sampled discrete differential-contraction penalty on
 19# the recurrent state transition. The task head and all non-intervention
 20# architecture are exactly bench's rnn_small.
 21def device_and_seed(seed):
 22    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 23    dev = 'cuda' if torch.cuda.is_available() else 'cpu'
 24    try:
 25        if dev == 'cuda': torch.cuda.empty_cache()
 26    except Exception:
 27        dev = 'cpu'
 28    return dev
 29
 30def transition_jacobian(net, x, h):
 31    """Differentiable closed-form Jacobian of one GRU step wrt h."""
 32    c = net.rnn; wi, wh = c.weight_ih_l0, c.weight_hh_l0
 33    gi = x @ wi.T + c.bias_ih_l0; gh = h @ wh.T + c.bias_hh_l0
 34    ir, iz, inn = gi.chunk(3, -1); hr, hz, hnn = gh.chunk(3, -1)
 35    r = torch.sigmoid(ir + hr); z = torch.sigmoid(iz + hz)
 36    n = torch.tanh(inn + r * hnn)
 37    d = h.shape[-1]
 38    Dr = torch.diag_embed(r * (1-r)); Dz = torch.diag_embed(z * (1-z))
 39    Dn = torch.diag_embed(1-n*n); D1mz = torch.diag_embed(1-z)
 40    Dh = torch.diag_embed(h.squeeze(1) - n.squeeze(1))
 41    # PyTorch GRU reset-after-matrix convention: n=tanh(inn+r*(Wh_n h+b)).
 42    Jr = Dr @ wh[:d]
 43    Jz = Dz @ wh[d:2*d]
 44    Jn = Dn @ (wh[2*d:] + torch.diag_embed(hnn.squeeze(1)) @ Jr)
 45    return (D1mz @ Jn - Dh @ Jz + torch.diag_embed(z.squeeze(1)))[0]
 46
 47def recurrent_outputs(net, x):
 48    seq = x.view(x.shape[0], -1, 3)
 49    out, _ = net.rnn(seq)
 50    return seq, out
 51
 52def contraction_penalty(net, x, alpha=ALPHA):
 53    seq, hs = recurrent_outputs(net, x)
 54    vals = []
 55    I = torch.eye(hs.shape[-1], device=x.device)
 56    for b in range(min(2, x.shape[0])):
 57        for t in (seq.shape[1] - 1,):
 58            M = transition_jacobian(net, seq[b:b+1, t], hs[b, t:t+1])
 59            lam = torch.linalg.eigvalsh(M.T @ M - (1.0-alpha) * I)[-1]
 60            vals.append(torch.nn.functional.softplus(lam + MARGIN) ** 2)
 61    return torch.stack(vals).mean()
 62
 63def train_idea(net, d, epochs=EPOCHS, lr=3e-3, seed=0):
 64    dev = next(net.parameters()).device
 65    net.train(); opt = torch.optim.Adam(net.parameters(), lr=lr)
 66    x, y = d['xtr'].to(dev), d['ytr'].to(dev)
 67    gen = torch.Generator(device=dev).manual_seed(seed + 91)
 68    n = len(x)
 69    for _ in range(epochs):
 70        for ix in torch.randperm(n, generator=gen, device=dev).split(BATCH):
 71            xb, yb = x[ix], y[ix]
 72            opt.zero_grad(); pred = net(xb)
 73            task = ((pred - yb) ** 2).mean()
 74            pen = contraction_penalty(net, xb)
 75            (task + PENALTY_WEIGHT * pen).backward()
 76            torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0)
 77            opt.step()
 78    net.eval()
 79    with torch.no_grad(): return float(((net(d['xte'].to(dev))-d['yte'].to(dev))**2).mean().cpu())
 80
 81def run_idea(cfg):
 82    def fn(seed):
 83        dev = device_and_seed(seed)
 84        d = get_dataset('dynamics', seed, n_train=400, n_test=400)
 85        try:
 86            net = make_model('rnn_small', d['input_shape'], d['out_dim']).to(dev)
 87            return train_idea(net, d, lr=cfg['lr'], seed=seed)
 88        except Exception:
 89            if dev != 'cuda': raise
 90            net = make_model('rnn_small', d['input_shape'], d['out_dim']).cpu()
 91            return train_idea(net, {k:(v.cpu() if torch.is_tensor(v) else v) for k,v in d.items()}, lr=cfg['lr'], seed=seed)
 92    return fn
 93
 94def run_base(cfg):
 95    def fn(seed):
 96        device_and_seed(seed)
 97        d = get_dataset('dynamics', seed, n_train=400, n_test=400)
 98        # Required standard path; bench itself supplies GPU->CPU fallback.
 99        _, metric, _ = train_model(make_model('rnn_small', d['input_shape'], d['out_dim']), d, epochs=EPOCHS, lr=cfg['lr'])
100        return float(metric)
101    return fn
102
103def signature(seed, cfg, idea):
104    dev = device_and_seed(seed); d = get_dataset('dynamics', seed, 400, 400)
105    try:
106        net = make_model('rnn_small', d['input_shape'], d['out_dim']).to(dev)
107        if idea: train_idea(net, d, lr=cfg['lr'], seed=seed)
108        else:
109            train_model(net, d, epochs=EPOCHS, lr=cfg['lr'])
110        x = d['xte'][:16].to(dev)
111        with torch.no_grad(): seq, hs = recurrent_outputs(net, x)
112        vals=[]; ratios=[]; I=torch.eye(hs.shape[-1], device=dev)
113        for b in range(4):
114            M=transition_jacobian(net,seq[b:b+1,0],hs[b,0:1])
115            vals.append(float(torch.linalg.eigvalsh(M.T@M-(1-ALPHA)*I)[-1].detach().cpu()))
116            h1=hs[b,0:1].detach(); h2=h1+0.01*torch.ones_like(h1)
117            with torch.no_grad():
118                a=transition_jacobian(net, seq[b:b+1,0], h1) @ (h1-h1).T
119                # finite perturbation through the actual GRU recurrence
120                def step(q):
121                    return recurrent_outputs(net, torch.cat([seq[b:b+1,0:1], seq[b:b+1,0:1]], 1))[1][:,0:1]
122                c=net.rnn; xx=seq[b:b+1,0]
123                def actual(q):
124                    gi=xx@c.weight_ih_l0.T+c.bias_ih_l0; gh=q@c.weight_hh_l0.T+c.bias_hh_l0
125                    ir,iz,inn=gi.chunk(3,-1); hr,hz,hnn=gh.chunk(3,-1)
126                    rr=torch.sigmoid(ir+hr); zz=torch.sigmoid(iz+hz); nnx=torch.tanh(inn+rr*hnn)
127                    return (1-zz)*nnx+zz*q
128                aa,zz=actual(h1),actual(h2)
129            ratios.append(float((torch.linalg.vector_norm(zz-aa)/torch.linalg.vector_norm(h2-h1)).cpu()))
130        return {'predicted_ratio_bound': math.sqrt(1-ALPHA), 'observed_mean_ratio': float(np.mean(ratios)), 'predicted_lmi_bound': 0.0, 'observed_max_lmi': max(vals), 'confirmed': float(np.mean(ratios)) <= math.sqrt(1-ALPHA)+0.03 and max(vals) <= 0.03}
131    except Exception as e:
132        return {'predicted_ratio_bound': math.sqrt(1-ALPHA), 'observed_mean_ratio': None, 'predicted_lmi_bound': 0.0, 'observed_max_lmi': None, 'confirmed': False, 'error': repr(e)}
133
134def main():
135    grid=[{'lr':1e-3},{'lr':3e-3},{'lr':1e-2}]
136    base=sweep_baseline(run_base, grid, seeds=SWEEP_SEEDS)
137    idea_runs=[]
138    for cfg in grid:
139        r=evaluate(run_idea(cfg), seeds=SEEDS)
140        idea_runs.append((r,cfg))
141    best_idea,best_cfg=min(idea_runs, key=lambda z:z[0]['mean'])
142    report=make_report('dynamics','rnn_small',base,best_idea,{'mechanism_signature': signature(0,best_cfg,True), 'idea_sweep':[{'cfg':c,'result':r} for r,c in idea_runs], 'structural_match':'Dynamics track directly tests recurrent trajectory sensitivity and long-horizon stability.'})
143    report['baseline']['union_grid']=grid
144    report['selected_idea_cfg']=best_cfg
145    Path('bench_report.json').write_text(json.dumps(report,indent=2))
146    print(json.dumps(report,indent=2))
147if __name__=='__main__': main()