1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6
  7SEED=17
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9try:
 10    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11    if device.type=='cuda': torch.zeros(1,device=device)
 12except Exception:
 13    device=torch.device('cpu')
 14
 15torch.set_default_dtype(torch.float32)
 16
 17class Coupling(nn.Module):
 18    def __init__(self,d,mask,cond=2,hidden=32,clamp=1.5):
 19        super().__init__(); self.d=d; self.register_buffer('mask',mask.float()); self.clamp=clamp
 20        inp=d+cond+1
 21        self.net=nn.Sequential(nn.Linear(inp,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,2*d))
 22        # near identity initialization, but trainable
 23        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
 24    def params_at(self,x,c,t):
 25        u=torch.cat([x*self.mask, c, t],-1); out=self.net(u)
 26        raw_s,b=out[...,:self.d],out[...,self.d:]
 27        s=self.clamp*torch.tanh(raw_s)* (1-self.mask)
 28        b=b*(1-self.mask)
 29        return s,b
 30    def forward(self,x,c,t):
 31        s,b=self.params_at(x,c,t); y=x*self.mask+(1-self.mask)*(x*torch.exp(s)+b)
 32        return y, s.sum(-1)
 33    def inverse(self,y,c,t):
 34        s,b=self.params_at(y,c,t); x=y*self.mask+(1-self.mask)*(y-b)*torch.exp(-s)
 35        return x, -s.sum(-1)
 36
 37class Flow(nn.Module):
 38    def __init__(self,d,K=4,cond=2):
 39        super().__init__(); masks=[]
 40        for k in range(K): masks.append(torch.tensor([(i+k)%2 for i in range(d)]))
 41        self.layers=nn.ModuleList([Coupling(d,m,cond) for m in masks]); self.d=d
 42    def forward(self,z,c,t):
 43        h=z; ld=torch.zeros(z.shape[0],device=z.device)
 44        for layer in self.layers:
 45            h,a=layer(h,c,t); ld=ld+a
 46        return h,ld
 47    def inverse(self,x,c,t):
 48        h=x; ld=torch.zeros(x.shape[0],device=x.device)
 49        for layer in reversed(self.layers):
 50            h,a=layer.inverse(h,c,t); ld=ld+a
 51        return h,ld
 52    def log_prob(self,x,c,t):
 53        z, invld=self.inverse(x,c,t)
 54        logq=-0.5*(z*z+math.log(2*math.pi)).sum(-1)
 55        return logq+invld
 56    def score(self,x,c,t):
 57        xx=x.detach().requires_grad_(True); lp=self.log_prob(xx,c,t).sum()
 58        return torch.autograd.grad(lp,xx)[0]
 59
 60class DiagMLP(nn.Module):
 61    def __init__(self,d,cond=2):
 62        super().__init__(); self.d=d; self.net=nn.Sequential(nn.Linear(cond+1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,2*d))
 63    def log_prob(self,x,c,t):
 64        o=self.net(torch.cat([c,t],-1)); mu,raw=o[...,:self.d],o[...,self.d:]
 65        logstd=2*torch.tanh(raw); return (-0.5*((x-mu)/logstd.exp())**2-logstd-0.5*math.log(2*math.pi)).sum(-1)
 66    def score(self,x,c,t):
 67        xx=x.detach().requires_grad_(True); lp=self.log_prob(xx,c,t).sum(); return torch.autograd.grad(lp,xx)[0]
 68
 69def sample_data(n,d,device):
 70    # task c=(separation, parity feature); conditional symmetric two-component mixture
 71    sep=torch.rand(n,1,device=device)*1.5+1.5; aux=torch.sin(sep)
 72    c=torch.cat([sep,aux],-1); t=torch.zeros(n,1,device=device)
 73    sign=torch.where(torch.rand(n,1,device=device)<.5,-1.,1.)
 74    direction=torch.ones(1,d,device=device); direction[:,1::2]=-1
 75    means=sign*sep*direction
 76    x=means+0.35*torch.randn(n,d,device=device)
 77    return x,c,t
 78
 79def analytic_score(x,c):
 80    sep=c[:,0:1]; direction=torch.ones(1,x.shape[1],device=x.device); direction[:,1::2]=-1
 81    m=sep*direction; var=.35**2
 82    lp=torch.stack([-0.5*((x-mm)**2).sum(-1)/var for mm in [m,-m]],-1)
 83    w=torch.softmax(lp,-1)
 84    scores=torch.stack([-(x-mm)/var for mm in [m,-m]],1)
 85    return (w.unsqueeze(-1)*scores).sum(1)
 86
 87def train(model,d,steps=700):
 88    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
 89    for i in range(steps):
 90        x,c,t=sample_data(96,d,device); loss=-model.log_prob(x,c,t).mean(); opt.zero_grad(); loss.backward(); opt.step()
 91    with torch.no_grad():
 92        x,c,t=sample_data(512,d,device); nll=-model.log_prob(x,c,t).mean().item()
 93    # score needs graph, use fixed batch
 94    x,c,t=sample_data(256,d,device); pred=model.score(x,c,t); truth=analytic_score(x,c)
 95    return nll, (pred-truth).pow(2).mean().sqrt().item()
 96
 97def math_checks():
 98    d=8; x=torch.randn(64,d,device=device); c=torch.randn(64,2,device=device); t=torch.randn(64,1,device=device)
 99    rows=[]
100    # constant alpha is imposed through a tiny custom coupling: test the exact formulas directly
101    for gamma in [0.,.25,.5,1.,1.5]:
102        mask=torch.tensor([1,0,1,0,1,0,1,0],device=device).float(); nB=int((1-mask).sum())
103        alpha=gamma*torch.ones_like(x)*(1-mask); beta=.13*torch.ones_like(x)*(1-mask)
104        y=x*mask+(1-mask)*(x*torch.exp(alpha)+beta)
105        xr=y*mask+(1-mask)*(y-beta)*torch.exp(-alpha)
106        recon=(xr-x).abs().max().item(); observed=y.shape[0]
107        # logdet is nB*gamma; additive K=3 prediction is 3*nB*gamma
108        one=alpha.sum(-1).mean().item(); comp=(3*alpha.sum(-1)).mean().item()
109        rows.append({'gamma':gamma,'pred_logdet':nB*gamma,'obs_logdet':one,'pred_3x':3*nB*gamma,'obs_3x':comp,'recon_max':recon})
110    # Direct Jacobian check on the implemented nonlinear coupling (single sample).
111    layer=Coupling(d, torch.tensor([1,0,1,0,1,0,1,0],device=device).float(), cond=2).to(device)
112    with torch.no_grad():
113        layer.net[-1].bias[:d].fill_(0.35)
114        layer.net[-1].bias[d:].fill_(0.11)
115    xx=torch.randn(d,device=device,requires_grad=True); cc=torch.randn(1,2,device=device); tt=torch.randn(1,1,device=device)
116    yy,ld=layer(xx[None],cc,tt); xr,ild=layer.inverse(yy,cc,tt)
117    J=torch.autograd.functional.jacobian(lambda u: layer(u[None],cc,tt)[0][0],xx)
118    sign,actual=torch.linalg.slogdet(J)
119    rows.append({'nonlinear_jacobian_logdet':actual.item(),'analytic_logdet':ld.item(),
120                 'jacobian_abs_error':abs(actual.item()-ld.item()),'jacobian_sign':sign.item(),
121                 'nonlinear_inverse_max_error':(xr[0]-xx).abs().max().item()})
122    return rows
123
124def main():
125    checks=math_checks(); results={'device':str(device),'math_checks':checks,'models':{}}
126    for d in [2,8]:
127        # reset reproducibly per dimension
128        torch.manual_seed(SEED+d)
129        results['models'][str(d)]={}
130        results['models'][str(d)]['flow']=train(Flow(d),d)
131        torch.manual_seed(SEED+d)
132        results['models'][str(d)]['diag_mlp']=train(DiagMLP(d),d)
133    Path('results.json').write_text(json.dumps(results,indent=2))
134    print(json.dumps(results,indent=2))
135
136if __name__=='__main__': main()