PDE Sinkhorn with asymmetric geometric boundaries / experiment.py
Mechanism confirmed, baseline not beaten
1import json, math
2import numpy as np
3from scipy.linalg import expm
4
5SEED=1094
6np.random.seed(SEED)
7
8# ---------- Core geometric calculation ----------
9def oblique_sweep():
10 # At a vertical face n=(1,0), G=[[1,c],[c,1]].
11 # r=Gn/(n^T G n)=(1,c), so r.grad=0 iff grad_x+c grad_y=0.
12 cs=np.array([0., .2, .5, .8, .95])
13 # A deliberately oblique tangent-compatible gradient: grad=(-c,1).
14 oblique=np.abs((-cs)+cs*1.)
15 # The normal Neumann condition grad_x=0 leaves residual c in the oblique BC.
16 normal_res=np.abs(cs)
17 # Normal diffusive flux n^T G grad has the same residual for this test.
18 flux_res=np.abs(cs+(-cs)) # exact zero under oblique condition
19 # Use a second gradient to show the normal condition fails flux cancellation.
20 wrong_flux=np.abs(cs)
21 return {"correlations":cs.tolist(), "oblique_residual":oblique.tolist(),
22 "normal_bc_residual":normal_res.tolist(), "oblique_flux_residual":flux_res.tolist(),
23 "normal_bc_flux_residual":wrong_flux.tolist()}
24
25# ---------- Finite-volume reflecting diffusion ----------
26def reflected_generator(n, D):
27 # Cell-centered 1D finite-volume generator for q_t=D q_xx, zero face flux.
28 # Column vector q, Q has column sums zero and nonnegative off diagonal.
29 h=1.0/n; Q=np.zeros((n,n))
30 rate=D/(h*h)
31 for j in range(n):
32 if j>0: Q[j-1,j]+=rate; Q[j,j]-=rate
33 if j<n-1: Q[j+1,j]+=rate; Q[j,j]-=rate
34 return Q
35
36def conservation_sweep():
37 n=80; x=(np.arange(n)+.5)/n
38 q=np.exp(-((x-.23)/.07)**2); q/=q.sum()
39 rows=[]
40 for D in [0., 1e-3, .03, .3, 1.0]:
41 Q=reflected_generator(n,D)
42 q1=expm(Q*.37)@q
43 # Boundary face current is imposed to exactly zero by construction.
44 mass_err=abs(q1.sum()-1.)
45 minq=q1.min()
46 rows.append({"D":D,"mass_error":float(mass_err),"min_density":float(minq),
47 "boundary_flux_left":0.,"boundary_flux_right":0.})
48 return rows
49
50def stability_sweep():
51 n=80; D=1.; h=1/n
52 Q=reflected_generator(n,D)
53 lam=np.linalg.eigvalsh(Q).min()
54 dt_spec=2/abs(lam)
55 # For forward Euler, positivity requires dt <= 1/max_j(-Q_jj),
56 # while l2 stability only requires dt <= 2/|lambda_min|.
57 dt_pos=1/np.max(-np.diag(Q))
58 # Positive alternating data strongly excites the highest-frequency mode.
59 z0=1.+.9*(-1.)**np.arange(n); z0/=z0.sum()
60 ratios=[.5,.9,1.0,1.05,1.5]
61 out=[]
62 for ratio in ratios:
63 dt=ratio*dt_pos; z=z0.copy()
64 for _ in range(80): z=z+dt*(Q@z)
65 out.append({"ratio_to_positivity_bound":ratio,"dt":dt,
66 "min_after_80":float(z.min()),"l1_norm":float(np.linalg.norm(z,1)),
67 "nonnegative_observed":bool(z.min()>=-1e-12)})
68 return {"lambda_min":float(lam),"predicted_positivity_dt":float(dt_pos),
69 "predicted_spectral_dt":float(dt_spec),"sweep":out}
70
71# ---------- Small endpoint scaling experiment ----------
72def kl(p,q):
73 return float(np.sum(p*np.log((p+1e-15)/(q+1e-15))))
74
75def sinkhorn(P, p0, p1, iters=100):
76 # P maps initial mass to terminal mass: terminal=P @ initial.
77 a=np.ones(len(p0)); b=np.ones(len(p1))
78 # joint diag(a) P diag(b), alternating endpoint marginal fitting
79 for k in range(iters):
80 a=p0/(P.T@b+1e-15)
81 b=p1/(P@a+1e-15)
82 return a,b
83
84def sinkhorn_demo():
85 n=60; x=(np.arange(n)+.5)/n; h=1/n
86 p0=np.exp(-.5*((x-.22)/.07)**2); p0/=p0.sum()
87 p1=np.exp(-.5*((x-.78)/.07)**2); p1/=p1.sum()
88 D=.015; T=.45
89 # PDE/reflected kernel from finite-volume Fokker-Planck operator.
90 Pref=expm(reflected_generator(n,D)*T)
91 # Unconstrained Gaussian transition kernel, row-normalized only for comparison.
92 K=np.exp(-(x[:,None]-x[None,:])**2/(4*D*T)); K/=K.sum(axis=0,keepdims=True)
93 results=[]
94 for name,P in [("reflected_PDE",Pref),("unconstrained_kernel",K)]:
95 a,b=sinkhorn(P,p0,p1,80)
96 left=a*(P.T@b); right=b*(P@a)
97 results.append({"method":name,"initial_KL":kl(left/left.sum(),p0),
98 "terminal_KL":kl(right/right.sum(),p1),
99 "boundary_mass_initial":float(left[0]+left[-1]),
100 "boundary_mass_terminal":float(right[0]+right[-1])})
101 return results
102
103def main():
104 report={"seed":SEED,"oblique_geometry":oblique_sweep(),
105 "finite_volume_no_flux":conservation_sweep(),"explicit_stability":stability_sweep(),
106 "sinkhorn_demo":sinkhorn_demo()}
107 with open("results.json","w") as f: json.dump(report,f,indent=2)
108 print(json.dumps(report,indent=2))
109
110if __name__=="__main__": main()