1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=467
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9DT=torch.float64
 10
 11# Finite latent space makes the partition function exact (no Monte Carlo ambiguity).
 12xvals=torch.tensor([-2.,-1.,0.,1.,2.],dtype=DT)
 13prior=torch.tensor([.08,.20,.44,.20,.08],dtype=DT); prior/=prior.sum()
 14sigma=.65
 15
 16def make_data(n, seed):
 17    g=torch.Generator().manual_seed(seed)
 18    ix=torch.multinomial(prior,n,replacement=True,generator=g)
 19    x=xvals[ix]
 20    y=x+sigma*torch.randn(n,generator=g,dtype=DT)
 21    return x,y
 22
 23def base_features(x,y):
 24    # frozen nonlinear feature map; feature dimension is deliberately small
 25    z=torch.stack([x, y, x*y, x*x, y*y, torch.sin(y), torch.cos(y)],-1)
 26    W=torch.tensor([[.7,-.2,.4,.1,-.3,.5,.2],[-.4,.8,.2,-.5,.3,.1,-.6],
 27                    [.3,.2,-.7,.6,.1,-.4,.5],[.5,.1,.3,.4,-.6,.2,.7],
 28                    [-.2,.6,.5,.2,.4,-.5,.1], [.4,-.3,.1,-.2,.7,.6,-.4]],dtype=DT)
 29    b=torch.tensor([.1,-.2,.05,.3,-.1,.15],dtype=DT)
 30    return torch.tanh(z@W.T+b)
 31
 32def phi_particles(y):
 33    # [B,K,D], with each observation paired to all prior-supported x
 34    B=y.numel(); K=xvals.numel()
 35    return base_features(xvals[None,:].expand(B,K), y[:,None].expand(B,K)).reshape(B,K,-1)
 36
 37def exact_loss(theta, x, y):
 38    ph=phi_particles(y); target=base_features(x,y)
 39    logits=-torch.einsum('bkd,d->bk',ph,theta)+torch.log(prior)[None,:]
 40    return (torch.einsum('d,bd->b',theta,target)+torch.logsumexp(logits,1)).mean()
 41
 42def posterior(theta,y):
 43    ph=phi_particles(y)
 44    return torch.softmax(-torch.einsum('bkd,d->bk',ph,theta)+torch.log(prior)[None,:],1)
 45
 46class NeuralScore(nn.Module):
 47    def __init__(self,d=6,seed=0):
 48        super().__init__(); torch.manual_seed(seed)
 49        self.net=nn.Sequential(nn.Linear(7,16),nn.Tanh(),nn.Linear(16,1)).double()
 50    def forward(self,x,y):
 51        z=torch.stack([x,y,x*y,x*x,y*y,torch.sin(y),torch.cos(y)],-1)
 52        return self.net(z).squeeze(-1)
 53
 54def neural_loss(model,x,y):
 55    B=y.numel(); xx=xvals[None,:].expand(B,-1); yy=y[:,None].expand(B,xvals.numel())
 56    s=model(xx.reshape(-1),yy.reshape(-1)).reshape(B,-1)
 57    st=model(x,y)
 58    return (st+torch.logsumexp(-s+torch.log(prior)[None,:],1)).mean()
 59
 60def true_post(x_unused,y):
 61    # generative posterior p(x|y) under the known Gaussian observation model
 62    ll=-.5*((y[:,None]-xvals[None,:])/sigma)**2
 63    return torch.softmax(ll+torch.log(prior)[None,:],1)
 64
 65def posterior_kl(pred,y):
 66    t=true_post(None,y); return (t*(torch.log(t+1e-12)-torch.log(pred+1e-12))).sum(1).mean().item()
 67
 68def coverage(pred,y,level=.9):
 69    # shortest contiguous x interval containing requested posterior mass, then test true x
 70    # This is a discrete, transparent calibration proxy.
 71    t=true_post(None,y); covered=[]
 72    for i in range(y.numel()):
 73        p=pred[i]; best=None
 74        for lo in range(5):
 75            for hi in range(lo,5):
 76                if p[lo:hi+1].sum()>=level and (best is None or hi-lo<best[0]): best=(hi-lo,lo,hi)
 77        # expected coverage under true posterior, not sampled latent, is lower variance
 78        covered.append(t[i,best[1]:best[2]+1].sum().item())
 79    return float(np.mean(covered))
 80
 81def convex_fit(x,y,init):
 82    # Newton with covariance Hessian; damping only for numerical stability.
 83    th=init.clone(); history=[]
 84    for it in range(40):
 85        th.requires_grad_(True)
 86        loss=exact_loss(th,x,y); g=torch.autograd.grad(loss,th,create_graph=False)[0]
 87        ph=phi_particles(y); r=posterior(th,y); mu=(r[:,:,None]*ph).sum(1)
 88        centered=ph-mu[:,None,:]
 89        H=torch.einsum('bk,bkd,bke->de',r,centered,centered)/y.numel()
 90        H=H+1e-7*torch.eye(th.numel(),dtype=DT)
 91        try: step=torch.linalg.solve(H,g)
 92        except RuntimeError: step=torch.linalg.lstsq(H,g[:,None]).solution[:,0]
 93        # backtracking preserves objective descent
 94        old=loss.item(); alpha=1.
 95        with torch.no_grad():
 96            while alpha>1e-5 and exact_loss(th-alpha*step,x,y).item()>old-1e-4*alpha*(g@step).item(): alpha*=.5
 97            th=(th-alpha*step).detach()
 98        history.append(old)
 99        if g.norm().item()<2e-8: break
100    return th,history
101
102def run():
103    # Core calculus check at a random parameter and random direction.
104    x,y=make_data(64,11); theta=torch.randn(6,dtype=DT)*.3; theta.requires_grad_()
105    L=exact_loss(theta,x,y); grad=torch.autograd.grad(L,theta,create_graph=True)[0]
106    H=torch.autograd.functional.hessian(lambda t: exact_loss(t,x,y),theta).detach().numpy()
107    eig=np.linalg.eigvalsh((H+H.T)/2)
108    # finite difference directional derivative agrees with Hessian-vector product
109    v=torch.randn(6,dtype=DT); v/=v.norm(); eps=1e-4
110    tp=(theta.detach()+eps*v).requires_grad_(True)
111    tm=(theta.detach()-eps*v).requires_grad_(True)
112    gp=torch.autograd.grad(exact_loss(tp,x,y),tp)[0]
113    gm=torch.autograd.grad(exact_loss(tm,x,y),tm)[0]
114    fd=((gp-gm)/(2*eps)).dot(v).item(); quad=(v@torch.tensor(H)@v).item()
115    check={'hessian_min_eigenvalue':float(eig[0]),'gradient_hvp_fd':fd,'gradient_hvp_formula':quad,'hvp_abs_error':abs(fd-quad)}
116
117    trainx,trainy=make_data(256,21); testx,testy=make_data(512,22)
118    convex=[]; histories=[]
119    for s in [0,1,2,3,4]:
120        th,h=convex_fit(trainx,trainy,torch.randn(6,dtype=DT)*2*(s+1)); histories.append({'start':s,'iterations':len(h),'final_train_objective':h[-1]})
121        pr=posterior(th,testy); convex.append({'objective':exact_loss(th,testx,testy).item(),'posterior_KL':posterior_kl(pr,testy),'coverage90':coverage(pr,testy),'theta':th.tolist()})
122    neural=[]
123    for s in [0,1,2,3,4]:
124        model=NeuralScore(seed=100+s); opt=torch.optim.Adam(model.parameters(),lr=.025)
125        hs=[]
126        for it in range(500):
127            opt.zero_grad(); loss=neural_loss(model,trainx,trainy); loss.backward(); opt.step(); hs.append(loss.item())
128        B=testy.numel(); xx=xvals[None,:].expand(B,-1); yy=testy[:,None].expand(B,xvals.numel())
129        with torch.no_grad(): pr=torch.softmax(-model(xx.reshape(-1),yy.reshape(-1)).reshape(B,-1)+torch.log(prior)[None,:],1)
130        neural.append({'objective':neural_loss(model,testx,testy).item(),'posterior_KL':posterior_kl(pr,testy),'coverage90':coverage(pr,testy),'final_train_objective':hs[-1]})
131    out={'math_check':check,'convex_runs':convex,'convex_histories':histories,'neural_runs':neural,
132         'summary':{'convex_objective_mean_sd':[float(np.mean([z['objective'] for z in convex])),float(np.std([z['objective'] for z in convex]))],
133                    'neural_objective_mean_sd':[float(np.mean([z['objective'] for z in neural])),float(np.std([z['objective'] for z in neural]))],
134                    'convex_KL_mean_sd':[float(np.mean([z['posterior_KL'] for z in convex])),float(np.std([z['posterior_KL'] for z in convex]))],
135                    'neural_KL_mean_sd':[float(np.mean([z['posterior_KL'] for z in neural])),float(np.std([z['posterior_KL'] for z in neural]))],
136                    'convex_coverage_mean_sd':[float(np.mean([z['coverage90'] for z in convex])),float(np.std([z['coverage90'] for z in convex]))],
137                    'neural_coverage_mean_sd':[float(np.mean([z['coverage90'] for z in neural])),float(np.std([z['coverage90'] for z in neural]))]}}
138    with open('results.json','w') as f: json.dump(out,f,indent=2)
139    print(json.dumps(out,indent=2))
140if __name__=='__main__': run()