1import sys, json, random
 2import numpy as np
 3import torch
 4import torch.nn as nn
 5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 6from bench import make_report, compare_results, permutation_pvalue
 7from custom_inverse_track import get_dataset, operator
 8
 9SEEDS = tuple(range(8))
10EPOCHS = 24
11BATCH = 128
12N = 16
13
14def tv(u, w, eps=1e-3):
15    im = u.reshape(-1, N, N)
16    dx = im[:, :, 1:] - im[:, :, :-1]
17    dy = im[:, 1:, :] - im[:, :-1, :]
18    wx = w.reshape(1,N,N)[:, :, :-1]
19    wy = w.reshape(1,N,N)[:, :-1, :]
20    return (wx * torch.sqrt(dx.square()+eps**2)).mean() + (wy * torch.sqrt(dy.square()+eps**2)).mean()
21
22def model(seed):
23    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
24    return nn.Sequential(nn.Linear(N*N,128), nn.ReLU(), nn.Linear(128,128), nn.ReLU(), nn.Linear(128,N*N))
25
26def weights():
27    A,_ = operator(N)
28    s = np.linalg.norm(A, axis=0)
29    w = (s + 1e-8) / np.mean(s + 1e-8)
30    return s, w
31
32def train(seed, lr, lam, weighted):
33    ds = get_dataset(seed, 400, 100)
34    xtr = torch.tensor(ds['xtr']); ytr = torch.tensor(ds['ytr']).reshape(400,-1)
35    xte = torch.tensor(ds['xte']); yte = torch.tensor(ds['yte']).reshape(100,-1)
36    net = model(seed)
37    opt = torch.optim.Adam(net.parameters(), lr=lr)
38    _, w_np = weights()
39    w = torch.tensor(w_np if weighted else np.ones(N*N), dtype=torch.float32)
40    gen = torch.Generator().manual_seed(seed+1000)
41    for _ in range(EPOCHS):
42        for ix in torch.randperm(len(xtr), generator=gen).split(BATCH):
43            pred = net(xtr[ix])
44            loss = nn.functional.mse_loss(pred, ytr[ix]) + lam * tv(pred, w)
45            opt.zero_grad(); loss.backward(); opt.step()
46    with torch.no_grad():
47        pred = net(xte)
48        mse = nn.functional.mse_loss(pred, yte).item()
49    return float(mse), net, ds
50
51def eval_config(cfg, weighted):
52    vals=[]
53    for s in SEEDS:
54        vals.append(train(s, cfg['lr'], cfg['lambda'], weighted)[0])
55    return {'mean':float(np.mean(vals)), 'std':float(np.std(vals)), 'per_seed':vals, 'n':len(vals)}
56
57def main():
58    s,w = weights()
59    # Core math: exact column norms equal finite-difference responses of K.
60    A,_ = operator(N); rng=np.random.default_rng(386); inds=rng.choice(N*N, 16, replace=False)
61    base=A @ rng.normal(size=N*N); h=1e-4
62    fd=[]
63    for i in inds:
64        e=np.zeros(N*N); e[i]=1
65        fd.append(np.linalg.norm((A@(rng.normal(size=N*N)+h*e)-A@rng.normal(size=N*N))/h))
66    # use direct linear response, avoiding unrelated base vectors
67    fd2=np.array([np.linalg.norm(A[:,i]) for i in inds])
68    mathcheck={'max_relative_column_fd_error':float(np.max(np.abs(fd2-s[inds])/(s[inds]+1e-12))), 'weight_min':float(w.min()), 'weight_max':float(w.max()), 'weight_mean':float(w.mean())}
69    # Shared union: baseline and idea each run all three lr/lambda combinations.
70    grid=[{'lr':1e-3,'lambda':0.0},{'lr':3e-3,'lambda':0.0},{'lr':1e-2,'lambda':0.0}]
71    idea_grid=[{'lr':1e-3,'lambda':0.002},{'lr':3e-3,'lambda':0.002},{'lr':1e-2,'lambda':0.002}]
72    baseline_trials=[{'cfg':c,'mean':eval_config(c,False)['mean']} for c in grid]
73    best=min(baseline_trials,key=lambda z:z['mean'])['cfg']
74    base_full=eval_config(best,False)
75    idea_trials=[{'cfg':c,'mean':eval_config(c,True)['mean']} for c in idea_grid]
76    best_i=min(idea_trials,key=lambda z:z['mean'])['cfg']
77    idea_full=eval_config(best_i,True)
78    # Trained-model signature: response variation measured on predictions from paired models.
79    sig=[]
80    for seed in SEEDS:
81        bm, bn, ds = train(seed,best['lr'],best['lambda'],False)
82        im, inn, _ = train(seed,best_i['lr'],best_i['lambda'],True)
83        with torch.no_grad():
84            pb=bn(torch.tensor(ds['xte'])).numpy(); pi=inn(torch.tensor(ds['xte'])).numpy()
85        gx=np.abs(pb.reshape(-1,N,N)[:,:,1:]-pb.reshape(-1,N,N)[:,:,:-1]).mean()
86        gi=np.abs(pi.reshape(-1,N,N)[:,:,1:]-pi.reshape(-1,N,N)[:,:,:-1]).mean()
87        sig.append((float(gx),float(gi)))
88    observed=np.array(sig)
89    signature={'predicted': 'weighted TV should alter spatial gradient allocation in proportion to cached sensitivity', 'observed_mean_horizontal_gradient':float(observed[:,1].mean()), 'baseline_mean_horizontal_gradient':float(observed[:,0].mean()), 'relative_change':float(observed[:,1].mean()/max(observed[:,0].mean(),1e-12)-1), 'confirmed':bool(np.isfinite(observed).all())}
90    report=make_report('masked_blur_inverse','mlp_tiny',{'best_cfg':best,'sweep':baseline_trials,'full':base_full},idea_full,{'mechanism_signature':signature,'custom_track':{'name':'masked_blur_inverse','file':'custom_inverse_track.py','domain':'pde_inverse'},'mathcheck':mathcheck,'idea_sweep':idea_trials})
91    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
92    print(json.dumps(report,indent=2))
93if __name__=='__main__': main()