Feasible High-Order Neural ODE Solver / experiment.py
Mechanism failed
1import json
2import numpy as np
3from scipy.optimize import root, minimize
4
5# Two-stage Gauss-Legendre Runge-Kutta method, formal order four.
6A = np.array([[0.25, 0.25 - np.sqrt(3.0) / 6.0],
7 [0.25 + np.sqrt(3.0) / 6.0, 0.25]])
8b = np.array([0.5, 0.5])
9c = np.array([0.5 - np.sqrt(3.0) / 6.0,
10 0.5 + np.sqrt(3.0) / 6.0])
11
12
13def gauss_unconstrained(y, t, h, f):
14 y = np.asarray(y, dtype=float)
15 d = y.size
16
17 def residual(x):
18 Y = x.reshape(2, d)
19 vals = np.array([f(t + c[i] * h, Y[i]) for i in range(2)])
20 return (Y - y[None, :] - h * A.dot(vals)).ravel()
21
22 sol = root(residual, np.tile(y, 2), method="hybr")
23 if not sol.success:
24 raise RuntimeError("implicit solve failed: " + sol.message)
25 Y = sol.x.reshape(2, d)
26 vals = np.array([f(t + c[i] * h, Y[i]) for i in range(2)])
27 z = y + h * b.dot(vals)
28 return z, Y, float(np.linalg.norm(residual(sol.x))), 2
29
30
31def gauss_constrained(y, t, h, f, lower, upper, rho=1.0):
32 """Paper objective with explicit box feasibility on stages and final state."""
33 y = np.asarray(y, dtype=float)
34 lower = np.broadcast_to(np.asarray(lower, dtype=float), y.shape)
35 upper = np.broadcast_to(np.asarray(upper, dtype=float), y.shape)
36 d = y.size
37
38 def unpack(x):
39 return x[:2*d].reshape(2, d), x[2*d:]
40
41 def residuals(x):
42 Y, z = unpack(x)
43 vals = np.array([f(t + c[i] * h, Y[i]) for i in range(2)])
44 rs = Y - y[None, :] - h * A.dot(vals)
45 rz = z - y - h * b.dot(vals)
46 return rs, rz
47
48 def objective(x):
49 rs, rz = residuals(x)
50 return 0.5 * (np.sum(rs * rs) + rho * np.sum(rz * rz))
51
52 # Explicit bounds are equivalent to projected optimization for this box.
53 x0 = np.concatenate([np.tile(np.clip(y, lower, upper), 2),
54 np.clip(y, lower, upper)])
55 bounds = [(float(lower[j]), float(upper[j])) for _ in range(3) for j in range(d)]
56 sol = minimize(objective, x0, method="L-BFGS-B", bounds=bounds,
57 options={"maxiter": 250, "ftol": 1e-14, "gtol": 1e-10, "maxls": 30})
58 Y, z = unpack(sol.x)
59 rs, rz = residuals(sol.x)
60 return z, Y, float(np.sqrt(np.sum(rs * rs) + rho * np.sum(rz * rz))), 2, int(sol.nit), bool(sol.success)
61
62
63def integrate(method, y0, f, T, n, box=None):
64 y = np.asarray(y0, dtype=float).copy()
65 t = 0.0
66 max_violation = 0.0
67 residual_sum = 0.0
68 iterations = 0
69 for _ in range(n):
70 h = T / n
71 if method == "unconstrained":
72 z, Y, r, _ = gauss_unconstrained(y, t, h, f)
73 elif method == "clip":
74 z, Y, r, _ = gauss_unconstrained(y, t, h, f)
75 z = np.clip(z, box[0], box[1])
76 Y = np.clip(Y, box[0], box[1])
77 elif method == "constrained":
78 z, Y, r, _, nit, ok = gauss_constrained(y, t, h, f, box[0], box[1])
79 iterations += nit
80 if not ok:
81 raise RuntimeError("constrained optimizer failed")
82 else:
83 raise ValueError(method)
84 if box is not None:
85 max_violation = max(max_violation,
86 float(max(np.max(box[0] - Y), np.max(Y - box[1]),
87 np.max(box[0] - z), np.max(z - box[1]), 0.0)))
88 residual_sum += r
89 y, t = z, t + T / n
90 return y, max_violation, residual_sum / n, iterations
91
92
93def main():
94 np.random.seed(7)
95 results = {"math_check": [], "stiff_test": []}
96
97 # Claimed high-order behavior away from the boundary: y' = y, exact exp(T).
98 f_exp = lambda t, y: y
99 exact = np.exp(1.0)
100 errors = []
101 for n in [2, 4, 8, 16, 32]:
102 y, _, _, _ = integrate("unconstrained", [1.0], f_exp, 1.0, n)
103 errors.append([n, float(abs(y[0] - exact))])
104 orders = [float(np.log(errors[i-1][1] / errors[i][1]) / np.log(2.0))
105 for i in range(1, len(errors))]
106 results["math_check"] = {"errors": errors, "observed_orders": orders,
107 "expected_order": 4}
108
109 # Stiff positive scalar benchmark. Exact solution remains in [0, 1].
110 # The large positive rate makes unconstrained RK stages leave the box.
111 lam = 40.0
112 f_stiff = lambda t, y: lam * (1.0 - y)
113 T, n, y0 = 1.0, 10, np.array([0.0])
114 box = (np.array([0.0]), np.array([1.0]))
115 exact_final = 1.0 - np.exp(-lam * T)
116 for method in ["unconstrained", "clip", "constrained"]:
117 y, violation, residual, iters = integrate(method, y0, f_stiff, T, n, box)
118 results["stiff_test"].append({
119 "method": method, "final": float(y[0]),
120 "abs_error": float(abs(y[0] - exact_final)),
121 "max_constraint_violation": violation,
122 "mean_residual": residual, "inner_iterations": iters})
123
124 with open("results.json", "w") as fp:
125 json.dump(results, fp, indent=2)
126 print(json.dumps(results, indent=2))
127
128
129if __name__ == "__main__":
130 main()