1import json, time, math
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5from scipy.stats import wasserstein_distance
  6
  7torch.set_default_dtype(torch.float64)
  8SEED=7
  9np.random.seed(SEED); torch.manual_seed(SEED)
 10device='cuda' if torch.cuda.is_available() else 'cpu'
 11try:
 12    if device=='cuda': torch.zeros(1,device='cuda')
 13except Exception:
 14    device='cpu'
 15
 16# Periodic spectral grid; the distributions are negligible at the boundaries.
 17L=8.0; n=256; dx=2*L/n
 18x=torch.linspace(-L,L-dx,n,device=device)
 19k=2*math.pi*torch.fft.fftfreq(n,d=dx).to(device)
 20hbar=.45; mass=1.; T=1.0; K=48; dt=T/K
 21
 22def normalize(r): return r/(torch.sum(r)*dx)
 23def gaussian(mu,sig): return torch.exp(-.5*((x-mu)/sig)**2)/(sig*math.sqrt(2*math.pi))
 24rho0=normalize(gaussian(0.,1.0))
 25rhot=normalize(.5*gaussian(-2.,.48)+.5*gaussian(2.,.48))
 26
 27def lap(f): return torch.fft.ifft(-(k*k)*torch.fft.fft(f)).real
 28def grad(f): return torch.fft.ifft(1j*k*torch.fft.fft(f)).real
 29
 30def split_step(psi, steps=K):
 31    # symmetric split-step implementation of the free Schrodinger equation
 32    kin=torch.exp(-1j*(hbar*k*k/(2*mass))*(dt/2))
 33    for _ in range(steps):
 34        z=torch.fft.ifft(torch.fft.fft(psi)*kin)
 35        z=torch.fft.ifft(torch.fft.fft(z)*kin)
 36    return z
 37
 38class PhaseNet(nn.Module):
 39    def __init__(self):
 40        super().__init__(); self.net=nn.Sequential(nn.Linear(1,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,1))
 41    def forward(self,z): return 3.0*self.net(z[:,None]).squeeze(-1)
 42
 43class VelocityNet(nn.Module):
 44    def __init__(self):
 45        super().__init__(); self.net=nn.Sequential(nn.Linear(2,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,1))
 46    def forward(self,xv,t): return self.net(torch.stack((xv,t.expand_as(xv)),1)).squeeze(-1)
 47
 48def sample_target(N):
 49    c=torch.randint(0,2,(N,),device=device); return torch.randn(N,device=device)*.48 + torch.where(c==0,-2.,2.)
 50def sample_source(N): return torch.randn(N,device=device)
 51
 52def stats(a,b):
 53    aa=a.detach().cpu().numpy(); bb=b.detach().cpu().numpy()
 54    w=wasserstein_distance(aa,bb)
 55    # RBF MMD, computed on subsample for stable, quick reporting
 56    aa=aa[:1000]; bb=bb[:1000]
 57    Daa=(aa[:,None]-aa[None,:])**2; Dbb=(bb[:,None]-bb[None,:])**2; Dab=(aa[:,None]-bb[None,:])**2
 58    mmd=float(np.mean(np.exp(-Daa/2))+np.mean(np.exp(-Dbb/2))-2*np.mean(np.exp(-Dab/2)))
 59    cov=float(((aa<-1.) .any() and (aa>1.).any()))
 60    return w,mmd,cov
 61
 62def verification():
 63    # Exact instantaneous Madelung residuals from a smooth nontrivial wavefunction.
 64    # Smooth periodic test state avoids boundary artifacts of the spectral derivative.
 65    rho_check=normalize(1.0 + 0.25*torch.cos(math.pi*x/L) + 0.10*torch.sin(2*math.pi*x/L))
 66    phase=.35*torch.sin(math.pi*x/L) + .12*torch.cos(2*math.pi*x/L)
 67    psi=torch.sqrt(rho_check)*torch.exp(1j*phase/hbar)
 68    psi_t=1j*hbar/(2*mass)*torch.fft.ifft(-(k*k)*torch.fft.fft(psi))
 69    rho=abs(psi)**2
 70    rho_t=2*torch.real(torch.conj(psi)*psi_t)
 71    S_t=hbar*torch.imag(psi_t/psi)
 72    S=phase; v=grad(S)/mass
 73    cont=rho_t+grad(rho*v)
 74    Q=-(hbar*hbar/(2*mass))*lap(torch.sqrt(rho))/(torch.sqrt(rho)+1e-12)
 75    hj=S_t+grad(S)**2/(2*mass)+Q
 76    # Conservative mass and particle-flow agreement for the same velocity field.
 77    with torch.no_grad():
 78        z=psi.clone(); times=[z]
 79        kin=torch.exp(-1j*(hbar*k*k/(2*mass))*(dt/2))
 80        for _ in range(K):
 81            z=torch.fft.ifft(torch.fft.fft(z)*kin)
 82            z=torch.fft.ifft(torch.fft.fft(z)*kin)
 83            times.append(z)
 84        rng=np.random.default_rng(SEED)
 85        cdf0=np.cumsum(rho_check.cpu().numpy())*dx
 86        u=rng.random(4000)
 87        p=np.interp(u, np.r_[0.,cdf0], np.r_[-L,x.cpu().numpy()])
 88        def interp(arr,q):
 89            q=np.mod(q+L,2*L)-L; uu=(q+L)/dx
 90            jj=np.floor(uu).astype(int)%n; ff=uu-np.floor(uu)
 91            return arr[jj]*(1-ff)+arr[(jj+1)%n]*ff
 92        def velocity(zz,q):
 93            dz=torch.fft.ifft(1j*k*torch.fft.fft(zz))
 94            rr=(abs(zz)**2).cpu().numpy()
 95            cur=(hbar*torch.imag(torch.conj(zz)*dz)/mass).cpu().numpy()
 96            return interp(cur/(rr+1e-12),q)
 97        sub=4; h=dt/sub
 98        for j in range(K):
 99            z0,z1=times[j],times[j+1]
100            # Linear interpolation in time is adequate at this refined check step.
101            def vv(q,frac): return velocity(z0*(1-frac)+z1*frac,q)
102            for qstep in range(sub):
103                f=qstep/sub
104                a=vv(p,f); b=vv(p+h*a/2,min(1.,f+.5/sub))
105                c=vv(p+h*b/2,min(1.,f+.5/sub)); d=vv(p+h*c,min(1.,f+1/sub))
106                p += h*(a+2*b+2*c+d)/6
107        hist,_=np.histogram(p,bins=n,range=(-L,L),density=True)
108        dens=(abs(times[-1].cpu().numpy())**2); dens=dens/(dens.sum()*dx)
109        particle_l1=float(np.sum(np.abs(hist-dens))*dx)
110    return {'continuity_max':float(torch.max(abs(cont))), 'HJ_max':float(torch.max(abs(hj))), 'mass0':float(torch.sum(rho)*dx), 'particle_density_L1':particle_l1}
111
112def train_phase():
113    model=PhaseNet().to(device); opt=torch.optim.Adam(model.parameters(),lr=.025)
114    target=rhot.detach(); t0=time.time()
115    for it in range(260):
116        ph=model(x); psi=torch.sqrt(rho0)*torch.exp(1j*ph/hbar); out=split_step(psi)
117        dens=abs(out)**2; loss=torch.mean((dens-target)**2)*20
118        # mild phase regularizer avoids unresolved grid oscillations
119        loss=loss+1e-4*torch.mean(grad(ph)**2)
120        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
121    with torch.no_grad():
122        out=split_step(torch.sqrt(rho0)*torch.exp(1j*model(x)/hbar)); d=abs(out)**2; cdf=torch.cumsum(d,0)*dx; u=torch.rand(3000,device=device); inds=torch.searchsorted(cdf,u).clamp(0,n-1); samples=x[inds]
123    return stats(samples,sample_target(len(samples))),float(time.time()-t0),float(loss)
124
125def train_flow():
126    model=VelocityNet().to(device); opt=torch.optim.Adam(model.parameters(),lr=.01); t0=time.time()
127    for it in range(500):
128        a=sample_source(256); b=sample_target(256); tt=torch.rand(256,device=device); z=(1-tt)*a+tt*b; true=b-a
129        loss=torch.mean((model(z,tt)-true)**2); opt.zero_grad(); loss.backward(); opt.step()
130    with torch.no_grad():
131        z=sample_source(3000)
132        for j in range(K):
133            tt=torch.full_like(z,j/K); z=z+dt*model(z,tt)
134    return stats(z,sample_target(len(z))),float(time.time()-t0),float(loss)
135
136if __name__=='__main__':
137    v=verification(); p=train_phase(); b=train_flow()
138    result={'device':device,'verification':v,'phase_only':{'W1':p[0][0],'MMD':p[0][1],'mode_coverage':p[0][2],'seconds':p[1],'train_loss':p[2]},'flow_matching':{'W1':b[0][0],'MMD':b[0][1],'mode_coverage':b[0][2],'seconds':b[1],'train_loss':b[2]}}
139    print(json.dumps(result,indent=2))