Padé-Hermite Neural ODE Integrator / experiment.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def pade_target(s, z):
  8    p = sum(math.factorial(2*s-k)*math.factorial(s-1) /
  9            (math.factorial(2*s)*math.factorial(k)*math.factorial(s-1-k)) * z**k
 10            for k in range(s))
 11    q = sum(math.factorial(2*s-k)*math.factorial(s+1) /
 12            (math.factorial(2*s)*math.factorial(k)*math.factorial(s+1-k)) * (-z)**k
 13            for k in range(s+2))
 14    return p / q
 15
 16
 17def make_coeffs(s=2, nodes=None):
 18    if nodes is None:
 19        nodes = np.linspace(0.0, 1.0, s + 1)
 20    q = np.zeros(2*s + 1)
 21    for k in range(s + 2):
 22        q[k] = (math.factorial(2*s-k)*math.factorial(s+1) /
 23                (math.factorial(2*s)*math.factorial(k)*math.factorial(s+1-k)) * (-1)**k)
 24    pi = np.zeros((s+1, 2*s + 1))
 25    for j, c in enumerate(nodes):
 26        for r in range(2*s + 1):
 27            pi[j, r] = sum(q[k] * c**(r-k) / math.factorial(r-k)
 28                            for k in range(min(s+1, r)+1))
 29    # B represents Pi_j and z Pi_j in P_{2s+1}.
 30    basis = [np.r_[pi[j], 0.0] for j in range(s+1)]
 31    basis += [np.r_[0.0, pi[j]] for j in range(s+1)]
 32    B = np.column_stack(basis)
 33    aa = np.zeros((s, s+1)); hh = np.zeros_like(aa)
 34    for i in range(1, s+1):
 35        # (Pi_i-Q)/z: constant difference is zero; represent through degree 2s+1.
 36        rhs = np.r_[pi[i, 1:] - q[1:], 0.0, 0.0]
 37        sol = np.linalg.solve(B, rhs)
 38        aa[i-1] = sol[:s+1]
 39        hh[i-1] = sol[s+1:]
 40    return np.asarray(nodes), aa, hh, q
 41
 42
 43def step_ph(u, t, h, fun, jac, A, H, nodes, tol=1e-11, maxit=100):
 44    s = len(nodes)-1; u = np.asarray(u, float)
 45    U = np.array([u + nodes[i]*h*fun(t, u) for i in range(1, s+1)])
 46    for it in range(maxit):
 47        states = np.vstack([u, U])
 48        fs = np.array([fun(t + nodes[j]*h, states[j]) for j in range(s+1)])
 49        gs = np.array([jac(t + nodes[j]*h, states[j]) @ fs[j] for j in range(s+1)])
 50        new = np.array([u + h*(A[i] @ fs) + h*h*(H[i] @ gs) for i in range(s)])
 51        err = np.max(np.abs(new-U)) / max(1.0, np.max(np.abs(new)))
 52        U = 0.7*new + 0.3*U
 53        if err < tol:
 54            return U[-1], it+1
 55    raise RuntimeError('nonlinear stage solve failed')
 56
 57
 58def ph_integrate(u, t0, t1, h, fun, jac, A, H, nodes):
 59    n = int(round((t1-t0)/h)); x = np.asarray(u, float); t=t0; its=0
 60    for _ in range(n):
 61        x, k = step_ph(x, t, h, fun, jac, A, H, nodes)
 62        its += k; t += h
 63    return x, its
 64
 65
 66def rk4(u, t, h, f):
 67    k1=f(t,u); k2=f(t+h/2,u+h*k1/2); k3=f(t+h/2,u+h*k2/2); k4=f(t+h,u+h*k3)
 68    return u+h*(k1+2*k2+2*k3+k4)/6
 69
 70
 71def main():
 72    np.random.seed(3)
 73    nodes,A,H,q = make_coeffs(2)
 74    zs = np.array([-1., -5., -10., -100., -1000.])
 75    pade_err = [abs(pade_target(2,z) - (1+z/4)/(1-3*z/4+z*z/4-z**3/24)) for z in zs]
 76    order_rows=[]
 77    for h in [0.2,0.1,0.05,0.025]:
 78        f=lambda t,x: np.array([-x[0]])
 79        j=lambda t,x: np.array([[-1.]])
 80        y,_=ph_integrate([1.],0.,1.,h,f,j,A,H,nodes)
 81        order_rows.append((h, float(abs(y[0]-math.exp(-1)))))
 82    orders=[math.log(a[1]/b[1],2) for a,b in zip(order_rows[:-1],order_rows[1:])]
 83    stability=[]
 84    for z in [-1,-2,-3,-5,-10,-50,-100]:
 85        r=pade_target(2,z)
 86        stability.append({'z':z,'R':float(r),'euler_abs':abs(1+z)})
 87    lam=80.
 88    def f(t,x): return np.array([-lam*x[0] - 0.2*x[1]**3, -x[1] + 0.1*x[0]])
 89    def jac(t,x): return np.array([[-lam, -0.6*x[1]**2],[0.1,-1.]])
 90    x0=np.array([1.,1.]); exact_ref,_=ph_integrate(x0,0.,1.,0.002,f,jac,A,H,nodes)
 91    cmp=[]
 92    for h in [0.05,0.025,0.0125]:
 93        try:
 94            yp,it=ph_integrate(x0,0.,1.,h,f,jac,A,H,nodes)
 95            pe=float(np.linalg.norm(yp-exact_ref))
 96        except Exception:
 97            pe=float('inf'); it=-1
 98        y=x0.copy(); ok=True
 99        for k in range(int(round(1/h))):
100            y=rk4(y,k*h,h,f)
101            if not np.all(np.isfinite(y)): ok=False
102        cmp.append({'h':h,'ph_error':pe,'ph_iterations':it,
103                    'rk4_error':float(np.linalg.norm(y-exact_ref)) if ok else float('inf')})
104    # Quantitative math checks: Hermite moments and A-stability grid.
105    moment_residual=0.0
106    for i,c in enumerate(nodes[1:]):
107        for m in range(1,5):
108            lhs=sum(A[i,j]*nodes[j]**(m-1)/math.factorial(m-1) +
109                    H[i,j]*(m-1)*nodes[j]**(m-2)/math.factorial(m-1) if m>=2 else
110                    A[i,j]*nodes[j]**(m-1)/math.factorial(m-1)
111                    for j in range(3))
112            # Integral of t^(m-1) from 0 to c, with derivative-data term
113            target=c**m/math.factorial(m)
114            moment_residual=max(moment_residual,abs(lhs-target))
115    grid=[]
116    for re in np.linspace(-20,0,41):
117        for im in np.linspace(-20,20,81):
118            grid.append(abs(pade_target(2,re+1j*im)))
119    a_stab_max=max(grid)
120    # One-step amplification from the actual anchored stages, across node choices.
121    node_sweep=[]
122    for nd in [np.array([0.,.5,1.]), np.array([0.,.25,1.]), np.array([0.,.75,1.])]:
123        nn,aa,hh,_=make_coeffs(2,nd)
124        vals=[]
125        for z in [-1.,-10.,-100.]:
126            # Exact solve of the scalar linear stage equations, avoiding
127            # convergence assumptions in this stability-only diagnostic.
128            c=nn; m=len(c)-1
129            M=np.eye(m)-z*aa[ :,1:]-z*z*hh[:,1:]
130            rhs=np.ones(m)+z*aa[:,0]+z*z*hh[:,0]
131            yy=np.linalg.solve(M,rhs)[-1]
132            vals.append(float(yy))
133        node_sweep.append({'nodes':nd.tolist(),'R':vals})
134    # Predicted explicit-Euler boundary is z=-2; measure the transition.
135    euler_boundary=[]
136    for z in [-1.9,-2.0,-2.1]:
137        euler_boundary.append({'z':z,'abs_R':abs(1+z)})
138    asymptotic_slope=(math.log(abs(pade_target(2,-1000.)/pade_target(2,-100.)),10))
139    out={'s':2,'nodes':nodes.tolist(),'A':A.tolist(),'H':H.tolist(),
140         'pade_negative_axis_errors':pade_err,'order_errors':order_rows,
141         'observed_orders':orders,'stability_sweep':stability,
142         'node_independence':node_sweep,'euler_boundary':euler_boundary,
143         'observed_log10_decay_slope_100_to_1000':asymptotic_slope,
144         'max_moment_residual':moment_residual,'max_left_half_plane_abs_R':float(a_stab_max),
145         'stiff_comparison':cmp}
146    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
147
148if __name__=='__main__': main()