Lyapunov-sign-preserving neural time stepping / lyapunov_controller.py

Mechanism confirmed, baseline not beaten

Raw ⬇ ZIP
  1"""Lyapunov-sign-preserving stepping MVP and quantitative verification."""
  2import json, math, os
  3import numpy as np
  4import torch
  5
  6SEED = 7
  7np.random.seed(SEED)
  8torch.manual_seed(SEED)
  9
 10
 11def euler_step(z, drift, h):
 12    return z + h * drift(z)
 13
 14
 15def semi_implicit_linear_step(z, A, h):
 16    return np.linalg.solve(np.eye(A.shape[0]) - h * A, z)
 17
 18
 19def exponent_from_jacobians(jacobians, h):
 20    q = np.ones(jacobians[0].shape[0] if np.ndim(jacobians[0]) > 0 else 1)
 21    q /= np.linalg.norm(q)
 22    total = 0.0
 23    for J in jacobians:
 24        v = (float(J) * q) if np.ndim(J) == 0 else np.asarray(J) @ q
 25        norm = np.linalg.norm(v)
 26        total += math.log(max(norm, 1e-15))
 27        q = v / max(norm, 1e-15)
 28    return total / (len(jacobians) * h)
 29
 30
 31def estimate_with_torch(drift, z, h, method="euler", A=None):
 32    """One exact JVP of a discrete map, suitable for a neural drift."""
 33    z = z.detach().clone().requires_grad_(True)
 34    q = torch.randn_like(z)
 35    q = q / q.norm()
 36    if method == "euler":
 37        step = lambda x: x + h * drift(x)
 38    elif method == "semi_implicit" and A is not None:
 39        eye = torch.eye(len(z), dtype=z.dtype, device=z.device)
 40        step = lambda x: torch.linalg.solve(eye - h * A, x)
 41    else:
 42        raise ValueError(method)
 43    out, jvp = torch.autograd.functional.jvp(step, z, q)
 44    return out.detach(), float(torch.log(jvp.norm().clamp_min(1e-15)).item()), q.detach()
 45
 46
 47class SignPreservingController:
 48    """Halve h when a nominally stable continuous exponent is estimated positive."""
 49    def __init__(self, h, stable=True, tolerance=0.0, min_h=1e-5):
 50        self.h = float(h)
 51        self.stable = stable
 52        self.tolerance = tolerance
 53        self.min_h = min_h
 54        self.history = []
 55
 56    def decide(self, lambda_hat):
 57        self.history.append((self.h, float(lambda_hat)))
 58        if self.stable and lambda_hat > self.tolerance and self.h > self.min_h:
 59            self.h *= 0.5
 60            return "halve"
 61        return "keep"
 62
 63
 64def scalar_results(a=1.0, hs=None, n=200000):
 65    """Exact scalar contraction dz=-a z; continuous lambda=-a.
 66
 67    Euler lambda_h=log(abs(1-a h))/h and implicit lambda_h=-log(1+a h)/h.
 68    """
 69    if hs is None:
 70        hs = np.array([0.25, 0.5, 1.0, 1.5, 1.9, 2.01, 2.5, 3.0])
 71    rows = []
 72    for h in hs:
 73        le = math.log(abs(1-a*h))/h if abs(1-a*h) > 0 else -math.inf
 74        li = -math.log(1+a*h)/h
 75        rows.append({"h": float(h), "euler": le, "implicit": li,
 76                     "euler_sign": int(np.sign(le)), "implicit_sign": int(np.sign(li)),
 77                     "euler_error": abs(le+a), "implicit_error": abs(li+a)})
 78    return rows
 79
 80
 81def matrix_results(hs=None):
 82    # Stable damped oscillator: continuous top exponent is real part eigenvalue=-0.2.
 83    A = np.array([[-0.2, -2.0], [2.0, -0.2]])
 84    true = -0.2
 85    if hs is None: hs = np.array([0.05, 0.1, 0.2, 0.4, 0.8])
 86    rows=[]
 87    for h in hs:
 88        Je = np.eye(2)+h*A
 89        Ji = np.linalg.inv(np.eye(2)-h*A)
 90        # Long-run exponent is log spectral radius / h; agrees with QR estimate.
 91        le = math.log(max(abs(np.linalg.eigvals(Je))))/h
 92        li = math.log(max(abs(np.linalg.eigvals(Ji))))/h
 93        rows.append({"h":float(h), "euler":float(le), "implicit":float(li),
 94                     "euler_error":abs(le-true), "implicit_error":abs(li-true),
 95                     "euler_sign":int(np.sign(le)), "implicit_sign":int(np.sign(li))})
 96    return true, rows
 97
 98
 99def main():
100    true, matrix = matrix_results()
101    scalar = scalar_results()
102    # Prediction 1: scalar Euler changes sign exactly at h*=2/a.
103    crossing = next(r["h"] for r in scalar if r["euler"] > 0)
104    boundary_sweep = []
105    for a in [0.5, 1.0, 2.0, 4.0]:
106        predicted = 2.0/a
107        grid = np.linspace(0.8*predicted, 1.2*predicted, 401)
108        observed = next(float(x) for x in grid if math.log(abs(1-a*x))/x > 0)
109        boundary_sweep.append({"a":a, "predicted_h_star":predicted,
110                               "observed_grid_h_star":observed,
111                               "relative_grid_error":abs(observed-predicted)/predicted})
112    # Empirical order from small scalar h: error ratio approximately 2 for h -> h/2.
113    small = scalar_results(hs=np.array([0.4, 0.2, 0.1, 0.05]))
114    ratios = [small[i]["euler_error"] / small[i+1]["euler_error"] for i in range(3)]
115    controller = SignPreservingController(2.5, stable=True)
116    decisions=[]
117    h=controller.h
118    while True:
119        lam=math.log(abs(1-h))/h
120        action=controller.decide(lam)
121        decisions.append({"h":h,"lambda_hat":lam,"action":action})
122        if action == "keep": break
123        h=controller.h
124    # Prediction 3: exact JVP should equal the discrete Jacobian action.
125    z = torch.tensor([0.7, -0.2], dtype=torch.float64)
126    A_t = torch.tensor([[-1.0, 0.0], [0.0, -2.0]], dtype=torch.float64)
127    drift = lambda x: A_t @ x
128    _, jvp_log, q_used = estimate_with_torch(drift, z, 0.3, method="euler")
129    exact_jvp = (torch.eye(2, dtype=torch.float64) + 0.3*A_t) @ q_used
130    expected_jvp_log = float(torch.log(exact_jvp.norm()).item())
131    out={"continuous_scalar":-1.0, "predicted_euler_boundary":2.0,
132         "observed_first_positive_grid_h":crossing, "boundary_parameter_sweep":boundary_sweep,
133         "scalar":scalar,
134         "jvp_log_error_for_known_linear_map":abs(jvp_log-expected_jvp_log),
135         "error_ratios_h_to_h2":ratios, "oscillator_true":true,
136         "oscillator":matrix, "controller":decisions}
137    with open("results.json","w") as f: json.dump(out,f,indent=2)
138    print(json.dumps(out, indent=2))
139
140if __name__ == "__main__": main()