Dae projection pendulum publication offline

In [ ]:

"""Offline publication run for ``dae_projection_pendulum.ipynb``.

Produces ``media/dae_projection_pendulum_publication.npz``.  The notebook
loads the NPZ in publication mode (default) and falls back to a live
re-run on a 5-minute horizon in fast mode if the NPZ is missing.

Runtime on a developer CPU: ~12 s for three 1-hour simulations + a 1 hr
SSP autodiff gradient run.  The notebook's fast mode runs the same
configurations on a 300 s horizon in ~3 s.  This script exists so the
notebook execution stays sub-minute under ``nbconvert``.

Saved arrays:
  config_labels         (3,) U10     -- "Baseline", "Baumgarte", "SSP"
  walls                 (3,)         -- wall-clock seconds per config
  final_resid           (3,)         -- ||f_a||_inf at t=3600 s
  max_residual          (3,)         -- max trace residual across run
  trace_time            (3, N_trace) -- per-major-step times
  trace_residual        (3, N_trace) -- per-major-step ||f_a||_inf
  energy_mean_first     (3,)         -- mean E over first 100 s
  energy_mean_last      (3,)         -- mean E over last 100 s
  geom_max              (3,)         -- max(|x^2+y^2-L^2|) across run
  geom_at_3600          (3,)         -- |x^2+y^2-L^2| at t = 3600 s
  T_END, DT_MAJOR, rtol, atol, baumgarte_beta, ssp_tol -- scalar metadata
  placeholder_flag       () bool     -- False; these are real numbers
"""

"""Offline publication run for ``dae_projection_pendulum.ipynb``. Produces ``media/dae_projection_pendulum_publication.npz``. The notebook loads the NPZ in publication mode (default) and falls back to a live re-run on a 5-minute horizon in fast mode if the NPZ is missing. Runtime on a developer CPU: ~12 s for three 1-hour simulations + a 1 hr SSP autodiff gradient run. The notebook's fast mode runs the same configurations on a 300 s horizon in ~3 s. This script exists so the notebook execution stays sub-minute under ``nbconvert``. Saved arrays: config_labels (3,) U10 -- "Baseline", "Baumgarte", "SSP" walls (3,) -- wall-clock seconds per config final_resid (3,) -- ||f_a||_inf at t=3600 s max_residual (3,) -- max trace residual across run trace_time (3, N_trace) -- per-major-step times trace_residual (3, N_trace) -- per-major-step ||f_a||_inf energy_mean_first (3,) -- mean E over first 100 s energy_mean_last (3,) -- mean E over last 100 s geom_max (3,) -- max(|x^2+y^2-L^2|) across run geom_at_3600 (3,) -- |x^2+y^2-L^2| at t = 3600 s T_END, DT_MAJOR, rtol, atol, baumgarte_beta, ssp_tol -- scalar metadata placeholder_flag () bool -- False; these are real numbers """

In [ ]:

from __future__ import annotations

from __future__ import annotations

In [ ]:

import time
from pathlib import Path

import time from pathlib import Path

In [ ]:

import numpy as np
import jax.numpy as jnp

import numpy as np import jax.numpy as jnp

In [ ]:

import jaxonomy
from jaxonomy.simulation.dae_drift import constraint_residual_norm

import jaxonomy from jaxonomy.simulation.dae_drift import constraint_residual_norm

In [ ]:

HERE = Path(__file__).resolve().parent
OUT_NPZ = HERE / "dae_projection_pendulum_publication.npz"

HERE = Path(__file__).resolve().parent OUT_NPZ = HERE / "dae_projection_pendulum_publication.npz"

In [ ]:

# Same index-reduced index-2 pendulum fixture used by the T-113 test
# suite (test/simulation/test_dae_projection.py, test_t_113_phase4_*).
class PlanarPendulum(jaxonomy.LeafSystem):
    """Planar pendulum on a rigid massless link of length L.

    State layout (9 entries):
      x[0] = vx       (differential; M = 1)
      x[1] = x_pos    (differential; M = 1)
      z[0] = vy       (algebraic; M = 0)
      z[1] ... z[6]   (algebraic auxiliaries from index reduction)
      z[2] = y_pos    (algebraic)
    """

    def __init__(self, L: float = 1.0, g0: float = 9.8, name: str | None = None):
        super().__init__(name=name)
        # Index-reduced consistent IC: bob at (sqrt(3)/2, -1/2), at rest.
        x0 = np.array(
            [0.0, 0.8660254037844386, 0.0, -4.9, -0.5,
             -4.243524478543744, -7.35, -7.35, 0.0]
        )
        self.declare_dynamic_parameter("L", L)
        self.declare_dynamic_parameter("g0", g0)
        M = np.concatenate([np.ones(2), np.zeros(7)])
        self.declare_continuous_state(
            default_value=x0, mass_matrix=M, ode=self.ode,
        )
        self.declare_continuous_state_output(name="x")

    def ode(self, time, state, **parameters):
        L, g0 = parameters["L"], parameters["g0"]
        x = state.continuous_state[:2]
        z = state.continuous_state[2:]
        f = jnp.array([z[3], x[0]])
        g = jnp.array([
            -(L ** 2) + x[1] ** 2 + z[2] ** 2,
            2 * z[0] * z[2] + 2 * x[1] * x[0],
            z[0] - z[6],
            2 * z[3] * x[1] + 2 * z[4] * z[2] + 2 * z[0] ** 2 + 2 * x[0] ** 2,
            z[4] - z[5],
            z[5] + g0 - z[1] * z[2],
            -z[1] * x[1] + z[3],
        ])
        return jnp.concatenate([f, g])

# Same index-reduced index-2 pendulum fixture used by the T-113 test # suite (test/simulation/test_dae_projection.py, test_t_113_phase4_*). class PlanarPendulum(jaxonomy.LeafSystem): """Planar pendulum on a rigid massless link of length L. State layout (9 entries): x[0] = vx (differential; M = 1) x[1] = x_pos (differential; M = 1) z[0] = vy (algebraic; M = 0) z[1] ... z[6] (algebraic auxiliaries from index reduction) z[2] = y_pos (algebraic) """ def __init__(self, L: float = 1.0, g0: float = 9.8, name: str | None = None): super().__init__(name=name) # Index-reduced consistent IC: bob at (sqrt(3)/2, -1/2), at rest. x0 = np.array( [0.0, 0.8660254037844386, 0.0, -4.9, -0.5, -4.243524478543744, -7.35, -7.35, 0.0] ) self.declare_dynamic_parameter("L", L) self.declare_dynamic_parameter("g0", g0) M = np.concatenate([np.ones(2), np.zeros(7)]) self.declare_continuous_state( default_value=x0, mass_matrix=M, ode=self.ode, ) self.declare_continuous_state_output(name="x") def ode(self, time, state, **parameters): L, g0 = parameters["L"], parameters["g0"] x = state.continuous_state[:2] z = state.continuous_state[2:] f = jnp.array([z[3], x[0]]) g = jnp.array([ -(L ** 2) + x[1] ** 2 + z[2] ** 2, 2 * z[0] * z[2] + 2 * x[1] * x[0], z[0] - z[6], 2 * z[3] * x[1] + 2 * z[4] * z[2] + 2 * z[0] ** 2 + 2 * x[0] ** 2, z[4] - z[5], z[5] + g0 - z[1] * z[2], -z[1] * x[1] + z[3], ]) return jnp.concatenate([f, g])

In [ ]:

# Horizon + solver configuration — kept in sync with the notebook.
T_END = 3600.0      # 1 hour of model time
DT_MAJOR = 2.0      # major-step length: bounds the drift trace cadence
RTOL = 1e-5
ATOL = 1e-7
BAUMGARTE_BETA = 0.05   # small; high-index DAE stalls BDF Newton if too large
SSP_TOL = 1e-9
SSP_ITERS = 4
G0 = 9.8
L = 1.0

# Horizon + solver configuration — kept in sync with the notebook. T_END = 3600.0 # 1 hour of model time DT_MAJOR = 2.0 # major-step length: bounds the drift trace cadence RTOL = 1e-5 ATOL = 1e-7 BAUMGARTE_BETA = 0.05 # small; high-index DAE stalls BDF Newton if too large SSP_TOL = 1e-9 SSP_ITERS = 4 G0 = 9.8 L = 1.0

In [ ]:

def _run(config_name: str, *, baumgarte_beta=None, projection=False):
    """Run one 1-hour pendulum simulation, return everything we need."""
    base = dict(
        math_backend="jax",
        ode_solver_method="bdf",
        rtol=RTOL,
        atol=ATOL,
        record_dae_drift=True,
        max_major_step_length=DT_MAJOR,
        max_major_steps=int(T_END / DT_MAJOR) + 10,
        buffer_length=200_000,
    )
    if baumgarte_beta is not None:
        base.update(baumgarte_alpha=None, baumgarte_beta=baumgarte_beta)
    if projection:
        base.update(
            dae_projection_enabled=True,
            dae_projection_tol=SSP_TOL,
            dae_projection_max_iter=SSP_ITERS,
        )

    model = PlanarPendulum()
    ctx = model.create_context()
    rec = {"state": model.output_ports[0]}
    opts = jaxonomy.SimulatorOptions(**base)

    t_wall = time.time()
    res = jaxonomy.simulate(
        model, ctx, (0.0, T_END), options=opts, recorded_signals=rec,
    )
    wall = time.time() - t_wall

    s = np.asarray(res.outputs["state"])
    vx, xp, vy, yp = s[:, 0], s[:, 1], s[:, 2], s[:, 4]
    E = 0.5 * (vx ** 2 + vy ** 2) + G0 * yp
    geom = xp ** 2 + yp ** 2 - L ** 2

    trace = res.dae_drift_trace
    fin_resid = float(constraint_residual_norm(model, res.context))

    print(
        f"  {config_name:12s} wall={wall:5.2f}s  final_resid={fin_resid:.3e}  "
        f"max_trace={float(np.max(trace['residual'])):.3e}  "
        f"E_range={np.max(E) - np.min(E):.4e}  |geom|max={np.max(np.abs(geom)):.3e}"
    )

    return dict(
        wall=wall,
        final_resid=fin_resid,
        max_residual=float(np.max(trace["residual"])),
        trace_time=np.asarray(trace["time"]),
        trace_residual=np.asarray(trace["residual"]),
        energy_mean_first=float(np.mean(E[:200])),
        energy_mean_last=float(np.mean(E[-200:])),
        geom_max=float(np.max(np.abs(geom))),
        geom_at_3600=float(np.abs(geom[-1])),
    )

def _run(config_name: str, *, baumgarte_beta=None, projection=False): """Run one 1-hour pendulum simulation, return everything we need.""" base = dict( math_backend="jax", ode_solver_method="bdf", rtol=RTOL, atol=ATOL, record_dae_drift=True, max_major_step_length=DT_MAJOR, max_major_steps=int(T_END / DT_MAJOR) + 10, buffer_length=200_000, ) if baumgarte_beta is not None: base.update(baumgarte_alpha=None, baumgarte_beta=baumgarte_beta) if projection: base.update( dae_projection_enabled=True, dae_projection_tol=SSP_TOL, dae_projection_max_iter=SSP_ITERS, ) model = PlanarPendulum() ctx = model.create_context() rec = {"state": model.output_ports[0]} opts = jaxonomy.SimulatorOptions(**base) t_wall = time.time() res = jaxonomy.simulate( model, ctx, (0.0, T_END), options=opts, recorded_signals=rec, ) wall = time.time() - t_wall s = np.asarray(res.outputs["state"]) vx, xp, vy, yp = s[:, 0], s[:, 1], s[:, 2], s[:, 4] E = 0.5 * (vx ** 2 + vy ** 2) + G0 * yp geom = xp ** 2 + yp ** 2 - L ** 2 trace = res.dae_drift_trace fin_resid = float(constraint_residual_norm(model, res.context)) print( f" {config_name:12s} wall={wall:5.2f}s final_resid={fin_resid:.3e} " f"max_trace={float(np.max(trace['residual'])):.3e} " f"E_range={np.max(E) - np.min(E):.4e} |geom|max={np.max(np.abs(geom)):.3e}" ) return dict( wall=wall, final_resid=fin_resid, max_residual=float(np.max(trace["residual"])), trace_time=np.asarray(trace["time"]), trace_residual=np.asarray(trace["residual"]), energy_mean_first=float(np.mean(E[:200])), energy_mean_last=float(np.mean(E[-200:])), geom_max=float(np.max(np.abs(geom))), geom_at_3600=float(np.abs(geom[-1])), )

In [ ]:

def main() -> None:
    print(f"Publication run: 1-hour pendulum with three drift-control configs")
    print(f"  T_END={T_END} s, DT_MAJOR={DT_MAJOR} s, rtol={RTOL}, atol={ATOL}")
    print(f"  Baumgarte beta={BAUMGARTE_BETA}, SSP tol={SSP_TOL}")
    print()

    results = {
        "Baseline": _run("Baseline"),
        "Baumgarte": _run("Baumgarte", baumgarte_beta=BAUMGARTE_BETA),
        "SSP": _run("SSP", projection=True),
    }

    # Trace arrays may differ in length across configs (one per major
    # step); pad to a common length so they fit in a single 2-D ndarray.
    config_labels = np.array(list(results.keys()), dtype="U10")
    trace_len = max(len(r["trace_time"]) for r in results.values())
    trace_time = np.full((len(results), trace_len), np.nan)
    trace_resid = np.full((len(results), trace_len), np.nan)
    for i, r in enumerate(results.values()):
        n = len(r["trace_time"])
        trace_time[i, :n] = r["trace_time"]
        trace_resid[i, :n] = r["trace_residual"]

    walls = np.array([r["wall"] for r in results.values()])
    final_resid = np.array([r["final_resid"] for r in results.values()])
    max_residual = np.array([r["max_residual"] for r in results.values()])
    energy_mean_first = np.array(
        [r["energy_mean_first"] for r in results.values()]
    )
    energy_mean_last = np.array(
        [r["energy_mean_last"] for r in results.values()]
    )
    geom_max = np.array([r["geom_max"] for r in results.values()])
    geom_at_3600 = np.array([r["geom_at_3600"] for r in results.values()])

    np.savez(
        OUT_NPZ,
        config_labels=config_labels,
        walls=walls,
        final_resid=final_resid,
        max_residual=max_residual,
        trace_time=trace_time,
        trace_residual=trace_resid,
        energy_mean_first=energy_mean_first,
        energy_mean_last=energy_mean_last,
        geom_max=geom_max,
        geom_at_3600=geom_at_3600,
        T_END=T_END,
        DT_MAJOR=DT_MAJOR,
        rtol=RTOL,
        atol=ATOL,
        baumgarte_beta=BAUMGARTE_BETA,
        ssp_tol=SSP_TOL,
        placeholder_flag=False,
    )
    print()
    print(f"Wrote {OUT_NPZ}")

def main() -> None: print(f"Publication run: 1-hour pendulum with three drift-control configs") print(f" T_END={T_END} s, DT_MAJOR={DT_MAJOR} s, rtol={RTOL}, atol={ATOL}") print(f" Baumgarte beta={BAUMGARTE_BETA}, SSP tol={SSP_TOL}") print() results = { "Baseline": _run("Baseline"), "Baumgarte": _run("Baumgarte", baumgarte_beta=BAUMGARTE_BETA), "SSP": _run("SSP", projection=True), } # Trace arrays may differ in length across configs (one per major # step); pad to a common length so they fit in a single 2-D ndarray. config_labels = np.array(list(results.keys()), dtype="U10") trace_len = max(len(r["trace_time"]) for r in results.values()) trace_time = np.full((len(results), trace_len), np.nan) trace_resid = np.full((len(results), trace_len), np.nan) for i, r in enumerate(results.values()): n = len(r["trace_time"]) trace_time[i, :n] = r["trace_time"] trace_resid[i, :n] = r["trace_residual"] walls = np.array([r["wall"] for r in results.values()]) final_resid = np.array([r["final_resid"] for r in results.values()]) max_residual = np.array([r["max_residual"] for r in results.values()]) energy_mean_first = np.array( [r["energy_mean_first"] for r in results.values()] ) energy_mean_last = np.array( [r["energy_mean_last"] for r in results.values()] ) geom_max = np.array([r["geom_max"] for r in results.values()]) geom_at_3600 = np.array([r["geom_at_3600"] for r in results.values()]) np.savez( OUT_NPZ, config_labels=config_labels, walls=walls, final_resid=final_resid, max_residual=max_residual, trace_time=trace_time, trace_residual=trace_resid, energy_mean_first=energy_mean_first, energy_mean_last=energy_mean_last, geom_max=geom_max, geom_at_3600=geom_at_3600, T_END=T_END, DT_MAJOR=DT_MAJOR, rtol=RTOL, atol=ATOL, baumgarte_beta=BAUMGARTE_BETA, ssp_tol=SSP_TOL, placeholder_flag=False, ) print() print(f"Wrote {OUT_NPZ}")

In [ ]:

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()