Feasible High-Order Neural ODE Solver / experiment.py

Mechanism failed

Raw ⬇ ZIP
  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()