1import os, sys, json, random
 2import numpy as np
 3import torch
 4from torch import nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
 7
 8SEEDS=tuple(range(8)); SWEEP_SEEDS=(0,1,2,3); EPOCHS=10; BATCH=128
 9LRS=[1e-3,3e-3,1e-2]
10
11def seed_all(s):
12    random.seed(s); np.random.seed(s); torch.manual_seed(s)
13    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
14
15def baseline_run(lr, seed, return_model=False):
16    seed_all(seed); ds=get_dataset('dynamics', seed=seed, n_train=400, n_test=200)
17    net=make_model('rnn_small', ds['input_shape'], ds['out_dim'])
18    net, metric, hist=train_model(net, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *a,**k:None)
19    if return_model: return float(metric), net, ds
20    return float(metric)
21
22class PICopies(nn.Module):
23    def __init__(self, hidden=64, copies=3, kp=0.5, ki=0.05):
24        super().__init__(); self.copies=copies; self.hidden=hidden; self.kp=kp; self.ki=ki
25        self.cells=nn.ModuleList([nn.GRUCell(3,hidden) for _ in range(copies)])
26        self.head=nn.Linear(hidden,1)
27        # deliberate replica-specific constant biases (heterogeneity)
28        self.bias=nn.Parameter(torch.tensor([[0.03],[-0.02],[0.04]]).repeat(1,hidden), requires_grad=False)
29    def forward(self,x, return_states=False):
30        seq=x.view(x.shape[0],-1,3); b=x.shape[0]; dev=x.device
31        hs=[torch.zeros(b,self.hidden,device=dev) for _ in range(self.copies)]
32        z=[torch.zeros_like(hs[0]) for _ in range(self.copies)]
33        states=[]
34        for t in range(seq.shape[1]):
35            raw=[self.cells[i](seq[:,t],hs[i])+self.bias[i] for i in range(self.copies)]
36            mean=torch.stack(raw).mean(0)
37            # Complete-graph zero-row-sum Laplacian output: mean - replica.
38            p=[raw[i]-mean for i in range(self.copies)]
39            mean_p=torch.stack(p).mean(0)
40            z=[z[i]+self.ki*(raw[i]-mean) for i in range(self.copies)]
41            hs=[raw[i]-self.kp*p[i]-z[i] for i in range(self.copies)]
42            states.append(torch.stack(hs))
43        hbar=torch.stack(hs).mean(0)
44        out=self.head(hbar)
45        if return_states: return out, torch.stack(states)
46        return out
47
48def idea_run(lr, seed, kp=0.5, ki=0.05, return_model=False):
49    seed_all(seed); ds=get_dataset('dynamics', seed=seed, n_train=400, n_test=200)
50    net=PICopies(kp=kp,ki=ki)
51    dev='cuda' if torch.cuda.is_available() else 'cpu'
52    try:
53        net=net.to(dev); xtr,ytr=ds['xtr'].to(dev),ds['ytr'].to(dev); xte,yte=ds['xte'].to(dev),ds['yte'].to(dev)
54        opt=torch.optim.Adam(net.parameters(),lr=lr); loss_fn=nn.MSELoss()
55        for _ in range(EPOCHS):
56            net.train(); perm=torch.randperm(len(xtr),device=dev)
57            for ix in perm.split(BATCH):
58                opt.zero_grad(); loss=loss_fn(net(xtr[ix]),ytr[ix]); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
59        net.eval()
60        with torch.no_grad(): metric=float(loss_fn(net(xte),yte).item())
61    except Exception:
62        net=net.cpu(); xtr,ytr=ds['xtr'],ds['ytr']; xte,yte=ds['xte'],ds['yte']
63        opt=torch.optim.Adam(net.parameters(),lr=lr); loss_fn=nn.MSELoss()
64        for _ in range(EPOCHS):
65            for ix in torch.randperm(len(xtr)).split(BATCH):
66                opt.zero_grad(); loss=loss_fn(net(xtr[ix]),ytr[ix]); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
67        with torch.no_grad(): metric=float(loss_fn(net(xte),yte).item())
68    if return_model:return metric,net,ds
69    return metric
70
71def signature():
72    # Measured on trained benchmark models, not the analytic toy system.
73    vals=[]
74    for s in (0,1,2,3):
75        m,net,ds=idea_run(3e-3,s,return_model=True)
76        dev=next(net.parameters()).device
77        with torch.no_grad():
78            _,st=net(ds['xte'].to(dev),return_states=True)
79            dis=float(st.std(dim=1).mean().cpu())
80        vals.append(dis)
81    obs=float(np.mean(vals))
82    # PI's predicted qualitative signature is contraction of transverse replicas.
83    return {'prediction':'integral coupling should yield small transverse hidden disagreement in trained heterogeneous copies',
84            'observed_mean_pairwise_hidden_std':obs,'per_seed':vals,
85            'threshold':0.20,'confirmed':bool(np.isfinite(obs) and obs<0.20)}
86
87def main():
88    grid=[{'lr':v} for v in LRS]
89    base=sweep_baseline(lambda c: lambda s: baseline_run(c['lr'],s),grid,seeds=SWEEP_SEEDS)
90    trials=[{'cfg':{'lr':lr},'result':evaluate(lambda s,lr=lr: idea_run(lr,s),SEEDS)} for lr in LRS]
91    best=min(trials,key=lambda q:q['result']['mean'])
92    rep=make_report('dynamics','rnn_small',base,best['result'],{
93      'idea_config':best['cfg'],'idea_sweep':trials,
94      'mechanism_signature':signature(),
95      'protocol_note':'baseline sweep and idea sweep use identical learning-rate union; all task metrics are test MSE'})
96    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
97    print(json.dumps(rep,indent=2))
98if __name__=='__main__': main()