1import sys, json, random
 2from pathlib import Path
 3import numpy as np
 4import torch
 5from torch import nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import train_model, evaluate, sweep_baseline, make_report, get_dataset
 8from custom_inverse_gain_track import META
 9
10SEEDS=tuple(range(8)); SWEEP=tuple(range(4)); EPOCHS=15
11# Union is used on both sides: baseline and idea see every lr tried.
12GRID=[{'lr':1e-3},{'lr':2e-3},{'lr':3e-3}]
13
14class DirectGRU(nn.Module):
15    def __init__(self):
16        super().__init__(); self.rnn=nn.GRU(3,64,batch_first=True); self.head=nn.Linear(64,1)
17    def forward(self,x):
18        _,h=self.rnn(x.view(x.shape[0],-1,3)); return self.head(h[-1])
19
20class StructuredGRU(nn.Module):
21    def __init__(self, qlo=.5, qhi=4.0):
22        super().__init__(); self.rnn=nn.GRU(3,64,batch_first=True); self.qhead=nn.Linear(64,1)
23        self.qlo=qlo; self.qhi=qhi
24    def forward(self,x, return_q=False):
25        seq=x.view(x.shape[0],-1,3); _,h=self.rnn(seq); q=self.qlo+(self.qhi-self.qlo)*torch.sigmoid(self.qhead(h[-1]))
26        last=seq[:,-1,:]; prev=seq[:,-2,0]
27        delta=(last[:,0]-prev)/.05; z=1.5*last[:,2]-1.5*last[:,0]-delta
28        q=q[:,0]; out=(last[:,1]+q*z).unsqueeze(1)
29        return (out,q,z) if return_q else out
30
31def seed_all(seed):
32    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
33
34def ds_for(seed):
35    d=get_dataset('inverse_gain_control', seed, 400, 400)
36    for k in ('xtr','ytr','xte','yte'): d[k]=torch.as_tensor(d[k],dtype=torch.float32)
37    return d
38
39def run(kind, seed, lr):
40    seed_all(seed+917)
41    d=ds_for(seed)
42    model=DirectGRU() if kind=='baseline' else StructuredGRU()
43    _,metric,_=train_model(model,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
44    return float(metric)
45
46def baseline_fn(cfg):
47    return lambda seed: run('baseline',seed,cfg['lr'])
48def idea_fn(cfg):
49    return lambda seed: run('idea',seed,cfg['lr'])
50
51# Canonical baseline sweep (four seeds), then full paired evaluation.
52base=sweep_baseline(baseline_fn,GRID,seeds=SWEEP)
53# Idea sweep has exactly the same three configurations and seed budget.
54idea_trials=[]
55for cfg in GRID:
56    r=evaluate(idea_fn(cfg),SWEEP)
57    idea_trials.append({'cfg':cfg,'mean':r['mean']})
58best_idea_cfg=min(idea_trials,key=lambda z:z['mean'])['cfg']
59idea_full=evaluate(idea_fn(best_idea_cfg),SEEDS)
60rep=make_report('inverse_gain_control','rnn_small',base,idea_full,extra={})
61rep['custom_track']={'name':META['name'],'file':'custom_inverse_gain_track.py','domain':META['domain']}
62rep['idea_sweep']=idea_trials
63rep['idea_best_cfg']=best_idea_cfg
64rep['protocol_note']='Both systems use identical GRU(3,64), Adam, epochs, batch, data, and the same lr union; only the output parameterization differs.'
65
66# NN-scale mechanism signature from a freshly trained held-out model, not a toy graph.
67# For each test sample, infer observed q from the held-out expert target and measured z.
68seed=0; lr=best_idea_cfg['lr']; seed_all(seed+917); d=ds_for(seed); net=StructuredGRU()
69net,_,_=train_model(net,d,epochs=EPOCHS,lr=lr,batch=128,log=lambda *_:None)
70net=net.cpu(); net.eval()
71with torch.no_grad():
72    pred,q,z=net(d['xte'],return_q=True)
73    pred=pred[:,0]; target=d['yte'][:,0]; qobs=torch.where(z.abs()>1e-5,(target-d['xte'].view(-1,8,3)[:,-1,1])/z,torch.zeros_like(z))
74    lhs=pred-target; rhs=(q-qobs)*z
75    mask=z.abs()>1e-4
76    x=rhs[mask].numpy(); y=lhs[mask].numpy()
77    slope=float(np.dot(x,y)/max(np.dot(x,x),1e-12)); corr=float(np.corrcoef(x,y)[0,1])
78    rel=float(torch.mean(torch.abs(lhs[mask]-rhs[mask]))/(torch.mean(torch.abs(lhs[mask]))+1e-8))
79    q_mae=float(torch.mean(torch.abs(q-qobs)))
80sig={'n_test':int(mask.sum()),'observed_action_error_mean_abs':float(np.mean(np.abs(y))),
81     'predicted_factor_error_mean_abs':float(np.mean(np.abs(x))),'slope_observed_on_predicted':slope,
82     'correlation':corr,'relative_residual':rel,'inverse_gain_mae':q_mae,
83     'prediction':'action_error=(qhat-q_observed)*z','confirmed':bool(abs(slope-1)<0.05 and corr>0.995 and rel<0.05)}
84rep['mechanism_signature']=sig
85Path('bench_report.json').write_text(json.dumps(rep,indent=2))
86print(json.dumps(rep,indent=2))