1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6# Jacobian-Frozen Stable Rollouts: small, reproducible numerical study.
  7SEED = 7
  8rng = np.random.default_rng(SEED)
  9
 10
 11def g(x, u):
 12    """A small nonlinear neural-state-space-like transition."""
 13    x = np.asarray(x, dtype=float)
 14    u = np.asarray(u, dtype=float)
 15    # Stable rotation plus smooth nonlinear state and input interactions.
 16    y0 = 0.93 * x[0] + 0.16 * x[1] + 0.08 * np.tanh(x[0] * x[1]) + 0.10 * u[0]
 17    y1 = -0.16 * x[0] + 0.93 * x[1] + 0.06 * np.tanh(x[0] ** 2) - 0.07 * u[0]
 18    return np.array([y0, y1])
 19
 20
 21def jacobian_xu(x, u):
 22    """Analytic A,B for g, avoiding any hidden finite-difference claim."""
 23    x0, x1 = x
 24    q = 1.0 - np.tanh(x0 * x1) ** 2
 25    A = np.array([
 26        [0.93 + 0.08 * x1 * q, 0.16 + 0.08 * x0 * q],
 27        [-0.16 + 0.12 * x0 * (1.0 - np.tanh(x0 ** 2) ** 2), 0.93],
 28    ])
 29    B = np.array([[0.10], [-0.07]])
 30    return A, B
 31
 32
 33def frozen_parameters(x, u):
 34    A, B = jacobian_xu(x, u)
 35    c = g(x, u) - A @ x - B @ u
 36    return A, B, c
 37
 38
 39def rollout_nonlinear(x0, us):
 40    xs = [np.asarray(x0, dtype=float).copy()]
 41    for u in us:
 42        xs.append(g(xs[-1], u))
 43    return np.asarray(xs)
 44
 45
 46def rollout_frozen(x0, us, clip_radius=None):
 47    x0 = np.asarray(x0, dtype=float)
 48    A, B, c = frozen_parameters(x0, us[0])
 49    rho = max(abs(np.linalg.eigvals(A)))
 50    if clip_radius is not None and rho > clip_radius:
 51        # Radial matrix scaling preserves the locally affine direction while
 52        # imposing the requested spectral-radius monitor threshold.
 53        A = A * (clip_radius / rho)
 54        rho = max(abs(np.linalg.eigvals(A)))
 55        c = g(x0, us[0]) - A @ x0 - B @ us[0]
 56    xs = [x0.copy()]
 57    for u in us:
 58        xs.append(A @ xs[-1] + B @ u + c)
 59    return np.asarray(xs), rho
 60
 61
 62def remainder_scaling():
 63    """Check ||g(x+d)-g(x)-Jd|| = O(||d||^2) around a nonzero point."""
 64    x = np.array([0.42, -0.31]); u = np.array([0.13])
 65    A, B = jacobian_xu(x, u)
 66    direction = np.array([0.8, -0.6]); direction /= np.linalg.norm(direction)
 67    eps = np.array([0.2, 0.1, 0.05, 0.025, 0.0125])
 68    errors = []
 69    for e in eps:
 70        d = e * direction
 71        errors.append(np.linalg.norm(g(x + d, u) - g(x, u) - A @ d))
 72    slope = np.polyfit(np.log(eps), np.log(errors), 1)[0]
 73    ratios = (np.asarray(errors[:-1]) / np.asarray(errors[1:])).tolist()
 74    return {"eps": eps.tolist(), "errors": np.asarray(errors).tolist(),
 75            "loglog_slope": float(slope), "halve_error_ratios": ratios}
 76
 77
 78def stability_boundary():
 79    """Linear scalar control check: x_{t+1}=a*x_t has boundary |rho|=1."""
 80    rows = []
 81    for a in [0.80, 0.95, 0.99, 1.00, 1.01, 1.10]:
 82        x = 1e-3
 83        for _ in range(40):
 84            x = a * x
 85        rows.append({"a": a, "rho": abs(a), "growth_factor": abs(x / 1e-3),
 86                     "decays": bool(abs(x) < 1e-3)})
 87    return rows
 88
 89
 90def mini_experiment():
 91    # Constant but nonzero input makes this a genuine state-space rollout.
 92    H = 30
 93    us = np.tile(np.array([[0.12]]), (H, 1))
 94    records = []
 95    for radius in [0.03, 0.10, 0.25, 0.50, 0.90]:
 96        x0 = radius * np.array([0.8, -0.6])
 97        truth = rollout_nonlinear(x0, us)
 98        frozen, rho = rollout_frozen(x0, us)
 99        clipped, clipped_rho = rollout_frozen(x0, us, clip_radius=0.98)
100        err = float(np.sqrt(np.mean((frozen - truth) ** 2)))
101        clip_err = float(np.sqrt(np.mean((clipped - truth) ** 2)))
102        records.append({"radius": radius, "rho_at_start": float(rho),
103                        "nonlinear_final_norm": float(np.linalg.norm(truth[-1])),
104                        "frozen_final_norm": float(np.linalg.norm(frozen[-1])),
105                        "frozen_rmse": err, "clipped_frozen_rmse": clip_err,
106                        "frozen_max_norm": float(np.max(np.linalg.norm(frozen, axis=1))),
107                        "nonlinear_max_norm": float(np.max(np.linalg.norm(truth, axis=1)))})
108    # A rough operation count illustrates why freezing can be cheaper: one
109    # nonlinear evaluation per step versus one Jacobian plus affine matvecs.
110    return records
111
112
113def main():
114    out = {"seed": SEED, "remainder_scaling": remainder_scaling(),
115           "stability_boundary": stability_boundary(),
116           "mini_experiment": mini_experiment(),
117           "notes": "Frozen rollout uses exact local analytic Jacobian; clipping rescales A and recomputes c."}
118    Path("results.json").write_text(json.dumps(out, indent=2))
119    print(json.dumps(out, indent=2))
120
121
122if __name__ == "__main__":
123    main()