File size: 4,401 Bytes
d8f9639 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | """Independent numerical checks of the lecture's identities and saved runs."""
import argparse,json
from pathlib import Path
import numpy as np
import torch
from scipy.linalg import expm
from ot_sbm_examples import *
from learned_bridges import sf2m_targets
def main():
parser=argparse.ArgumentParser();parser.add_argument('--results',default='results');args=parser.parse_args();out=Path(args.results);out.mkdir(parents=True,exist_ok=True);checks={}
def check(name,value):
checks[name]=bool(value)
if not value:raise AssertionError(name)
ot=transport_example();check('independent LP and dual certificate',abs(ot['ot_cost']-ot['dual'])<1e-10 and abs(ot['dual']-3.4)<1e-10)
check('Sinkhorn endpoint constraints',np.allclose(ot['entropic_plan'].sum(1),ot['a']) and np.allclose(ot['entropic_plan'].sum(0),ot['b']))
br=finite_bridge_example();check('Doob transition normalization',all(np.allclose(q.sum(1),1) for q in br['Q']))
check('bridge endpoint and path normalization',np.allclose(br['marginals'][-1],[.2,.8]) and abs(br['bridge'].sum()-1)<1e-10)
# Conditional path laws are identical when the two endpoints are fixed.
for a in [0,1]:
for b in [0,1]:
ix=(br['paths'][:,0]==a)&(br['paths'][:,-1]==b)
check(f'reference conditional bridges {a}{b}',np.allclose(br['bridge'][ix]/br['bridge'][ix].sum(),br['reference'][ix]/br['reference'][ix].sum()))
imf=discrete_imf_example();check('IMF converges to full path law',imf['final_path_l1']<1e-10)
check('IMF path KL is nonincreasing',np.max(np.diff(imf['kl_to_bridge']))<1e-12)
# Independently differentiate analytic marginals to check the master equation.
G=np.array([[-2.,2.],[1.,-1.]]);a=np.array([.5,.5]);b=np.array([.2,.8]);pi,f,g,_=sinkhorn(a,b,np.log(a[:,None]*expm(G)))
p=lambda t:((a*f)@expm(t*G))*(expm((1-t)*G)@g)
h=expm(.5*G)@g;Q=G*h[None,:]/h[:,None];np.fill_diagonal(Q,0);np.fill_diagonal(Q,-Q.sum(1))
check('CTMC rates satisfy the independent master equation',np.allclose((p(.500001)-p(.499999))/.000002,p(.5)@Q,atol=1e-8))
t=torch.tensor([[.1],[.3],[.8]],dtype=torch.float64);x0=torch.tensor([[0.],[1.],[-1.]],dtype=torch.float64);x1=x0+2;z=torch.tensor([[.7],[-.4],[1.2]],dtype=torch.float64)
x,v,s,sd=sf2m_targets(t,x0,x1,z)
check('conditional velocity plus score equals bridge drift',torch.allclose(v+.5*s,(x1-x)/(1-t)))
rw=reward_bridge_example();check('exact path importance identity',np.allclose(rw['weighted_paths'],rw['target_paths']))
bg=branching_example();check('analytic branch allocation is positive and conserved',np.min(bg['weights'])>=0 and np.allclose(bg['total_mass'],1))
geo=entangled_geometry_example();check('cone projection and finite-step bound',abs(geo['orthogonal']@geo['unit_direction'])<1e-12 and geo['distance2_after']<geo['distance2_before'])
d=geo['direction'];v=geo['bias'];dt=geo['max_dt']*1.1
check('oversized inward step can increase distance',(d-dt*v)@(d-dt*v)>d@d)
for key in ['dsb','dsbm','sf2m','tr2d2','branch','entangled']:
path=out/(key+'_learned.json')
if not path.exists():continue
r=json.load(open(path))
if key=='dsb':q=r['history'][-1];check('learned DSB Gaussian endpoint',abs(q['terminal_mean']-2)<.1 and abs(q['terminal_variance']-1)<.15)
if key=='dsbm':check('learned DSBM Gaussian endpoint',abs(r['terminal_mean']-2)<.1 and abs(r['terminal_variance']-1)<.15)
if key=='sf2m':check('learned SF2M fields and endpoint',r['velocity_mse']<.05 and r['score_mse']<.05 and abs(r['terminal_variance']-1)<.15)
if key=='tr2d2':check('full-support search WDCE improves reference reward',r['history'][-1]['reward']>r['reference_reward'] and r['history'][-1]['target_l1']<.5)
if key=='branch':check('soft branch constraints have small residuals',r['max_mass_error']<.02 and r['minimum_weight']>-.002)
if key=='entangled':check('learned cone bias retains local alignment',r['min_bias_alignment']>=-1e-6)
path=out/'discrete_learned_results.json'
if path.exists():
r=json.load(open(path));check('learned DDSBM rate projection',r['ddsbm']['relative_mse']<.08);check('learned CSBM transition projection',r['csbm']['maximum_transition_error']<.001)
(out/'checks.json').write_text(json.dumps(checks,indent=2));print(len(checks),'checks passed')
if __name__=='__main__':main()
|