Hybrid ml physics publication offline

In [ ]:

"""Offline publication-quality run for ``hybrid_ml_physics_predictor.ipynb``.

Produces ``media/hybrid_ml_physics_publication.npz`` containing:
  - the trained MLP-residual parameters (Equinox-MLP weight pytree, flattened),
  - the three closed-loop simulation trajectories (pure-physics, pure-ML,
    hybrid) on the in-distribution test scenario,
  - the OOD test trajectories,
  - the training-loss curve.

The notebook in default ``publication`` mode loads from this NPZ and plots
in seconds. In ``fast`` mode (or when the NPZ is missing) the notebook
re-runs a smaller training step inline.

Run from the repo root::

    JAXONOMY_DISABLE_PROFILING=1 \
        python docs/examples/media/hybrid_ml_physics_publication_offline.py

Expected wall-time: ~30 s on a developer machine.

The wedge: in a production setting steps 1-3 would happen in a PyTorch /
TensorFlow training pipeline that the user already owns; the trained
checkpoint would land here as ``hybrid_residual_model.pt`` and the notebook
would load it via ``jaxonomy.library.PyTorch(file_name=...)``. Neither
torch nor tensorflow is installed in CI, so we ship the JAX/Equinox
equivalent and document the swap in the notebook.
"""
from __future__ import annotations

"""Offline publication-quality run for ``hybrid_ml_physics_predictor.ipynb``. Produces ``media/hybrid_ml_physics_publication.npz`` containing: - the trained MLP-residual parameters (Equinox-MLP weight pytree, flattened), - the three closed-loop simulation trajectories (pure-physics, pure-ML, hybrid) on the in-distribution test scenario, - the OOD test trajectories, - the training-loss curve. The notebook in default ``publication`` mode loads from this NPZ and plots in seconds. In ``fast`` mode (or when the NPZ is missing) the notebook re-runs a smaller training step inline. Run from the repo root:: JAXONOMY_DISABLE_PROFILING=1 \ python docs/examples/media/hybrid_ml_physics_publication_offline.py Expected wall-time: ~30 s on a developer machine. The wedge: in a production setting steps 1-3 would happen in a PyTorch / TensorFlow training pipeline that the user already owns; the trained checkpoint would land here as ``hybrid_residual_model.pt`` and the notebook would load it via ``jaxonomy.library.PyTorch(file_name=...)``. Neither torch nor tensorflow is installed in CI, so we ship the JAX/Equinox equivalent and document the swap in the notebook. """ from __future__ import annotations

In [ ]:

from pathlib import Path
import time

from pathlib import Path import time

In [ ]:

import numpy as np
import jax
import jax.numpy as jnp
import equinox as eqx
import optax

import numpy as np import jax import jax.numpy as jnp import equinox as eqx import optax

In [ ]:

HERE = Path(__file__).resolve().parent
NPZ_OUT = HERE / "hybrid_ml_physics_publication.npz"

HERE = Path(__file__).resolve().parent NPZ_OUT = HERE / "hybrid_ml_physics_publication.npz"

In [ ]:

# ---------------------------------------------------------------------
# Top-level constants — must match the notebook.
# ---------------------------------------------------------------------
G = 9.81           # m/s^2
L = 1.0            # m  (pendulum length)
M = 1.0            # kg
B_VISC = 0.10      # N·m·s   viscous damping (known to the engineer)
F_COUL = 0.30      # N·m     Coulomb friction (the residual the MLP must learn)
COUL_EPS = 0.05    # rad/s   sign smoothing for the "true" Coulomb term

# --------------------------------------------------------------------- # Top-level constants — must match the notebook. # --------------------------------------------------------------------- G = 9.81 # m/s^2 L = 1.0 # m (pendulum length) M = 1.0 # kg B_VISC = 0.10 # N·m·s viscous damping (known to the engineer) F_COUL = 0.30 # N·m Coulomb friction (the residual the MLP must learn) COUL_EPS = 0.05 # rad/s sign smoothing for the "true" Coulomb term

In [ ]:

DT_SIM = 0.01      # s       integration timestep
T_TRAIN = 30.0     # s       training-data trajectory length
T_TEST = 12.0      # s       in-distribution test length
T_OOD = 8.0        # s       out-of-distribution test length

DT_SIM = 0.01 # s integration timestep T_TRAIN = 30.0 # s training-data trajectory length T_TEST = 12.0 # s in-distribution test length T_OOD = 8.0 # s out-of-distribution test length

In [ ]:

N_HIDDEN = 32
N_LAYERS = 2
N_EPOCHS = 4000
LR = 5e-3
NOISE_THETA_DDOT = 0.05  # rad/s^2

N_HIDDEN = 32 N_LAYERS = 2 N_EPOCHS = 4000 LR = 5e-3 NOISE_THETA_DDOT = 0.05 # rad/s^2

In [ ]:

SEED = 7

SEED = 7

In [ ]:

# ---------------------------------------------------------------------
# Physics: true vs known plants. We integrate both with a hand-rolled
# fixed-step RK4 here in the offline script for portability and speed;
# the notebook does the same kinematics inside jaxonomy LeafSystems so
# the reader sees the framework patterns.
# ---------------------------------------------------------------------
def smooth_sign(v: jnp.ndarray, eps: float = COUL_EPS) -> jnp.ndarray:
    """tanh-based smoothed sign for differentiability through the Coulomb term."""
    return jnp.tanh(v / eps)

# --------------------------------------------------------------------- # Physics: true vs known plants. We integrate both with a hand-rolled # fixed-step RK4 here in the offline script for portability and speed; # the notebook does the same kinematics inside jaxonomy LeafSystems so # the reader sees the framework patterns. # --------------------------------------------------------------------- def smooth_sign(v: jnp.ndarray, eps: float = COUL_EPS) -> jnp.ndarray: """tanh-based smoothed sign for differentiability through the Coulomb term.""" return jnp.tanh(v / eps)

In [ ]:

def true_accel(theta: jnp.ndarray, omega: jnp.ndarray, tau: jnp.ndarray) -> jnp.ndarray:
    """True plant ddot{theta} including Coulomb friction (engineer doesn't know F_COUL)."""
    grav = -(G / L) * jnp.sin(theta)
    visc = -(B_VISC / (M * L * L)) * omega
    coul = -(F_COUL / (M * L * L)) * smooth_sign(omega)
    drive = tau / (M * L * L)
    return grav + visc + coul + drive

def true_accel(theta: jnp.ndarray, omega: jnp.ndarray, tau: jnp.ndarray) -> jnp.ndarray: """True plant ddot{theta} including Coulomb friction (engineer doesn't know F_COUL).""" grav = -(G / L) * jnp.sin(theta) visc = -(B_VISC / (M * L * L)) * omega coul = -(F_COUL / (M * L * L)) * smooth_sign(omega) drive = tau / (M * L * L) return grav + visc + coul + drive

In [ ]:

def known_accel(theta: jnp.ndarray, omega: jnp.ndarray, tau: jnp.ndarray) -> jnp.ndarray:
    """Known-physics ddot{theta}: gravity + viscous damping + torque (no Coulomb)."""
    grav = -(G / L) * jnp.sin(theta)
    visc = -(B_VISC / (M * L * L)) * omega
    drive = tau / (M * L * L)
    return grav + visc + drive

def known_accel(theta: jnp.ndarray, omega: jnp.ndarray, tau: jnp.ndarray) -> jnp.ndarray: """Known-physics ddot{theta}: gravity + viscous damping + torque (no Coulomb).""" grav = -(G / L) * jnp.sin(theta) visc = -(B_VISC / (M * L * L)) * omega drive = tau / (M * L * L) return grav + visc + drive

In [ ]:

def rk4_step(state, tau, accel_fn, dt):
    """Fixed-step RK4 for a (theta, omega) two-state pendulum under accel_fn."""
    def deriv(s):
        th, om = s
        return jnp.array([om, accel_fn(th, om, tau)])
    k1 = deriv(state)
    k2 = deriv(state + 0.5 * dt * k1)
    k3 = deriv(state + 0.5 * dt * k2)
    k4 = deriv(state + dt * k3)
    return state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)

def rk4_step(state, tau, accel_fn, dt): """Fixed-step RK4 for a (theta, omega) two-state pendulum under accel_fn.""" def deriv(s): th, om = s return jnp.array([om, accel_fn(th, om, tau)]) k1 = deriv(state) k2 = deriv(state + 0.5 * dt * k1) k3 = deriv(state + 0.5 * dt * k2) k4 = deriv(state + dt * k3) return state + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)

In [ ]:

def rollout(state0, taus, accel_fn, dt):
    """Roll out a fixed-step trajectory of length len(taus) under accel_fn."""
    def step(carry, tau):
        s_next = rk4_step(carry, tau, accel_fn, dt)
        return s_next, s_next
    _, traj = jax.lax.scan(step, state0, taus)
    return jnp.concatenate([state0[None, :], traj], axis=0)

def rollout(state0, taus, accel_fn, dt): """Roll out a fixed-step trajectory of length len(taus) under accel_fn.""" def step(carry, tau): s_next = rk4_step(carry, tau, accel_fn, dt) return s_next, s_next _, traj = jax.lax.scan(step, state0, taus) return jnp.concatenate([state0[None, :], traj], axis=0)

In [ ]:

# ---------------------------------------------------------------------
# Training data: drive the TRUE plant with a low-frequency random torque
# excitation, sample at 100 Hz, build (theta, omega) -> residual targets.
# ---------------------------------------------------------------------
def make_training_excitation(key, n):
    """Smoothed-noise torque excitation: bounded, low-frequency, mean-zero."""
    raw = jax.random.normal(key, (n,)) * 0.6  # ~0.6 N·m std
    # smooth with a 50-sample moving average for low-freq character
    kernel = jnp.ones(50) / 50.0
    smoothed = jnp.convolve(raw, kernel, mode="same")
    return smoothed

# --------------------------------------------------------------------- # Training data: drive the TRUE plant with a low-frequency random torque # excitation, sample at 100 Hz, build (theta, omega) -> residual targets. # --------------------------------------------------------------------- def make_training_excitation(key, n): """Smoothed-noise torque excitation: bounded, low-frequency, mean-zero.""" raw = jax.random.normal(key, (n,)) * 0.6 # ~0.6 N·m std # smooth with a 50-sample moving average for low-freq character kernel = jnp.ones(50) / 50.0 smoothed = jnp.convolve(raw, kernel, mode="same") return smoothed

In [ ]:

def build_training_set(key):
    """Generate a (theta, omega, theta_ddot_residual) supervised dataset."""
    n_steps = int(T_TRAIN / DT_SIM)
    k_drv, k_init, k_noise = jax.random.split(key, 3)
    taus = make_training_excitation(k_drv, n_steps)
    theta0 = jax.random.uniform(k_init, minval=-0.5, maxval=0.5)
    omega0 = 0.0
    state0 = jnp.array([theta0, omega0])
    traj = rollout(state0, taus, true_accel, DT_SIM)  # (n_steps+1, 2)
    thetas = traj[:-1, 0]
    omegas = traj[:-1, 1]
    # ground-truth residual = true - known, evaluated under the true (theta, omega)
    true_a = jax.vmap(true_accel)(thetas, omegas, taus)
    known_a = jax.vmap(known_accel)(thetas, omegas, taus)
    residual = true_a - known_a  # what the MLP must learn
    # add measurement noise on the residual target (the engineer would measure
    # theta_ddot via finite differences on noisy theta, which contaminates the
    # residual estimate)
    noise = jax.random.normal(k_noise, residual.shape) * NOISE_THETA_DDOT
    return thetas, omegas, residual + noise, taus

def build_training_set(key): """Generate a (theta, omega, theta_ddot_residual) supervised dataset.""" n_steps = int(T_TRAIN / DT_SIM) k_drv, k_init, k_noise = jax.random.split(key, 3) taus = make_training_excitation(k_drv, n_steps) theta0 = jax.random.uniform(k_init, minval=-0.5, maxval=0.5) omega0 = 0.0 state0 = jnp.array([theta0, omega0]) traj = rollout(state0, taus, true_accel, DT_SIM) # (n_steps+1, 2) thetas = traj[:-1, 0] omegas = traj[:-1, 1] # ground-truth residual = true - known, evaluated under the true (theta, omega) true_a = jax.vmap(true_accel)(thetas, omegas, taus) known_a = jax.vmap(known_accel)(thetas, omegas, taus) residual = true_a - known_a # what the MLP must learn # add measurement noise on the residual target (the engineer would measure # theta_ddot via finite differences on noisy theta, which contaminates the # residual estimate) noise = jax.random.normal(k_noise, residual.shape) * NOISE_THETA_DDOT return thetas, omegas, residual + noise, taus

In [ ]:

# ---------------------------------------------------------------------
# MLP residual model: a small Equinox MLP. The notebook documents how to
# swap this for a `jaxonomy.library.PyTorch` / `TensorFlow` predictor.
# ---------------------------------------------------------------------
def make_mlp(key):
    return eqx.nn.MLP(
        in_size=2,
        out_size=1,
        width_size=N_HIDDEN,
        depth=N_LAYERS,
        activation=jax.nn.tanh,
        key=key,
    )

# --------------------------------------------------------------------- # MLP residual model: a small Equinox MLP. The notebook documents how to # swap this for a `jaxonomy.library.PyTorch` / `TensorFlow` predictor. # --------------------------------------------------------------------- def make_mlp(key): return eqx.nn.MLP( in_size=2, out_size=1, width_size=N_HIDDEN, depth=N_LAYERS, activation=jax.nn.tanh, key=key, )

In [ ]:

def train(thetas, omegas, residuals, key):
    mlp = make_mlp(key)
    inputs = jnp.stack([thetas, omegas], axis=-1)         # (N, 2)
    targets = residuals[:, None]                          # (N, 1)

    params, static = eqx.partition(mlp, eqx.is_array)
    opt = optax.adam(LR)
    opt_state = opt.init(params)

    @jax.jit
    def loss_fn(params, x, y):
        m = eqx.combine(params, static)
        pred = jax.vmap(m)(x)
        return jnp.mean((pred - y) ** 2)

    @jax.jit
    def step(params, opt_state, x, y):
        loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
        upd, opt_state = opt.update(grads, opt_state)
        params = optax.apply_updates(params, upd)
        return params, opt_state, loss

    loss_hist = []
    for epoch in range(N_EPOCHS):
        params, opt_state, loss = step(params, opt_state, inputs, targets)
        if epoch % 100 == 0:
            loss_hist.append(float(loss))
    return params, static, jnp.array(loss_hist)

def train(thetas, omegas, residuals, key): mlp = make_mlp(key) inputs = jnp.stack([thetas, omegas], axis=-1) # (N, 2) targets = residuals[:, None] # (N, 1) params, static = eqx.partition(mlp, eqx.is_array) opt = optax.adam(LR) opt_state = opt.init(params) @jax.jit def loss_fn(params, x, y): m = eqx.combine(params, static) pred = jax.vmap(m)(x) return jnp.mean((pred - y) ** 2) @jax.jit def step(params, opt_state, x, y): loss, grads = jax.value_and_grad(loss_fn)(params, x, y) upd, opt_state = opt.update(grads, opt_state) params = optax.apply_updates(params, upd) return params, opt_state, loss loss_hist = [] for epoch in range(N_EPOCHS): params, opt_state, loss = step(params, opt_state, inputs, targets) if epoch % 100 == 0: loss_hist.append(float(loss)) return params, static, jnp.array(loss_hist)

In [ ]:

# ---------------------------------------------------------------------
# Closed-loop comparison: pure-physics vs pure-ML vs hybrid under a
# setpoint-tracking PD controller, on the SAME plant the engineer doesn't
# know exactly (the true plant). The three predictors differ in how they
# compute the simulated ddot{theta}; ground truth is always the true plant.
# ---------------------------------------------------------------------
KP, KD = 6.0, 1.5
THETA_REF = 0.8  # rad — a modest setpoint inside training distribution

# --------------------------------------------------------------------- # Closed-loop comparison: pure-physics vs pure-ML vs hybrid under a # setpoint-tracking PD controller, on the SAME plant the engineer doesn't # know exactly (the true plant). The three predictors differ in how they # compute the simulated ddot{theta}; ground truth is always the true plant. # --------------------------------------------------------------------- KP, KD = 6.0, 1.5 THETA_REF = 0.8 # rad — a modest setpoint inside training distribution

In [ ]:

def control_torque(theta, omega, theta_ref):
    return KP * (theta_ref - theta) - KD * omega

def control_torque(theta, omega, theta_ref): return KP * (theta_ref - theta) - KD * omega

In [ ]:

def closed_loop(state0, accel_fn, t_end, theta_ref=THETA_REF):
    n_steps = int(t_end / DT_SIM)
    def step(carry, _):
        th, om = carry
        tau = control_torque(th, om, theta_ref)
        s_next = rk4_step(carry, tau, accel_fn, DT_SIM)
        return s_next, jnp.array([s_next[0], s_next[1], tau])
    _, packed = jax.lax.scan(step, state0, jnp.arange(n_steps))
    times = jnp.arange(n_steps) * DT_SIM
    return times, packed  # packed[:, 0]=theta, [:,1]=omega, [:,2]=tau

def closed_loop(state0, accel_fn, t_end, theta_ref=THETA_REF): n_steps = int(t_end / DT_SIM) def step(carry, _): th, om = carry tau = control_torque(th, om, theta_ref) s_next = rk4_step(carry, tau, accel_fn, DT_SIM) return s_next, jnp.array([s_next[0], s_next[1], tau]) _, packed = jax.lax.scan(step, state0, jnp.arange(n_steps)) times = jnp.arange(n_steps) * DT_SIM return times, packed # packed[:, 0]=theta, [:,1]=omega, [:,2]=tau

In [ ]:

def hybrid_accel_factory(params, static):
    mlp = eqx.combine(params, static)
    def hybrid_accel(theta, omega, tau):
        x = jnp.array([theta, omega])
        residual = mlp(x)[0]
        return known_accel(theta, omega, tau) + residual
    return hybrid_accel

def hybrid_accel_factory(params, static): mlp = eqx.combine(params, static) def hybrid_accel(theta, omega, tau): x = jnp.array([theta, omega]) residual = mlp(x)[0] return known_accel(theta, omega, tau) + residual return hybrid_accel

In [ ]:

def pure_ml_accel_factory(params, static, taus_for_baseline=None):
    """Pure-ML predicts the FULL ddot{theta} as f(theta, omega), ignoring tau.

    This is the deliberately-naive comparison point: an ML model that learned
    the unforced dynamics. We train it the same way as the residual MLP but
    target = true ddot{theta} (under tau=0) so the model has no notion of the
    applied torque. (See `pure_ml_train` below.)
    """
    mlp = eqx.combine(params, static)
    def ml_accel(theta, omega, tau):
        x = jnp.array([theta, omega])
        # ML is unaware of tau by construction — that's the honest weakness
        # we want to demonstrate.
        return mlp(x)[0] + (tau / (M * L * L))
    return ml_accel

def pure_ml_accel_factory(params, static, taus_for_baseline=None): """Pure-ML predicts the FULL ddot{theta} as f(theta, omega), ignoring tau. This is the deliberately-naive comparison point: an ML model that learned the unforced dynamics. We train it the same way as the residual MLP but target = true ddot{theta} (under tau=0) so the model has no notion of the applied torque. (See `pure_ml_train` below.) """ mlp = eqx.combine(params, static) def ml_accel(theta, omega, tau): x = jnp.array([theta, omega]) # ML is unaware of tau by construction — that's the honest weakness # we want to demonstrate. return mlp(x)[0] + (tau / (M * L * L)) return ml_accel

In [ ]:

def pure_ml_train(key):
    """Train an MLP to predict ddot{theta} | tau = 0 under the true plant,
    using the same (theta, omega) coverage as the residual training set.
    """
    k1, k2 = jax.random.split(key)
    thetas, omegas, residuals_unused, _ = build_training_set(k1)
    # rebuild targets: true acceleration at zero torque (no driving)
    true_a_no_tau = jax.vmap(lambda th, om: true_accel(th, om, jnp.array(0.0)))(thetas, omegas)
    noise = jax.random.normal(k2, true_a_no_tau.shape) * NOISE_THETA_DDOT
    targets = true_a_no_tau + noise
    return train(thetas, omegas, targets, jax.random.PRNGKey(SEED + 1))

def pure_ml_train(key): """Train an MLP to predict ddot{theta} | tau = 0 under the true plant, using the same (theta, omega) coverage as the residual training set. """ k1, k2 = jax.random.split(key) thetas, omegas, residuals_unused, _ = build_training_set(k1) # rebuild targets: true acceleration at zero torque (no driving) true_a_no_tau = jax.vmap(lambda th, om: true_accel(th, om, jnp.array(0.0)))(thetas, omegas) noise = jax.random.normal(k2, true_a_no_tau.shape) * NOISE_THETA_DDOT targets = true_a_no_tau + noise return train(thetas, omegas, targets, jax.random.PRNGKey(SEED + 1))

In [ ]:

def main():
    t0 = time.time()
    key = jax.random.PRNGKey(SEED)
    k_train, k_pure_ml = jax.random.split(key)

    print("--- Stage 1: build training data ---")
    thetas, omegas, residuals, train_taus = build_training_set(k_train)
    print(f"  N samples = {thetas.shape[0]}, residual stats: mean={float(jnp.mean(residuals)):+.4f}, "
          f"std={float(jnp.std(residuals)):.4f}")

    print("--- Stage 2: train residual MLP ---")
    params, static, loss_hist = train(thetas, omegas, residuals, jax.random.PRNGKey(SEED + 2))
    final_loss = float(loss_hist[-1])
    print(f"  Trained {N_EPOCHS} epochs, final MSE = {final_loss:.4e}")

    print("--- Stage 3: train pure-ML baseline ---")
    pure_ml_params, pure_ml_static, pure_ml_loss_hist = pure_ml_train(k_pure_ml)
    print(f"  Pure-ML final MSE = {float(pure_ml_loss_hist[-1]):.4e}")

    print("--- Stage 4: closed-loop comparison (in-distribution) ---")
    state0 = jnp.array([0.0, 0.0])
    times_true, traj_true = closed_loop(state0, true_accel, T_TEST)
    times_phys, traj_phys = closed_loop(state0, known_accel, T_TEST)
    times_hyb, traj_hyb = closed_loop(state0, hybrid_accel_factory(params, static), T_TEST)
    times_ml, traj_ml = closed_loop(state0, pure_ml_accel_factory(pure_ml_params, pure_ml_static), T_TEST)

    def rmse(traj, ref):
        return float(jnp.sqrt(jnp.mean((traj[:, 0] - ref[:, 0]) ** 2)))

    print(f"  RMSE vs true:  pure-physics={rmse(traj_phys, traj_true):.4f} rad")
    print(f"                 pure-ML     ={rmse(traj_ml, traj_true):.4f} rad")
    print(f"                 hybrid      ={rmse(traj_hyb, traj_true):.4f} rad")

    print("--- Stage 5: OOD test (large initial angle, no setpoint) ---")
    state0_ood = jnp.array([2.5, 0.0])  # well outside the training range of (-0.5, 0.5)
    times_true_ood, traj_true_ood = closed_loop(state0_ood, true_accel, T_OOD, theta_ref=0.0)
    times_phys_ood, traj_phys_ood = closed_loop(state0_ood, known_accel, T_OOD, theta_ref=0.0)
    times_hyb_ood, traj_hyb_ood = closed_loop(state0_ood, hybrid_accel_factory(params, static), T_OOD, theta_ref=0.0)
    times_ml_ood, traj_ml_ood = closed_loop(state0_ood, pure_ml_accel_factory(pure_ml_params, pure_ml_static), T_OOD, theta_ref=0.0)
    print(f"  OOD RMSE vs true: pure-physics={rmse(traj_phys_ood, traj_true_ood):.4f} rad")
    print(f"                    pure-ML     ={rmse(traj_ml_ood, traj_true_ood):.4f} rad")
    print(f"                    hybrid      ={rmse(traj_hyb_ood, traj_true_ood):.4f} rad")

    print("--- Stage 6: residual grid for heatmap visualization ---")
    th_grid = jnp.linspace(-1.0, 1.0, 41)
    om_grid = jnp.linspace(-2.0, 2.0, 41)
    TH, OM = jnp.meshgrid(th_grid, om_grid, indexing="xy")
    pts = jnp.stack([TH.ravel(), OM.ravel()], axis=-1)
    mlp_eval = eqx.combine(params, static)
    pred_res = jax.vmap(mlp_eval)(pts).reshape(TH.shape)
    true_res = jax.vmap(lambda th, om: -(F_COUL / (M * L * L)) * smooth_sign(om))(TH.ravel(), OM.ravel()).reshape(TH.shape)

    wall = time.time() - t0
    print(f"Total wall-time = {wall:.2f} s; writing NPZ to {NPZ_OUT}")

    # Flatten the MLP parameter pytree for portable storage.
    flat_params = jax.tree.leaves(params)
    flat_pure_ml = jax.tree.leaves(pure_ml_params)

    np.savez(
        NPZ_OUT,
        # training-data set
        train_thetas=np.asarray(thetas),
        train_omegas=np.asarray(omegas),
        train_residuals=np.asarray(residuals),
        train_taus=np.asarray(train_taus),
        # loss curves
        loss_hist=np.asarray(loss_hist),
        pure_ml_loss_hist=np.asarray(pure_ml_loss_hist),
        # in-distribution closed-loop
        time_id=np.asarray(times_true),
        traj_true=np.asarray(traj_true),
        traj_phys=np.asarray(traj_phys),
        traj_hyb=np.asarray(traj_hyb),
        traj_ml=np.asarray(traj_ml),
        # OOD closed-loop
        time_ood=np.asarray(times_true_ood),
        traj_true_ood=np.asarray(traj_true_ood),
        traj_phys_ood=np.asarray(traj_phys_ood),
        traj_hyb_ood=np.asarray(traj_hyb_ood),
        traj_ml_ood=np.asarray(traj_ml_ood),
        # residual heatmap data
        th_grid=np.asarray(th_grid),
        om_grid=np.asarray(om_grid),
        pred_res=np.asarray(pred_res),
        true_res=np.asarray(true_res),
        # MLP params (pytree-flattened; the notebook reconstructs alongside identical static)
        n_mlp_leaves=len(flat_params),
        n_pure_ml_leaves=len(flat_pure_ml),
        wall_time=wall,
        seed=SEED,
        **{f"mlp_leaf_{i}": np.asarray(p) for i, p in enumerate(flat_params)},
        **{f"pure_ml_leaf_{i}": np.asarray(p) for i, p in enumerate(flat_pure_ml)},
    )
    print(f"Done. NPZ size = {NPZ_OUT.stat().st_size/1024:.1f} KiB.")

def main(): t0 = time.time() key = jax.random.PRNGKey(SEED) k_train, k_pure_ml = jax.random.split(key) print("--- Stage 1: build training data ---") thetas, omegas, residuals, train_taus = build_training_set(k_train) print(f" N samples = {thetas.shape[0]}, residual stats: mean={float(jnp.mean(residuals)):+.4f}, " f"std={float(jnp.std(residuals)):.4f}") print("--- Stage 2: train residual MLP ---") params, static, loss_hist = train(thetas, omegas, residuals, jax.random.PRNGKey(SEED + 2)) final_loss = float(loss_hist[-1]) print(f" Trained {N_EPOCHS} epochs, final MSE = {final_loss:.4e}") print("--- Stage 3: train pure-ML baseline ---") pure_ml_params, pure_ml_static, pure_ml_loss_hist = pure_ml_train(k_pure_ml) print(f" Pure-ML final MSE = {float(pure_ml_loss_hist[-1]):.4e}") print("--- Stage 4: closed-loop comparison (in-distribution) ---") state0 = jnp.array([0.0, 0.0]) times_true, traj_true = closed_loop(state0, true_accel, T_TEST) times_phys, traj_phys = closed_loop(state0, known_accel, T_TEST) times_hyb, traj_hyb = closed_loop(state0, hybrid_accel_factory(params, static), T_TEST) times_ml, traj_ml = closed_loop(state0, pure_ml_accel_factory(pure_ml_params, pure_ml_static), T_TEST) def rmse(traj, ref): return float(jnp.sqrt(jnp.mean((traj[:, 0] - ref[:, 0]) ** 2))) print(f" RMSE vs true: pure-physics={rmse(traj_phys, traj_true):.4f} rad") print(f" pure-ML ={rmse(traj_ml, traj_true):.4f} rad") print(f" hybrid ={rmse(traj_hyb, traj_true):.4f} rad") print("--- Stage 5: OOD test (large initial angle, no setpoint) ---") state0_ood = jnp.array([2.5, 0.0]) # well outside the training range of (-0.5, 0.5) times_true_ood, traj_true_ood = closed_loop(state0_ood, true_accel, T_OOD, theta_ref=0.0) times_phys_ood, traj_phys_ood = closed_loop(state0_ood, known_accel, T_OOD, theta_ref=0.0) times_hyb_ood, traj_hyb_ood = closed_loop(state0_ood, hybrid_accel_factory(params, static), T_OOD, theta_ref=0.0) times_ml_ood, traj_ml_ood = closed_loop(state0_ood, pure_ml_accel_factory(pure_ml_params, pure_ml_static), T_OOD, theta_ref=0.0) print(f" OOD RMSE vs true: pure-physics={rmse(traj_phys_ood, traj_true_ood):.4f} rad") print(f" pure-ML ={rmse(traj_ml_ood, traj_true_ood):.4f} rad") print(f" hybrid ={rmse(traj_hyb_ood, traj_true_ood):.4f} rad") print("--- Stage 6: residual grid for heatmap visualization ---") th_grid = jnp.linspace(-1.0, 1.0, 41) om_grid = jnp.linspace(-2.0, 2.0, 41) TH, OM = jnp.meshgrid(th_grid, om_grid, indexing="xy") pts = jnp.stack([TH.ravel(), OM.ravel()], axis=-1) mlp_eval = eqx.combine(params, static) pred_res = jax.vmap(mlp_eval)(pts).reshape(TH.shape) true_res = jax.vmap(lambda th, om: -(F_COUL / (M * L * L)) * smooth_sign(om))(TH.ravel(), OM.ravel()).reshape(TH.shape) wall = time.time() - t0 print(f"Total wall-time = {wall:.2f} s; writing NPZ to {NPZ_OUT}") # Flatten the MLP parameter pytree for portable storage. flat_params = jax.tree.leaves(params) flat_pure_ml = jax.tree.leaves(pure_ml_params) np.savez( NPZ_OUT, # training-data set train_thetas=np.asarray(thetas), train_omegas=np.asarray(omegas), train_residuals=np.asarray(residuals), train_taus=np.asarray(train_taus), # loss curves loss_hist=np.asarray(loss_hist), pure_ml_loss_hist=np.asarray(pure_ml_loss_hist), # in-distribution closed-loop time_id=np.asarray(times_true), traj_true=np.asarray(traj_true), traj_phys=np.asarray(traj_phys), traj_hyb=np.asarray(traj_hyb), traj_ml=np.asarray(traj_ml), # OOD closed-loop time_ood=np.asarray(times_true_ood), traj_true_ood=np.asarray(traj_true_ood), traj_phys_ood=np.asarray(traj_phys_ood), traj_hyb_ood=np.asarray(traj_hyb_ood), traj_ml_ood=np.asarray(traj_ml_ood), # residual heatmap data th_grid=np.asarray(th_grid), om_grid=np.asarray(om_grid), pred_res=np.asarray(pred_res), true_res=np.asarray(true_res), # MLP params (pytree-flattened; the notebook reconstructs alongside identical static) n_mlp_leaves=len(flat_params), n_pure_ml_leaves=len(flat_pure_ml), wall_time=wall, seed=SEED, **{f"mlp_leaf_{i}": np.asarray(p) for i, p in enumerate(flat_params)}, **{f"pure_ml_leaf_{i}": np.asarray(p) for i, p in enumerate(flat_pure_ml)}, ) print(f"Done. NPZ size = {NPZ_OUT.stat().st_size/1024:.1f} KiB.")

In [ ]:

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()