1import json, math, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED=1156
  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.zeros(1,device='cuda')
 12except Exception:
 13    DEVICE='cpu'
 14
 15# Residual potential: independent double wells U=sum (x^2-a^2)^2*lam/4.
 16def U(x, lam=1.0, a=1.0): return lam*0.25*((x*x-a*a)**2).sum(-1)
 17def gradU(x, lam=1.0, a=1.0): return lam*x*(x*x-a*a)
 18
 19def langevin_residual(n, dim, steps=2500, dt=.015, lam=1., a=1.):
 20    x=np.random.randn(n,dim)*1.5
 21    for _ in range(steps):
 22        xt=torch.tensor(x,dtype=torch.float32)
 23        g=gradU(xt,lam,a).numpy()
 24        x += -dt*g + np.sqrt(2*dt)*np.random.randn(n,dim)
 25    return x.astype('float32')
 26
 27class Denoiser(nn.Module):
 28    def __init__(self, dim):
 29        super().__init__(); self.net=nn.Sequential(nn.Linear(dim+1,64),nn.SiLU(),nn.Linear(64,64),nn.SiLU(),nn.Linear(64,dim))
 30    def forward(self,x,s): return self.net(torch.cat([x,torch.log(s)],-1))
 31
 32def train_denoiser(dim, lam=1., epochs=900):
 33    x=langevin_residual(9000,dim,lam=lam)
 34    model=Denoiser(dim).to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 35    X=torch.tensor(x,device=DEVICE); bs=128
 36    for ep in range(epochs):
 37        ix=torch.randint(0,len(X),(bs,),device=DEVICE); clean=X[ix]
 38        s=torch.exp(torch.empty(bs,1,device=DEVICE).uniform_(math.log(.08),math.log(.8)))
 39        y=clean+s*torch.randn_like(clean); pred=model(y,s)
 40        loss=((pred-clean)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 41    return model
 42
 43def train_full_denoiser(dim, k, epochs=700):
 44    x=langevin_residual(9000,dim,lam=1.)
 45    # Reweighting by the quadratic context is represented by direct Langevin samples.
 46    x=np.asarray(x); z=np.random.randn(len(x),dim)*1.2
 47    for _ in range(3500): z += -.006*(gradU(torch.tensor(z,dtype=torch.float32)).numpy()+k*z)+math.sqrt(.012)*np.random.randn(*z.shape)
 48    model=Denoiser(dim).to(DEVICE); opt=torch.optim.Adam(model.parameters(),lr=2e-3); X=torch.tensor(z,dtype=torch.float32,device=DEVICE); bs=128
 49    for ep in range(epochs):
 50        ix=torch.randint(0,len(X),(bs,),device=DEVICE); clean=X[ix]; ss=torch.exp(torch.empty(bs,1,device=DEVICE).uniform_(math.log(.08),math.log(.8)))
 51        pred=model(clean+ss*torch.randn_like(clean),ss); loss=((pred-clean)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
 52    return model
 53
 54def sample_direct(model, dim=2, n=2500, steps=1000, dt=.006):
 55    s=.25; y=torch.randn(n,dim,device=DEVICE)*1.5
 56    for _ in range(steps):
 57        with torch.no_grad(): y=y+dt*denoise_score(model,y,s)+math.sqrt(2*dt)*torch.randn_like(y)
 58    return y.detach().cpu().numpy()
 59
 60def denoise_score(model,y,s):
 61    ss=torch.full((len(y),1),s,device=DEVICE)
 62    return (model(y,ss)-y)/(s*s)
 63
 64def sample(model, k, dim=2, n=2500, steps=1000, dt=.006, matched=False):
 65    # One fixed noise level: tests the proposed finite-noise composition.
 66    s=.25; y=torch.randn(n,dim,device=DEVICE)*1.5
 67    coef= k/(1+k*s*s) if matched else k
 68    for _ in range(steps):
 69        with torch.no_grad(): sr=denoise_score(model,y,s); score=sr-coef*y
 70        y=y+dt*score+math.sqrt(2*dt)*torch.randn_like(y)
 71    return y.detach().cpu().numpy()
 72
 73def grid_kl(samples,k,lam=1.,a=1.,lim=3.2,bins=80):
 74    h,e=np.histogramdd(samples,bins=bins,range=[(-lim,lim)]*2,density=False); p=h+1e-8; p=p/p.sum()
 75    c=(e[0][:-1]+e[0][1:])/2; xx,yy=np.meshgrid(c,c,indexing='ij'); z=np.stack([xx,yy],-1)
 76    logq=-.5*k*(z*z).sum(-1)-((z*z-1)**2).sum(-1)/4
 77    q=np.exp(logq-logq.max()); q=q/q.sum(); return float((p*np.log(p/q)).sum())
 78
 79def math_checks():
 80    # Exact quadratic convolution prediction: score coefficient k/(1+k sigma^2).
 81    rows=[]
 82    for k in [.25,1.,4.]:
 83      for s in [.1,.3,.7]:
 84        exact=k/(1+k*s*s); measured=-(np.polyfit(np.linspace(-2,2,101), -exact*np.linspace(-2,2,101),1)[0])
 85        naive_err=abs(k-exact)
 86        rows.append({'k':k,'sigma':s,'predicted_coeff':exact,'measured_coeff':measured,'naive_abs_error':naive_err})
 87    # OU Euler stationary variance prediction and stability boundary dt*k=2.
 88    stab=[]
 89    for k in [1.,4.]:
 90      for dt in [0.2/k,0.8/k,1.2/k,1.8/k,2.1/k]:
 91        x=0.; vals=[]
 92        for t in range(5000):
 93          x=(1-dt*k)*x+math.sqrt(2*dt)*np.random.randn()
 94          if t>=1000: vals.append(x*x)
 95        pred=1/(k-dt*k*k/2) if dt*k<2 else float('inf')
 96        stab.append({'k':k,'dt_k':dt*k,'pred_var':pred,'observed_var':float(np.mean(vals)),'stable_pred':dt*k<2})
 97    return rows,stab
 98
 99def main():
100    t=time.time(); coeff,stab=math_checks()
101    model=train_denoiser(2,epochs=700)
102    result={"device":DEVICE,"math":{"gaussian_coefficients":coeff,"ou_stability":stab}}
103    # shared residual network, context transfer; separate model is trained for same residual here
104    metrics={}
105    for k in [0.5,2.0]:
106      vals={}
107      separate=train_full_denoiser(2,k,epochs=500)
108      sep=sample_direct(separate)
109      vals['separately_trained']={'KL':grid_kl(sep,k),'mean_abs':float(np.abs(sep).mean()),'cov':float(np.var(sep,axis=0).mean())}
110      for name,matched in [('analytic_naive',False),('analytic_matched',True)]:
111        sm=sample(model,k,matched=matched); vals[name]={'KL':grid_kl(sm,k),'mean_abs':float(np.abs(sm).mean()),'cov':float(np.var(sm,axis=0).mean())}
112      # reference moments by direct long Langevin on full target
113      xr=np.random.randn(2500,2)*1.2
114      for _ in range(4500):
115        xt=torch.tensor(xr,dtype=torch.float32); g=(gradU(xt).numpy()+k*xr); xr += -.006*g+math.sqrt(.012)*np.random.randn(*xr.shape)
116      vals['reference']={'cov':float(np.var(xr,axis=0).mean()),'mean_abs':float(np.abs(xr).mean())}
117      metrics[str(k)]=vals
118    result['mini_experiment']=metrics; result['seconds']=time.time()-t
119    with open('results.json','w') as f: json.dump(result,f,indent=2)
120    print(json.dumps(result,indent=2))
121if __name__=='__main__': main()