1import sys, json, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7_LAST_NET = None
  8_LAST_X = None
  9sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 10from bench import get_dataset, make_model, evaluate, sweep_baseline, make_report
 11
 12# Pullback monitor: replay the same additive latent perturbation for K replicas.
 13def pairwise_rms(z):
 14    d = z[:, None, :] - z[None, :, :]
 15    return torch.sqrt(torch.mean(torch.sum(d*d, dim=-1)) + 1e-12)
 16
 17def pullback_penalty(net, x, steps=4, replicas=8, noise_scale=0.03, target_lambda=-0.08):
 18    # Use the trained GRU's recurrent cell on a short observed dynamics prefix.
 19    seq = x.view(x.shape[0], -1, 3)[:1]
 20    with torch.no_grad():
 21        _, h0 = net.rnn(seq)
 22    h = h0[-1].detach().repeat(replicas, 1)
 23    init = torch.randn_like(h)
 24    init = init / (init.norm(dim=1, keepdim=True) + 1e-8) * 0.5
 25    h = h + init
 26    # Same noise vector at every step for all replicas; noise is diagnostic-only.
 27    ds = []
 28    for r in range(steps):
 29        inp = torch.zeros(replicas, 3, device=h.device)
 30        h = net.rnn(inp.unsqueeze(1), h.unsqueeze(0))[1][-1]
 31        if noise_scale:
 32            eps = torch.randn(1, h.shape[1], device=h.device) * noise_scale
 33            h = h + eps
 34        ds.append(pairwise_rms(h))
 35    logs = torch.log(torch.stack(ds))
 36    # Penalize growth beyond the desired contraction envelope.
 37    t = torch.arange(1, steps + 1, device=h.device, dtype=h.dtype)
 38    excess = torch.relu(logs[1:] - logs[:-1] - target_lambda)
 39    return torch.mean(excess**2), float((logs[-1]-logs[0]).detach().cpu())
 40
 41def train_one(seed, cfg, idea):
 42    global _LAST_NET, _LAST_X
 43    torch.manual_seed(seed); np.random.seed(seed)
 44    d = get_dataset('dynamics', seed, n_train=400, n_test=160)
 45    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
 46    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 47    try:
 48        net.to(device)
 49        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
 50        lossf = nn.MSELoss(); xtr,ytr=d['xtr'].to(device),d['ytr'].to(device)
 51        for ep in range(cfg['epochs']):
 52            net.train(); perm=torch.randperm(len(xtr),device=device)
 53            for i in range(0,len(xtr),128):
 54                ix=perm[i:i+128]; pred=net(xtr[ix]); loss=lossf(pred,ytr[ix])
 55                if idea:
 56                    reg,_=pullback_penalty(net,xtr[ix],target_lambda=cfg['target_lambda'])
 57                    loss=loss+cfg['attr_weight']*reg
 58                opt.zero_grad(); loss.backward(); opt.step()
 59        net.eval()
 60        with torch.no_grad(): metric=float(lossf(net(d['xte'].to(device)),d['yte'].to(device)).cpu())
 61        _LAST_NET, _LAST_X = net, d['xte'][:1].to(device)
 62        return metric
 63    except RuntimeError:
 64        # CPU fallback, matching the harness safety requirement.
 65        net=net.cpu(); opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
 66        xtr,ytr=d['xtr'],d['ytr']
 67        for ep in range(cfg['epochs']):
 68            perm=torch.randperm(len(xtr))
 69            for i in range(0,len(xtr),128):
 70                ix=perm[i:i+128]; loss=lossf(net(xtr[ix]),ytr[ix])
 71                if idea: loss=loss+cfg['attr_weight']*pullback_penalty(net,xtr[ix],target_lambda=cfg['target_lambda'])[0]
 72                opt.zero_grad(); loss.backward(); opt.step()
 73        with torch.no_grad(): metric=float(lossf(net(d['xte']),d['yte']))
 74        _LAST_NET, _LAST_X = net, d['xte'][:1]
 75        return metric
 76
 77def signature(seed=0, cfg=None):
 78    # Uses the final trained model retained by evaluate(), not an analytic toy.
 79    global _LAST_NET, _LAST_X
 80    if _LAST_NET is None:
 81        return {'error': 'no trained model', 'confirmed': False}
 82    net, x = _LAST_NET, _LAST_X
 83    vals=[]
 84    old = torch.backends.cudnn.enabled
 85    try:
 86        torch.backends.cudnn.enabled = False
 87        net.eval()
 88        with torch.no_grad():
 89            seq=x.view(1,-1,3); _,h0=net.rnn(seq); h=h0[-1].repeat(16,1)
 90            q=torch.randn_like(h); q=q/(q.norm(dim=1,keepdim=True)+1e-8)*.5; h+=q
 91            for _ in range(6):
 92                inp=torch.zeros(16,1,3,device=h.device)
 93                h=net.rnn(inp,h.unsqueeze(0))[1][-1]
 94                vals.append(float(pairwise_rms(h).cpu()))
 95    except RuntimeError:
 96        net=net.cpu(); x=x.cpu(); torch.backends.cudnn.enabled=False
 97        with torch.no_grad():
 98            seq=x.view(1,-1,3); _,h0=net.rnn(seq); h=h0[-1].repeat(16,1)
 99            q=torch.randn_like(h); q=q/(q.norm(dim=1,keepdim=True)+1e-8)*.5; h+=q
100            for _ in range(6):
101                h=net.rnn(torch.zeros(16,1,3),h.unsqueeze(0))[1][-1]
102                vals.append(float(pairwise_rms(h)))
103    finally:
104        torch.backends.cudnn.enabled=old
105    observed=float(np.polyfit(np.arange(6),np.log(np.maximum(vals,1e-12)),1)[0])
106    return {'predicted_log_diameter_step': -0.08, 'observed_log_diameter_slope': observed,
107            'final_diameter': vals[-1], 'confirmed': bool(observed < 0 and abs(observed+0.08)<0.15)}
108
109def main():
110    # Union parity: every idea lr is also in baseline sweep.
111    grid=[{'lr':lr,'weight_decay':wd,'epochs':12,'attr_weight':0.0,'target_lambda':-0.08} for lr in (0.0015,0.003,0.006) for wd in (0.0,1e-4)]
112    base=sweep_baseline(lambda c: lambda s: train_one(s,c,False),grid)
113    idea_cfgs=[dict(c,attr_weight=w,target_lambda=t) for c in [base['best_cfg']] for w,t in ((0.001,-0.08),(0.01,-0.08),(0.03,-0.04))]
114    idea_runs=[]
115    for c in idea_cfgs:
116        idea_runs.append((c,evaluate(lambda s,c=c: train_one(s,c,True))))
117    best_cfg,best=min(idea_runs,key=lambda z:z[1]['mean'])
118    rep=make_report('dynamics','rnn_small',base,best,{'prediction':'trained hidden pullback diameter contracts exponentially under shared replay','config':best_cfg,**signature(0,best_cfg)})
119    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
120    rep['track_justification']='Dynamics is structurally matched: the idea diagnoses stability and contraction of recurrent latent trajectories.'
121    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
122    print(json.dumps(rep,indent=2))
123if __name__=='__main__': main()