import sys, os, json, numpy as np, torch
sys.path.insert(0,'.'); sys.path.insert(0,'/home/maxwelhelp/all/math2nn')
from run_bench import seed_all, make_ds, PlainMLP
from bench import train_model
seed_all(0); ds=make_ds(0); net=PlainMLP(enriched=True)
net,_,_=train_model(net,ds,epochs=24,lr=.006,batch=128,log=lambda *_:None)
net.eval()
r=np.logspace(-2.0,-0.7,160)
th=np.pi/4
x=torch.tensor(np.c_[r*np.cos(th),r*np.sin(th)],dtype=torch.float32)
with torch.no_grad(): p=net(x.to(next(net.parameters()).device)).detach().cpu().numpy().ravel()
# Fit away from zeros; use local positive ray where singular mode is nonzero.
mask=np.isfinite(p)&(np.abs(p)>1e-5)
slope=float(np.polyfit(np.log(r[mask]),np.log(np.abs(p[mask])),1)[0])
out=json.load(open('bench_report.json'))
out['mechanism_signature'].update({'observed_nn_loglog_slope':slope,'slope_abs_error':abs(slope-.5),'probe_points':int(mask.sum()),'confirmed':bool(abs(slope-.5)<.15)})
json.dump(out,open('bench_report.json','w'),indent=2)
print(json.dumps(out['mechanism_signature'],indent=2))