Exponential-Map Stochastic Residual Layer / run_experiment.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
 1import json, math
 2import numpy as np
 3from sphere_exp_residual import sphere_exp
 4np.set_printoptions(precision=8, suppress=True)
 5
 6def exp_np(x,v):
 7    r=np.linalg.norm(v)
 8    return math.cos(r)*x + (math.sin(r)/r if r>1e-14 else 1.)*v
 9
10def bfield(x):
11    c=np.array([.31,-.47,.22])
12    return c-x*np.dot(x,c)
13
14def intrinsic(x,h,n):
15    for _ in range(n): x=exp_np(x,h*bfield(x))
16    return x
17
18def chart(x): return x[:2]/(1-x[2])
19def invchart(q):
20    r=np.dot(q,q)
21    return np.array([2*q[0]/(1+r),2*q[1]/(1+r),(r-1)/(1+r)])
22def jacobian_inv(q):
23    J=np.empty((3,2)); eps=1e-6
24    for j in range(2):
25        e=np.zeros(2); e[j]=eps
26        J[:,j]=(invchart(q+e)-invchart(q-e))/(2*eps)
27    return J
28def transformed(q): return np.array([q[0]+.55*q[0]**2,q[1]+.30*q[0]*q[1]])
29def inv_transformed(p):
30    q=p.copy()
31    for _ in range(12):
32        f=transformed(q)-p
33        J=np.array([[1+1.1*q[0],0],[.30*q[1],1+.30*q[0]]])
34        q-=np.linalg.solve(J,f)
35    return q
36def jac_transformed(q): return np.array([[1+1.1*q[0],0],[.30*q[1],1+.30*q[0]]])
37def coordinate_euler(x0,h,n,which):
38    q=chart(x0) if which=='q' else transformed(chart(x0))
39    for _ in range(n):
40        qq=q if which=='q' else inv_transformed(q)
41        x=invchart(qq); dq=np.linalg.lstsq(jacobian_inv(qq),bfield(x),rcond=None)[0]
42        q=q+h*dq if which=='q' else q+h*(jac_transformed(qq)@dq)
43    return invchart(q if which=='q' else inv_transformed(q))
44def slope(hs,errs): return float(np.polyfit(np.log(hs),np.log(np.maximum(errs,1e-18)),1)[0])
45
46def parallel_transport_from_start(x0,v0,theta):
47    r=np.linalg.norm(v0); u=v0/r
48    return r*(-math.sin(theta)*x0+math.cos(theta)*u)
49
50def main():
51    x0=np.array([.35,-.25,.9027735]); x0/=np.linalg.norm(x0); T=.8
52    hs=np.array([.1,.05,.025,.0125])
53    ref=intrinsic(x0,1e-5,int(T/1e-5)); exp_err=[]; euler_err=[]; norm_exp=[]; norm_euler=[]; chart_gap=[]
54    for h in hs:
55        n=round(T/h); xe=intrinsic(x0,h,n); xu=x0.copy()
56        for _ in range(n): xu=xu+h*bfield(xu)
57        q1=coordinate_euler(x0,h,n,'q'); q2=coordinate_euler(x0,h,n,'p')
58        exp_err.append(np.linalg.norm(xe-ref)); euler_err.append(np.linalg.norm(xu-ref))
59        norm_exp.append(abs(np.linalg.norm(xe)-1)); norm_euler.append(abs(np.linalg.norm(xu)-1)); chart_gap.append(np.linalg.norm(q1-q2))
60    # Correct geodesic composition test: velocity is parallel transported at every point.
61    v=np.array([.2,.3,-.1]); v-=x0*np.dot(x0,v); r=np.linalg.norm(v); exact=exp_np(x0,v); geo_err=[]
62    for h in hs:
63        z=x0.copy()
64        for k in range(round(1/h)):
65            vk=parallel_transport_from_start(x0,v,k*h*r)
66            z=exp_np(z,h*vk)
67        geo_err.append(np.linalg.norm(z-exact))
68    # Tangent isotropic noise: empirical covariance in an orthonormal tangent frame.
69    rng=np.random.default_rng(7); e1=np.array([1.,0,0]); e1-=x0*np.dot(x0,e1); e1/=np.linalg.norm(e1); e2=np.cross(x0,e1)
70    Z=rng.normal(size=(200000,2)); cov=np.cov(Z,rowvar=False)
71    tangent_cov=np.cov((Z[:,0,None]*e1+Z[:,1,None]*e2).T)
72    out={'step_sizes':hs.tolist(),'predictions':{
73      'P1_norm_preservation':{'predicted':'Exp error roundoff; additive error nonzero','exp_abs_errors':norm_exp,'additive_abs_errors':norm_euler},
74      'P2_first_order_global_error':{'predicted':'global endpoint error O(h), slope near 1','exp_errors':exp_err,'additive_errors':euler_err,'exp_slope':slope(hs,exp_err),'additive_slope':slope(hs,euler_err)},
75      'P3_chart_effect':{'predicted':'coordinate Euler discrepancy O(h), slope near 1','chart_gaps':chart_gap,'slope':slope(hs,chart_gap)},
76      'P4_geodesic_exactness':{'predicted':'Exp with parallel-transported velocity is step independent','errors':geo_err},
77      'P5_tangent_noise':{'predicted':'orthonormal tangent coordinates have covariance I','coordinate_covariance':cov.tolist(),'ambient_tangent_covariance_error':float(np.linalg.norm(tangent_cov-(np.eye(3)-np.outer(x0,x0))/1.0))}
78    },'summary':{'max_exp_norm_error':max(norm_exp),'coarsest_additive_norm_error':norm_euler[0],'coarsest_exp_error':exp_err[0],'coarsest_additive_error':euler_err[0],'max_geodesic_error':max(geo_err)}}
79    print(json.dumps(out,indent=2))
80if __name__=='__main__': main()