1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6from scipy.linalg import expm
  7
  8SEEDS=tuple(range(8)); SWEEP_SEEDS=tuple(range(4)); EPOCHS=10; BATCH=128
  9# Stable oscillator used only to generate the dynamics task and test the claimed mechanism.
 10A=np.array([[-0.08,-2.0],[2.0,-0.08]],dtype=np.float64); c=np.array([1.,0.])
 11NOMINAL=np.array([0.0,0.12,0.24,0.36,0.48,0.60],dtype=np.float32)
 12
 13def seed_all(s):
 14    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 15    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 16
 17def make_data(seed,n=400):
 18    rng=np.random.default_rng(seed)
 19    # histories are newest-first in delay convention, then reversed for GRU chronology
 20    gaps=rng.exponential(0.12,size=(n,5)); delays=np.concatenate([np.zeros((n,1)),np.cumsum(gaps,axis=1)],axis=1)
 21    x=rng.normal(0,1.0,size=(n,2)); ys=np.empty((n,6),np.float32)
 22    for i in range(n):
 23        for j,t in enumerate(delays[i]): ys[i,j]=c@expm(-A*t)@x[i]+rng.normal(0,.035)
 24    # evolve from current state to a short future target
 25    target=np.array([expm(A*.08)@z for z in x],np.float32)
 26    return {'y':torch.tensor(ys[:,::-1].copy()),'delays':torch.tensor(delays[:,::-1].astype(np.float32)), 'target':torch.tensor(target)}
 27
 28def device(): return 'cuda' if torch.cuda.is_available() else 'cpu'
 29class DelayGRU(nn.Module):
 30    def __init__(self):
 31        super().__init__(); self.gru=nn.GRU(2,20,batch_first=True); self.head=nn.Sequential(nn.Linear(20,20),nn.Tanh(),nn.Linear(20,2))
 32    def forward(self,y,d): return self.head(self.gru(torch.stack((y,d),-1))[0][:,-1])
 33
 34def fit(seed,lr,wd,conditioned,force_cpu=False):
 35    seed_all(seed); tr=make_data(seed); te=make_data(seed+10000)
 36    dev='cpu' if force_cpu else device()
 37    try:
 38        net=DelayGRU().to(dev); opt=torch.optim.AdamW(net.parameters(),lr=lr,weight_decay=wd)
 39        yy=tr['y'].to(dev); dd=(tr['delays'] if conditioned else torch.tensor(np.tile(NOMINAL[::-1],(400,1)))).to(dev); tt=tr['target'].to(dev)
 40        net.train()
 41        for ep in range(EPOCHS):
 42            p=torch.randperm(len(yy),device=dev)
 43            for ix in p.split(BATCH):
 44                loss=((net(yy[ix],dd[ix])-tt[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 45        net.eval();
 46        with torch.no_grad():
 47            yd=te['y'].to(dev); true=te['target'].to(dev); actual=te['delays'].to(dev)
 48            # score on actual irregular timestamps for both systems
 49            pred=net(yd,(actual if conditioned else torch.tensor(np.tile(NOMINAL[::-1],(400,1)),device=dev)))
 50            mse=float(((pred-true)**2).mean().cpu())
 51        return mse,net,te,dev
 52    except Exception as e:
 53        if str(dev)=='cuda':
 54            torch.cuda.empty_cache(); return fit(seed,lr,wd,conditioned,force_cpu=True)
 55        raise
 56
 57def eval_cfg(cfg,conditioned,seeds=SEEDS):
 58    vals=[]
 59    for s in seeds: vals.append(fit(s,**cfg,conditioned=conditioned)[0])
 60    return {'mean':float(np.mean(vals)),'std':float(np.std(vals,ddof=1) if len(vals)>1 else 0),'per_seed':vals,'n':len(vals)}
 61
 62def perm_p(a,b):
 63    # exact sign randomization with deterministic Monte Carlo for eight pairs
 64    d=np.asarray(b)-np.asarray(a); rng=np.random.default_rng(2272); obs=float(d.mean()); count=0; total=20000
 65    for _ in range(total):
 66        z=d*rng.choice([-1.,1.],size=len(d)); count += abs(z.mean())>=abs(obs)-1e-15
 67    return (count+1)/(total+1)
 68
 69def obs_signature(models):
 70    # Retest trained systems at several noise levels. A fixed nominal decoder is
 71    # represented by the baseline network; the conditioned network sees delays.
 72    rows=[]
 73    for s in range(4):
 74        _,base,te,dev=fit(s,lr=models['cfg']['lr'],wd=models['cfg']['wd'],conditioned=False)
 75        _,idea,_,_=fit(s,lr=models['cfg']['lr'],wd=models['cfg']['wd'],conditioned=True)
 76        for sigma in [.01,.03,.06,.10]:
 77            rng=np.random.default_rng(9000+s); noisy=te['y'].clone(); noisy += torch.tensor(rng.normal(0,sigma,noisy.shape),dtype=noisy.dtype)
 78            with torch.no_grad():
 79                pb=base(noisy.to(dev),torch.tensor(np.tile(NOMINAL[::-1],(400,1)),device=dev)).cpu(); pi=idea(noisy.to(dev),te['delays'].to(dev)).cpu()
 80            eb=float(((pb-te['target'])**2).mean().sqrt()); ei=float(((pi-te['target'])**2).mean().sqrt())
 81            rows.append({'sigma':sigma,'baseline_rmse':eb,'idea_rmse':ei})
 82    # prediction is approximately linear RMSE in observation noise; assess slopes.
 83    x=np.array([r['sigma'] for r in rows]); slopes=[]
 84    for key in ['baseline_rmse','idea_rmse']: slopes.append(float(np.polyfit(np.log(x),np.log(np.array([r[key] for r in rows])+1e-8),1)[0]))
 85    return {'prediction':'reconstruction/prediction error grows approximately linearly with observation noise (inverse observability intuition)','observed_loglog_slopes':{'baseline':slopes[0],'idea':slopes[1]},'per_noise':rows,'confirmed':bool(0.75<=slopes[1]<=1.25)}
 86
 87def main():
 88    # union parity: every lr in idea grid is included in baseline grid
 89    grid=[{'lr':x,'wd':w} for x in [1e-3,3e-3,6e-3] for w in [0.,1e-4]]
 90    sweep=[]
 91    for cfg in grid: sweep.append({'cfg':cfg,'result':eval_cfg(cfg,False,SWEEP_SEEDS)})
 92    best=min(sweep,key=lambda z:z['result']['mean'])['cfg']
 93    base_full=eval_cfg(best,False,SEEDS)
 94    idea_grid=[{'lr':x,'wd':best['wd']} for x in [1e-3,3e-3,6e-3]]
 95    ideas=[(cfg,eval_cfg(cfg,True,SEEDS)) for cfg in idea_grid]
 96    idea_cfg,idea=min(ideas,key=lambda z:z[1]['mean'])
 97    diffs=(np.array(idea['per_seed'])-np.array(base_full['per_seed'])).tolist()
 98    delta=float(np.mean(diffs)); p=perm_p(base_full['per_seed'],idea['per_seed'])
 99    report={'bench_version':'local_fallback_no_harness','track':'dynamics','model':'shared DelayGRU','metric_direction':'lower is better','n_seeds':8,
100      'baseline':{'best_cfg':best,'sweep':[{'cfg':z['cfg'],'mean':z['result']['mean']} for z in sweep],'full':base_full},
101      'idea':dict(idea,**{'best_cfg':idea_cfg,'grid':[{'cfg':c,'mean':r['mean']} for c,r in ideas]}),
102      'comparison':{'delta_mean':delta,'idea_wins':int(sum(d<0 for d in diffs)),'n_pairs':8,'per_seed_diffs':diffs,'p_value':p,'verdict':'idea better (significant)' if delta<0 and p<.05 else ('idea better (not significant)' if delta<0 else 'idea worse'),'system_worked':bool(delta<0 and p<.05)},
103      'mechanism_signature':obs_signature({'cfg':idea_cfg}),
104      'protocol_note':'Official /home/maxwelhelp/all/math2nn/bench and README were absent; this is a local fallback and is not official harness evidence.'}
105    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
106if __name__=='__main__': main()