1import sys, json, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import make_model, get_dataset, evaluate, sweep_baseline, make_report
  8
  9SEEDS = tuple(range(8))
 10SWEEP_SEEDS = tuple(range(4))
 11GRID = [{"lr": 1e-3}, {"lr": 3e-3}, {"lr": 1e-2}]
 12EPOCHS, BATCH, GAMMA = 25, 64, 0.5
 13
 14def seed_all(s):
 15    random.seed(s); np.random.seed(s); torch.manual_seed(s)
 16    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 17
 18def stein_h(x, gamma=GAMMA):
 19    # Target is N(0,1), score s(x)=-x, and RBF exp(-gamma*(x-y)^2).
 20    dif=x[:,None,:]-x[None,:,:]
 21    r2=(dif*dif).sum(-1)
 22    dot=x@x.T
 23    d=1
 24    return torch.exp(-gamma*r2)*(dot + 2*gamma*d - (2*gamma+4*gamma*gamma)*r2)
 25
 26def objective(out, kind):
 27    h=stein_h(out)
 28    n=out.shape[0]
 29    if kind == 'v':
 30        return torch.sqrt(torch.relu(h.mean()) + 1e-8)
 31    u=(h.sum()-torch.diagonal(h).sum())/(n*(n-1))
 32    return torch.sqrt(torch.relu(u) + 1e-8)
 33
 34def train_system(seed, cfg, kind, return_model=False):
 35    seed_all(seed)
 36    ds=get_dataset('gaussian_score_matching_1d', seed, 400, 400)
 37    net=make_model('mlp_tiny', (1,), 1)
 38    # train_model cannot express the intervention, so this local loop changes
 39    # only the loss while retaining Adam, batch size, epochs, and architecture.
 40    use_cuda=torch.cuda.is_available()
 41    device=torch.device('cuda' if use_cuda else 'cpu')
 42    try:
 43        net.to(device)
 44        x= torch.as_tensor(ds['xtr'], dtype=torch.float32, device=device)
 45        opt=torch.optim.Adam(net.parameters(), lr=cfg['lr'])
 46        for _ in range(EPOCHS):
 47            perm=torch.randperm(len(x), device=device)
 48            for i in range(0,len(x),BATCH):
 49                o=net(x[perm[i:i+BATCH]])
 50                loss=objective(o,kind)
 51                opt.zero_grad(); loss.backward(); opt.step()
 52        with torch.no_grad():
 53            pred=net(torch.as_tensor(ds['xte'],dtype=torch.float32,device=device))
 54            metric=float(((pred-torch.as_tensor(ds['yte'],dtype=torch.float32,device=device))**2).mean().cpu())
 55        return (metric, net.cpu(), ds) if return_model else metric
 56    except Exception:
 57        # Explicit CPU fallback for CUDA/runtime failures.
 58        seed_all(seed); device=torch.device('cpu'); net=make_model('mlp_tiny',(1,),1).to(device)
 59        x=torch.as_tensor(ds['xtr'],dtype=torch.float32)
 60        opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'])
 61        for _ in range(EPOCHS):
 62            perm=torch.randperm(len(x))
 63            for i in range(0,len(x),BATCH):
 64                loss=objective(net(x[perm[i:i+BATCH]]),kind)
 65                opt.zero_grad(); loss.backward(); opt.step()
 66        with torch.no_grad():
 67            p=net(torch.as_tensor(ds['xte'],dtype=torch.float32))
 68            metric=float(((p-torch.as_tensor(ds['yte'],dtype=torch.float32))**2).mean())
 69        return (metric,net,ds) if return_model else metric
 70
 71def train_fn(kind, cfg):
 72    return lambda seed: train_system(seed,cfg,kind)
 73
 74def signature(cfg, kind):
 75    metric, net, ds=train_system(0,cfg,kind,True)
 76    with torch.no_grad(): out=net(torch.as_tensor(ds['xte'],dtype=torch.float32))
 77    h=stein_h(out).numpy(); n=len(out)
 78    diag=np.diag(h); off=(h.sum()-diag.sum())/(n*(n-1)); v=float(np.sqrt(max(h.mean(),0))); u=float(np.sqrt(max(off,0)))
 79    # Bootstrap the trained model's actual outputs, not an analytic toy graph.
 80    rng=np.random.RandomState(395); vs=[]; us=[]
 81    for _ in range(80):
 82        ix=rng.choice(n,64,replace=False); hh=h[np.ix_(ix,ix)]; dd=np.diag(hh)
 83        vs.append(np.sqrt(max(float(hh.mean()),0))); us.append(np.sqrt(max(float((hh.sum()-dd.sum())/(63*64)),0)))
 84    # Empirical covariance spectrum gives the predicted HS/trace fluctuation ratio.
 85    arr=out.numpy().reshape(-1); c=np.atleast_2d(np.cov(arr, rowvar=False)); ev=np.linalg.eigvalsh(c); ev=np.maximum(ev,0)
 86    pred=float(np.sqrt(np.sqrt((ev*ev).sum())/max(ev.sum(),1e-12)))
 87    obs=float(np.std(us,ddof=1)/max(np.std(vs,ddof=1),1e-12))
 88    return {"predicted_hs_over_trace_scale":pred,"observed_bootstrap_sd_U_over_V":obs,
 89            "trained_output_v":v,"trained_output_u":u,"test_mse":metric,
 90            "confirmed": bool(np.isfinite(obs) and abs(np.log(max(obs,1e-12)/max(pred,1e-12)))<1.0)}
 91
 92def main():
 93    base=sweep_baseline(lambda c: train_fn('v',c), GRID, seeds=SWEEP_SEEDS)
 94    # Equal-sized idea sweep on exactly the union of baseline learning rates.
 95    idea_trials=[]
 96    for c in GRID:
 97        r=evaluate(train_fn('u',c), seeds=SWEEP_SEEDS)
 98        idea_trials.append({"cfg":c,"mean":r['mean']})
 99    best=min(idea_trials,key=lambda q:q['mean'])['cfg']
100    idea=evaluate(train_fn('u',best), seeds=SEEDS)
101    sig=signature(best,'u')
102    rep=make_report('gaussian_score_matching_1d','mlp_tiny',base,idea,
103                    {"mechanism_signature":sig,"custom_track":{"name":"gaussian_score_matching_1d","file":"ksd_track_1d.py","domain":"loss"},
104                     "idea_sweep":idea_trials,"training":{"epochs":EPOCHS,"batch":BATCH,"gamma":GAMMA}})
105    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
106    print(json.dumps(rep,indent=2))
107if __name__=='__main__': main()