1import sys, json, math
  2import numpy as np
  3import torch
  4from torch import nn
  5sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  6from bench import get_dataset as bench_get_dataset, evaluate, sweep_baseline, make_report
  7
  8DEVICE='cuda' if torch.cuda.is_available() else 'cpu'
  9try: torch.set_num_threads(4)
 10except Exception: pass
 11
 12# Registered bench track; its target is a bounded 2-D energy surface.
 13TRACK='bounded_energy_regression'
 14K=64
 15side=8
 16axis=torch.linspace(-1.5,1.5,side)
 17grid=torch.cartesian_prod(axis,axis)
 18LOGPRIOR=torch.full((K,),-math.log(K))
 19
 20def seed_all(s):
 21    np.random.seed(s); torch.manual_seed(s)
 22    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)
 23
 24def data(seed):
 25    return bench_get_dataset(TRACK,seed,n_train=400,n_test=400)
 26
 27def phi(latent, obs):
 28    # Frozen neural-style features, shared by both separately trained systems.
 29    a,b=latent[...,0],latent[...,1]
 30    o=obs[...,0] if obs.shape[-1] else obs
 31    z=torch.stack((a,b,o,a*o,b*o,a*a,b*b,o*o,torch.sin(a),torch.cos(b)),-1)
 32    W=torch.tensor([[.31,-.22,.18,.11,-.14,.25,.07,-.09,.13,.21],
 33                    [-.12,.28,.24,-.19,.17,.08,-.16,.12,.22,-.11],
 34                    [.21,.07,-.31,.26,.09,-.12,.19,.15,-.18,.17],
 35                    [.16,.19,.05,.22,-.27,.13,.11,-.21,.14,.09],
 36                    [-.24,.11,.16,.08,.21,-.17,.23,.18,.07,-.15],
 37                    [.09,-.18,.27,.14,.12,.22,-.08,.16,-.11,.24]],device=latent.device)
 38    bias=torch.tensor([.1,-.2,.05,.3,-.1,.15],device=latent.device)
 39    return torch.tanh(z@W.t()+bias)
 40
 41class NeuralPotential(nn.Module):
 42    def __init__(self,hidden=32):
 43        super().__init__(); self.net=nn.Sequential(nn.Linear(10,hidden),nn.Tanh(),nn.Linear(hidden,1))
 44    def forward(self,latent,obs):
 45        a,b=latent[...,0],latent[...,1]; o=obs[...,0] if obs.shape[-1] else obs
 46        z=torch.stack((a,b,o,a*o,b*o,a*a,b*b,o*o,torch.sin(a),torch.cos(b)),-1)
 47        return self.net(z).squeeze(-1)
 48
 49def convex_loss(theta,latent_obs,latent_true,obs,ret=False):
 50    # observation is scalar target energy; latent_true is the coordinate pair.
 51    b=obs.shape[0]; particles=grid.to(obs.device)[None].expand(b,-1,-1); oo=obs[:,None].expand(b,K,1)
 52    p=phi(particles,oo); target=phi(latent_true,obs)
 53    logits=-torch.einsum('bkd,d->bk',p,theta)+LOGPRIOR.to(obs.device)
 54    loss=(torch.einsum('bd,d->b',target,theta)+torch.logsumexp(logits,1)).mean()
 55    return (loss,p,torch.softmax(logits,1)) if ret else loss
 56
 57def neural_loss(net,latent_true,obs,ret=False):
 58    b=obs.shape[0]; particles=grid.to(obs.device)[None].expand(b,-1,-1); oo=obs[:,None].expand(b,K,1)
 59    logits=-net(particles.reshape(-1,2),oo.reshape(-1,1)).reshape(b,K)+LOGPRIOR.to(obs.device)
 60    loss=(net(latent_true,obs)+torch.logsumexp(logits,1)).mean()
 61    return (loss,logits,torch.softmax(logits,1)) if ret else loss
 62
 63def posterior_mean(model,obs,idea):
 64    b=obs.shape[0]; p=grid.to(obs.device)[None].expand(b,-1,-1); oo=obs[:,None].expand(b,K,1)
 65    if idea: logits=-torch.einsum('bkd,d->bk',phi(p,oo),model)+LOGPRIOR.to(obs.device)
 66    else: logits=-model(p.reshape(-1,2),oo.reshape(-1,1)).reshape(b,K)+LOGPRIOR.to(obs.device)
 67    return (torch.softmax(logits,1)[...,None]*p).sum(1)
 68
 69def train_one(seed,lr,idea,epochs=28,return_model=False):
 70    seed_all(seed); d=data(seed)
 71    # Bench x is latent coordinate and y is observed energy.
 72    latent=torch.as_tensor(d['xtr'],device=DEVICE); obs=torch.as_tensor(d['ytr'],device=DEVICE)
 73    if idea:
 74        model=nn.Parameter(torch.randn(6,device=DEVICE)*.15); opt=torch.optim.Adam([model],lr=lr)
 75    else:
 76        model=NeuralPotential().to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=lr)
 77    n=len(obs)
 78    for _ in range(epochs):
 79        order=torch.randperm(n,device=DEVICE)
 80        for j in range(0,n,64):
 81            ix=order[j:j+64]; opt.zero_grad(set_to_none=True)
 82            loss=convex_loss(model,None,latent[ix],obs[ix]) if idea else neural_loss(model,latent[ix],obs[ix])
 83            loss.backward(); opt.step()
 84    lt=torch.as_tensor(d['xte'],device=DEVICE); ot=torch.as_tensor(d['yte'],device=DEVICE)
 85    with torch.no_grad(): metric=((posterior_mean(model,ot,idea)-lt)**2).mean().item()
 86    return (metric,model,d) if return_model else metric
 87
 88GRID_CFG=[{'lr':.01},{'lr':.03},{'lr':.08}]
 89def base_fn(cfg): return lambda s: train_one(s,cfg['lr'],False)
 90def idea_fn(cfg): return lambda s: train_one(s,cfg['lr'],True)
 91
 92def signature(seed,lr):
 93    _,_,d=train_one(seed,lr,False,return_model=True); _,theta,_=train_one(seed,lr,True,return_model=True)
 94    obs=torch.as_tensor(d['yte'][:96],device=DEVICE); true=torch.as_tensor(d['xte'][:96],device=DEVICE)
 95    with torch.no_grad():
 96        _,p,r=convex_loss(theta,None,true,obs,True)
 97        target=phi(true,obs); grad=(target-(r[...,None]*p).sum(1)).mean(0)
 98        mu=(r[...,None]*p).sum(1); c=p-mu[:,None,:]
 99        H=torch.einsum('bk,bkd,bke->de',r,c,c)/len(obs)
100        mineig=float(torch.linalg.eigvalsh((H+H.T)/2)[0])
101    return {'prediction':'trained convex posterior head has covariance Hessian PSD',
102            'observed_gradient_norm':float(grad.norm()),
103            'observed_covariance_min_eigenvalue':mineig,
104            'confirmed':bool(mineig>=-1e-5)}
105
106def main():
107    base=sweep_baseline(base_fn,GRID_CFG)
108    trials=[{'cfg':c,'result':evaluate(idea_fn(c))} for c in GRID_CFG]
109    best=min(trials,key=lambda z:z['result']['mean'])
110    rep=make_report(TRACK,'energy_potential_mlp',base,best['result'],{
111        'idea_sweep':trials,
112        'structural_match':'Registered bounded_energy_regression: 2-D energy surface; latent coordinates are normalized particles conditioned on observed energy.',
113        'mechanism_signature':signature(0,best['cfg']['lr'])})
114    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
115    print(json.dumps(rep,indent=2))
116if __name__=='__main__': main()