1import sys, copy, json, numpy as np, torch
 2import torch.nn as nn
 3sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
 4import bench
 5from stage2_bench import setup, make_blocks, choose_order, EPOCHS, BATCH
 6
 7# Retest the local prediction on one actually trained benchmark model.
 8def main():
 9    seed=0; cfg={'lr':0.0015,'weight_decay':0.0}
10    ds, model=setup(seed)
11    net, metric, hist=bench.train_model(model,ds,epochs=EPOCHS,lr=cfg['lr'],batch=BATCH,weight_decay=0.0,log=lambda *a,**k:None)
12    dev=next(net.parameters()).device
13    x,y=ds['xtr'].to(dev),ds['ytr'].to(dev)
14    loss_fn=lambda: nn.MSELoss()(net(x),y)
15    bs=make_blocks(net)
16    order,pred_rho,sens,mat=choose_order(net,loss_fn,bs,cfg['lr'])
17    # Apply the same sequential gradient map to two nearby trained states.
18    q=copy.deepcopy(net)
19    qbs=make_blocks(q); q.to(dev)
20    delta=[]
21    with torch.no_grad():
22        for p in q.parameters():
23            z=torch.randn_like(p); p.add_(1e-4*z/max(float(z.norm()),1.0))
24    def sweep(m, groups):
25        for b in order:
26            m.zero_grad(set_to_none=True)
27            l=nn.MSELoss()(m(x),y); l.backward()
28            with torch.no_grad():
29                for p in groups[b]:
30                    if p.grad is not None: p.add_(-cfg['lr']*p.grad)
31    before=torch.cat([p.detach().reshape(-1) for p in net.parameters()])
32    before_q=torch.cat([p.detach().reshape(-1) for p in q.parameters()])
33    sweep(net,bs); sweep(q,qbs)
34    after=torch.cat([p.detach().reshape(-1) for p in net.parameters()])
35    after_q=torch.cat([p.detach().reshape(-1) for p in q.parameters()])
36    observed=float((after_q-after).norm()/(before_q-before).norm())
37    sig={'predicted_rho':float(pred_rho),'observed_local_perturbation_ratio':observed,
38         'relative_gap_pct':float(abs(observed-pred_rho)/max(abs(pred_rho),1e-12)*100),
39         'selected_order':order,'trained_test_mse':float(metric),
40         'confirmed':bool(abs(observed-pred_rho)/max(abs(pred_rho),1e-12)<=0.20)}
41    with open('bench_report.json') as f: rep=json.load(f)
42    rep['mechanism_signature']=sig
43    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
44    print(json.dumps(sig,indent=2))
45if __name__=='__main__': main()