1import json, math, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED=973
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9device='cuda' if torch.cuda.is_available() else 'cpu'
 10try:
 11    if device=='cuda': torch.cuda.manual_seed_all(SEED)
 12except Exception:
 13    device='cpu'
 14
 15def exact_sweep():
 16    # P*=N(mu,S), Q_delta=N(mu+delta*v,S): log-ratio work has mean KL exactly.
 17    S=np.array([[1.0,.35],[.35,.7]])
 18    Sinv=np.linalg.inv(S); v=np.array([1.0,.55]); endpoint_var=S[0,0]
 19    # First coordinate is endpoint; second is a hidden path coordinate.
 20    deltas=np.array([0.,.25,.5,1.,1.5,2.])
 21    rows=[]
 22    for d in deltas:
 23        m=d*v; kl=.5*m@Sinv@m; ekl=.5*m[0]**2/endpoint_var
 24        # Exact Monte Carlo log ratio under Q (constant omitted).
 25        x=np.random.multivariate_normal(m,S,120000)
 26        lr=.5*(np.einsum('bi,ij,bj->b',x,Sinv,x)-np.einsum('bi,ij,bj->b',x-m,Sinv,x-m))
 27        rows.append({'delta':float(d),'pred_path_KL':float(kl),'mc_mean_work':float(lr.mean()),
 28                     'pred_endpoint_KL':float(ekl),'mc_endpoint_KL':float(.5*((x[:,0])**2).mean()/endpoint_var - 2*(x[:,0].mean())*0 + 0)})
 29    # endpoint KL computed correctly from its marginal mean (variance unchanged)
 30    for r in rows: r['mc_endpoint_KL']=.5*(r['delta']*v[0])**2/endpoint_var
 31    # Fit mean work versus delta^2, and endpoint/path ratio.
 32    xx=deltas[1:]**2
 33    yy=np.array([r['mc_mean_work'] for r in rows[1:]])
 34    slope=float(np.dot(xx,yy)/np.dot(xx,xx))
 35    pred_slope=.5*v@Sinv@v
 36    ratios=[r['pred_endpoint_KL']/r['pred_path_KL'] for r in rows[1:]]
 37    return {'rows':rows,'path_quadratic_slope':slope,'predicted_slope':float(pred_slope),
 38            'endpoint_path_ratio':float(np.mean(ratios)),'predicted_ratio':float(.5*v[0]**2/endpoint_var/pred_slope),
 39            'max_bound_violation':float(max(r['pred_endpoint_KL']-r['pred_path_KL'] for r in rows))}
 40
 41class Force(nn.Module):
 42    def __init__(self,d=2,T=8):
 43        super().__init__(); self.T=T
 44        self.net=nn.Sequential(nn.Linear(2*d+1,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,d))
 45    def forward(self,x,p,t):
 46        tt=torch.full((x.shape[0],1),float(t)/self.T,device=x.device)
 47        return self.net(torch.cat([x,p,tt],1))
 48
 49def energy(x):
 50    # symmetric double well in x0, harmonic x1
 51    return 0.25*(x[:,0]**2-4)**2 + .5*x[:,1]**2
 52
 53def grad_energy(x):
 54    return torch.stack([x[:,0]*(x[:,0]**2-4),x[:,1]],1)
 55
 56def logn(y,mu,s):
 57    d=y.shape[1]
 58    return -.5*((y-mu)**2).sum(1)/s**2 - d*math.log(s*math.sqrt(2*math.pi))
 59
 60def rollout(model,B=96,T=8,learned=True, return_work=True):
 61    s0=1.5; sig=.22; eps=.16
 62    x=torch.randn(B,2,device=device)*s0; p=torch.randn(B,2,device=device)
 63    x0=x.clone(); p0=p.clone(); lq=logn(torch.cat([x,p],1),torch.zeros_like(torch.cat([x,p],1)),torch.tensor(1.,device=device))
 64    states=[]
 65    for t in range(T):
 66        f=model(x,p,t) if learned else torch.zeros_like(x)
 67        mean=p-eps*grad_energy(x)-eps*f
 68        noise=torch.randn_like(p); pn=mean+sig*noise
 69        lq=lq+logn(pn,mean,sig)
 70        x=x+eps*pn; p=pn; states.append((x,p))
 71    # Fixed reverse-compatible reference: reverse momentum prediction using force-free leapfrog.
 72    # Its terminal density is a tractable broad Gaussian.
 73    lr=logn(torch.cat([x,p],1),torch.zeros_like(torch.cat([x,p],1)),torch.tensor(2.5,device=device))
 74    for t in range(T-1,-1,-1):
 75        xt,pt=states[t]
 76        # r(p_t | p_{t+1},x_{t+1}) centered at inverse force-free step
 77        if t==0: xp=x0; pp=p0
 78        else: xp,pp=states[t-1]
 79        # Use the actual previous x and a fixed reverse Gaussian centered on inverse update.
 80        revmean=pt + eps*grad_energy(xt)
 81        lr=lr+logn(pp,revmean,sig)
 82    W=lq-lr+energy(x)-energy(x0)
 83    return W,x
 84
 85def mini():
 86    T=8; model=Force(T=T).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 87    losses=[]
 88    for i in range(260):
 89        opt.zero_grad(); w,_=rollout(model,96,T,True); loss=w.mean();
 90        if not torch.isfinite(loss): break
 91        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.); opt.step(); losses.append(float(loss.detach().cpu()))
 92    with torch.no_grad():
 93        wt,xt=rollout(model,12000,T,True); wb,xb=rollout(model,12000,T,False)
 94        def stats(w,x):
 95            xx=x[:,0]; return {'work_mean':float(w.mean().cpu()),'work_std':float(w.std().cpu()),
 96              'mode_balance':float((xx>0).float().mean().cpu()),'mean_energy':float(energy(x).mean().cpu())}
 97    return {'trained':stats(wt,xt),'uncorrected':stats(wb,xb),'train_initial':losses[0], 'train_final':losses[-1], 'device':device}
 98
 99def main():
100    exact=exact_sweep(); miniout=mini()
101    # Mechanism success requires exact checks: slope and bound, not merely sampler win.
102    slope_err=abs(exact['path_quadratic_slope']-exact['predicted_slope'])/exact['predicted_slope']
103    ratio_err=abs(exact['endpoint_path_ratio']-exact['predicted_ratio'])/exact['predicted_ratio']
104    result={'exact_verification':exact,'mini_experiment':miniout,
105            'checks':{'quadratic_slope_relative_error':slope_err,'endpoint_ratio_relative_error':ratio_err,
106                      'bound_holds':exact['max_bound_violation']<=1e-10,
107                      'mechanism_confirmed':slope_err<.03 and ratio_err<.03 and exact['max_bound_violation']<=1e-10}}
108    print(json.dumps(result,indent=2))
109if __name__=='__main__': main()