import json, time, math import numpy as np import torch import torch.nn as nn from scipy.stats import wasserstein_distance torch.set_default_dtype(torch.float64) SEED=7 np.random.seed(SEED); torch.manual_seed(SEED) device='cuda' if torch.cuda.is_available() else 'cpu' try: if device=='cuda': torch.zeros(1,device='cuda') except Exception: device='cpu' # Periodic spectral grid; the distributions are negligible at the boundaries. L=8.0; n=256; dx=2*L/n x=torch.linspace(-L,L-dx,n,device=device) k=2*math.pi*torch.fft.fftfreq(n,d=dx).to(device) hbar=.45; mass=1.; T=1.0; K=48; dt=T/K def normalize(r): return r/(torch.sum(r)*dx) def gaussian(mu,sig): return torch.exp(-.5*((x-mu)/sig)**2)/(sig*math.sqrt(2*math.pi)) rho0=normalize(gaussian(0.,1.0)) rhot=normalize(.5*gaussian(-2.,.48)+.5*gaussian(2.,.48)) def lap(f): return torch.fft.ifft(-(k*k)*torch.fft.fft(f)).real def grad(f): return torch.fft.ifft(1j*k*torch.fft.fft(f)).real def split_step(psi, steps=K): # symmetric split-step implementation of the free Schrodinger equation kin=torch.exp(-1j*(hbar*k*k/(2*mass))*(dt/2)) for _ in range(steps): z=torch.fft.ifft(torch.fft.fft(psi)*kin) z=torch.fft.ifft(torch.fft.fft(z)*kin) return z class PhaseNet(nn.Module): def __init__(self): super().__init__(); self.net=nn.Sequential(nn.Linear(1,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,1)) def forward(self,z): return 3.0*self.net(z[:,None]).squeeze(-1) class VelocityNet(nn.Module): def __init__(self): super().__init__(); self.net=nn.Sequential(nn.Linear(2,48),nn.Tanh(),nn.Linear(48,48),nn.Tanh(),nn.Linear(48,1)) def forward(self,xv,t): return self.net(torch.stack((xv,t.expand_as(xv)),1)).squeeze(-1) def sample_target(N): c=torch.randint(0,2,(N,),device=device); return torch.randn(N,device=device)*.48 + torch.where(c==0,-2.,2.) def sample_source(N): return torch.randn(N,device=device) def stats(a,b): aa=a.detach().cpu().numpy(); bb=b.detach().cpu().numpy() w=wasserstein_distance(aa,bb) # RBF MMD, computed on subsample for stable, quick reporting aa=aa[:1000]; bb=bb[:1000] Daa=(aa[:,None]-aa[None,:])**2; Dbb=(bb[:,None]-bb[None,:])**2; Dab=(aa[:,None]-bb[None,:])**2 mmd=float(np.mean(np.exp(-Daa/2))+np.mean(np.exp(-Dbb/2))-2*np.mean(np.exp(-Dab/2))) cov=float(((aa<-1.) .any() and (aa>1.).any())) return w,mmd,cov def verification(): # Exact instantaneous Madelung residuals from a smooth nontrivial wavefunction. # Smooth periodic test state avoids boundary artifacts of the spectral derivative. rho_check=normalize(1.0 + 0.25*torch.cos(math.pi*x/L) + 0.10*torch.sin(2*math.pi*x/L)) phase=.35*torch.sin(math.pi*x/L) + .12*torch.cos(2*math.pi*x/L) psi=torch.sqrt(rho_check)*torch.exp(1j*phase/hbar) psi_t=1j*hbar/(2*mass)*torch.fft.ifft(-(k*k)*torch.fft.fft(psi)) rho=abs(psi)**2 rho_t=2*torch.real(torch.conj(psi)*psi_t) S_t=hbar*torch.imag(psi_t/psi) S=phase; v=grad(S)/mass cont=rho_t+grad(rho*v) Q=-(hbar*hbar/(2*mass))*lap(torch.sqrt(rho))/(torch.sqrt(rho)+1e-12) hj=S_t+grad(S)**2/(2*mass)+Q # Conservative mass and particle-flow agreement for the same velocity field. with torch.no_grad(): z=psi.clone(); times=[z] kin=torch.exp(-1j*(hbar*k*k/(2*mass))*(dt/2)) for _ in range(K): z=torch.fft.ifft(torch.fft.fft(z)*kin) z=torch.fft.ifft(torch.fft.fft(z)*kin) times.append(z) rng=np.random.default_rng(SEED) cdf0=np.cumsum(rho_check.cpu().numpy())*dx u=rng.random(4000) p=np.interp(u, np.r_[0.,cdf0], np.r_[-L,x.cpu().numpy()]) def interp(arr,q): q=np.mod(q+L,2*L)-L; uu=(q+L)/dx jj=np.floor(uu).astype(int)%n; ff=uu-np.floor(uu) return arr[jj]*(1-ff)+arr[(jj+1)%n]*ff def velocity(zz,q): dz=torch.fft.ifft(1j*k*torch.fft.fft(zz)) rr=(abs(zz)**2).cpu().numpy() cur=(hbar*torch.imag(torch.conj(zz)*dz)/mass).cpu().numpy() return interp(cur/(rr+1e-12),q) sub=4; h=dt/sub for j in range(K): z0,z1=times[j],times[j+1] # Linear interpolation in time is adequate at this refined check step. def vv(q,frac): return velocity(z0*(1-frac)+z1*frac,q) for qstep in range(sub): f=qstep/sub a=vv(p,f); b=vv(p+h*a/2,min(1.,f+.5/sub)) c=vv(p+h*b/2,min(1.,f+.5/sub)); d=vv(p+h*c,min(1.,f+1/sub)) p += h*(a+2*b+2*c+d)/6 hist,_=np.histogram(p,bins=n,range=(-L,L),density=True) dens=(abs(times[-1].cpu().numpy())**2); dens=dens/(dens.sum()*dx) particle_l1=float(np.sum(np.abs(hist-dens))*dx) 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} def train_phase(): model=PhaseNet().to(device); opt=torch.optim.Adam(model.parameters(),lr=.025) target=rhot.detach(); t0=time.time() for it in range(260): ph=model(x); psi=torch.sqrt(rho0)*torch.exp(1j*ph/hbar); out=split_step(psi) dens=abs(out)**2; loss=torch.mean((dens-target)**2)*20 # mild phase regularizer avoids unresolved grid oscillations loss=loss+1e-4*torch.mean(grad(ph)**2) opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step() with torch.no_grad(): 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] return stats(samples,sample_target(len(samples))),float(time.time()-t0),float(loss) def train_flow(): model=VelocityNet().to(device); opt=torch.optim.Adam(model.parameters(),lr=.01); t0=time.time() for it in range(500): a=sample_source(256); b=sample_target(256); tt=torch.rand(256,device=device); z=(1-tt)*a+tt*b; true=b-a loss=torch.mean((model(z,tt)-true)**2); opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): z=sample_source(3000) for j in range(K): tt=torch.full_like(z,j/K); z=z+dt*model(z,tt) return stats(z,sample_target(len(z))),float(time.time()-t0),float(loss) if __name__=='__main__': v=verification(); p=train_phase(); b=train_flow() 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]}} print(json.dumps(result,indent=2))