Second-order systems

In [1]:

import jax
import jax.numpy as np
import jax.random as jrandom
import matplotlib.pyplot as plt

from ode_filters.filters import ekf1_sqr_loop, rts_sqr_smoother_loop
from ode_filters.measurement import (
    Measurement,
    SecondOrderODEInformation,
    SecondOrderODEInformationWithHidden,
)
from ode_filters.priors import IWP, JointPrior, taylor_mode_initialization

jit_filter = jax.jit(ekf1_sqr_loop, static_argnums=(2, 3, 4, 5))
jit_smoother = jax.jit(rts_sqr_smoother_loop, static_argnums=(5,))

# Damped harmonic oscillator parameters
omega, zeta = 1.0, 0.1
gamma = 2 * zeta * omega
omega_d = np.sqrt(1 - zeta**2) * omega  # damped frequency
tspan, N = (0, 30), 100
ts = np.linspace(tspan[0], tspan[1], N + 1)

import jax import jax.numpy as np import jax.random as jrandom import matplotlib.pyplot as plt from ode_filters.filters import ekf1_sqr_loop, rts_sqr_smoother_loop from ode_filters.measurement import ( Measurement, SecondOrderODEInformation, SecondOrderODEInformationWithHidden, ) from ode_filters.priors import IWP, JointPrior, taylor_mode_initialization jit_filter = jax.jit(ekf1_sqr_loop, static_argnums=(2, 3, 4, 5)) jit_smoother = jax.jit(rts_sqr_smoother_loop, static_argnums=(5,)) # Damped harmonic oscillator parameters omega, zeta = 1.0, 0.1 gamma = 2 * zeta * omega omega_d = np.sqrt(1 - zeta**2) * omega # damped frequency tspan, N = (0, 30), 100 ts = np.linspace(tspan[0], tspan[1], N + 1)

In [2]:

# Basic second-order ODE: d^2x/dt^2 = -omega^2 * x - gamma * dx/dt
def vf(x, dx, *, t):
    return -(omega**2) * x - gamma * dx


x0, dx0 = np.array([1.0]), np.array([0.0])
prior = IWP(q=2, d=1, Xi=1 * np.eye(1))
mu_0, Sigma_0_sqr = taylor_mode_initialization(vf, (x0, dx0), q=2, order=2)
measure = SecondOrderODEInformation(vf, prior.E0, prior.E1, prior.E2)

# Basic second-order ODE: d^2x/dt^2 = -omega^2 * x - gamma * dx/dt def vf(x, dx, *, t): return -(omega**2) * x - gamma * dx x0, dx0 = np.array([1.0]), np.array([0.0]) prior = IWP(q=2, d=1, Xi=1 * np.eye(1)) mu_0, Sigma_0_sqr = taylor_mode_initialization(vf, (x0, dx0), q=2, order=2) measure = SecondOrderODEInformation(vf, prior.E0, prior.E1, prior.E2)

In [3]:

# Run filter and smoother
m_seq, P_sqr, _, _, G_back, d_back, P_back_sqr, *_ = jit_filter(
    mu_0, Sigma_0_sqr, prior, measure, tspan, N
)
m_smooth, P_smooth_sqr = jit_smoother(
    m_seq[-1], P_sqr[-1], np.array(G_back), np.array(d_back), np.array(P_back_sqr), N
)
m_smooth, P_smooth_sqr = np.array(m_smooth), np.array(P_smooth_sqr)

# Run filter and smoother m_seq, P_sqr, _, _, G_back, d_back, P_back_sqr, *_ = jit_filter( mu_0, Sigma_0_sqr, prior, measure, tspan, N ) m_smooth, P_smooth_sqr = jit_smoother( m_seq[-1], P_sqr[-1], np.array(G_back), np.array(d_back), np.array(P_back_sqr), N ) m_smooth, P_smooth_sqr = np.array(m_smooth), np.array(P_smooth_sqr)

In [4]:

# Plot smoothed estimate vs true solution
x_true = np.exp(-zeta * omega * ts) * (
    np.cos(omega_d * ts) + (zeta * omega / omega_d) * np.sin(omega_d * ts)
)

P_smooth = np.einsum("ijk,ijl->ikl", P_smooth_sqr, P_smooth_sqr)
margin = 2 * np.sqrt(P_smooth[:, 0, 0])

plt.figure(figsize=(10, 3))
plt.plot(ts, m_smooth[:, 0], label="smoothed")
plt.plot(ts, x_true, "k--", label="true")
plt.fill_between(ts, m_smooth[:, 0] - margin, m_smooth[:, 0] + margin, alpha=0.3)
plt.xlabel("t"), plt.ylabel("x(t)"), plt.legend()
plt.show()

# Plot smoothed estimate vs true solution x_true = np.exp(-zeta * omega * ts) * ( np.cos(omega_d * ts) + (zeta * omega / omega_d) * np.sin(omega_d * ts) ) P_smooth = np.einsum("ijk,ijl->ikl", P_smooth_sqr, P_smooth_sqr) margin = 2 * np.sqrt(P_smooth[:, 0, 0]) plt.figure(figsize=(10, 3)) plt.plot(ts, m_smooth[:, 0], label="smoothed") plt.plot(ts, x_true, "k--", label="true") plt.fill_between(ts, m_smooth[:, 0] - margin, m_smooth[:, 0] + margin, alpha=0.3) plt.xlabel("t"), plt.ylabel("x(t)"), plt.legend() plt.show()

In [5]:

# Joint state-parameter estimation: infer damping gamma from noisy observations
noise_std = 1e-2
key = jrandom.PRNGKey(42)
z = (x_true[1:] + noise_std * jrandom.normal(key, shape=(N,))).reshape(-1, 1)
z_t = ts[1:]

plt.figure(figsize=(10, 3))
plt.scatter(z_t, z, s=10, alpha=0.7, color="black", label="observations")
plt.plot(ts, x_true, "k--", label="true")
plt.xlabel("t"), plt.ylabel("x(t)"), plt.legend()
plt.show()

# Joint state-parameter estimation: infer damping gamma from noisy observations noise_std = 1e-2 key = jrandom.PRNGKey(42) z = (x_true[1:] + noise_std * jrandom.normal(key, shape=(N,))).reshape(-1, 1) z_t = ts[1:] plt.figure(figsize=(10, 3)) plt.scatter(z_t, z, s=10, alpha=0.7, color="black", label="observations") plt.plot(ts, x_true, "k--", label="true") plt.xlabel("t"), plt.ylabel("x(t)"), plt.legend() plt.show()

In [6]:

# Setup joint prior: x (q=2) + gamma (q=1)
prior_x = IWP(q=2, d=1, Xi=0.1 * np.eye(1))
prior_gamma = IWP(q=2, d=1, Xi=1e-2 * np.eye(1))
prior_joint = JointPrior(prior_x, prior_gamma)

# Initial conditions (start with true gamma for initialization)
gamma0 = np.array([gamma])
mu_0_x, Sig_0_x = taylor_mode_initialization(vf, (x0, dx0), q=2, order=2)
print(mu_0_x.shape, Sig_0_x.shape)
mu_0_gamma = np.concatenate([gamma0, np.zeros(2)])
Sig_0_gamma = 1e-6 * np.eye(3)

D_x, D_gamma = 3, 3
mu_0_joint = np.concatenate([mu_0_x, mu_0_gamma])
zeros = np.zeros((D_x, D_gamma))
Sig_0_joint = np.block([[Sig_0_x, zeros], [zeros.T, Sig_0_gamma]])


# Vector field with hidden state u = gamma
def vf_joint(x, v, u, *, t):
    return -(omega**2) * x - u * v


# Measurement model: ODE + observations
obs = Measurement(A=np.array([[1.0]]), z=z, z_t=z_t, noise=noise_std)
measure_joint = SecondOrderODEInformationWithHidden(
    vf_joint,
    E0=prior_joint.E0_x,
    E1=prior_joint.E1,
    E2=prior_joint.E2,
    E0_hidden=prior_joint.E0_hidden,
    constraints=[obs],
)

M = 1 * N
ts2 = np.linspace(tspan[0], tspan[1], M + 1)
# Filter and smooth
m_filt, P_filt_sqr, _, _, G, d, P_back, *_ = ekf1_sqr_loop(
    mu_0_joint, Sig_0_joint, prior_joint, measure_joint, tspan, M
)
m_smooth, P_smooth_sqr = rts_sqr_smoother_loop(
    m_filt[-1], P_filt_sqr[-1], np.array(G), np.array(d), np.array(P_back), M
)
m_smooth, P_smooth_sqr = np.array(m_smooth), np.array(P_smooth_sqr)

# Setup joint prior: x (q=2) + gamma (q=1) prior_x = IWP(q=2, d=1, Xi=0.1 * np.eye(1)) prior_gamma = IWP(q=2, d=1, Xi=1e-2 * np.eye(1)) prior_joint = JointPrior(prior_x, prior_gamma) # Initial conditions (start with true gamma for initialization) gamma0 = np.array([gamma]) mu_0_x, Sig_0_x = taylor_mode_initialization(vf, (x0, dx0), q=2, order=2) print(mu_0_x.shape, Sig_0_x.shape) mu_0_gamma = np.concatenate([gamma0, np.zeros(2)]) Sig_0_gamma = 1e-6 * np.eye(3) D_x, D_gamma = 3, 3 mu_0_joint = np.concatenate([mu_0_x, mu_0_gamma]) zeros = np.zeros((D_x, D_gamma)) Sig_0_joint = np.block([[Sig_0_x, zeros], [zeros.T, Sig_0_gamma]]) # Vector field with hidden state u = gamma def vf_joint(x, v, u, *, t): return -(omega**2) * x - u * v # Measurement model: ODE + observations obs = Measurement(A=np.array([[1.0]]), z=z, z_t=z_t, noise=noise_std) measure_joint = SecondOrderODEInformationWithHidden( vf_joint, E0=prior_joint.E0_x, E1=prior_joint.E1, E2=prior_joint.E2, E0_hidden=prior_joint.E0_hidden, constraints=[obs], ) M = 1 * N ts2 = np.linspace(tspan[0], tspan[1], M + 1) # Filter and smooth m_filt, P_filt_sqr, _, _, G, d, P_back, *_ = ekf1_sqr_loop( mu_0_joint, Sig_0_joint, prior_joint, measure_joint, tspan, M ) m_smooth, P_smooth_sqr = rts_sqr_smoother_loop( m_filt[-1], P_filt_sqr[-1], np.array(G), np.array(d), np.array(P_back), M ) m_smooth, P_smooth_sqr = np.array(m_smooth), np.array(P_smooth_sqr)

(3,) (3, 3)

In [7]:

# Plot results
P = np.einsum("ijk,ijl->ikl", P_smooth_sqr, P_smooth_sqr)


fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True)

# State x
ax1.scatter(z_t, z, s=10, alpha=0.7, color="black", label="observations")
ax1.plot(ts2, m_smooth[:, 0], label="smoothed")
ax1.plot(ts, x_true, "k--", label="true")
ax1.fill_between(
    ts2,
    m_smooth[:, 0] - 2 * np.sqrt(P[:, 0, 0]),
    m_smooth[:, 0] + 2 * np.sqrt(P[:, 0, 0]),
    alpha=0.3,
)
ax1.set_ylabel("x(t)"), ax1.legend()

# Inferred gamma (at index D_x=3)
ax2.plot(ts2, m_smooth[:, D_x], label="smoothed gamma")
ax2.axhline(gamma, color="k", linestyle="--", label="true gamma")
ax2.fill_between(
    ts2,
    m_smooth[:, D_x] - 2 * np.sqrt(P[:, D_x, D_x]),
    m_smooth[:, D_x] + 2 * np.sqrt(P[:, D_x, D_x]),
    alpha=0.3,
)
ax2.set_xlabel("t"), ax2.set_ylabel("gamma"), ax2.legend()
ax2.set_ylim(0, 1.0)

plt.tight_layout()
plt.show()

# Plot results P = np.einsum("ijk,ijl->ikl", P_smooth_sqr, P_smooth_sqr) fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True) # State x ax1.scatter(z_t, z, s=10, alpha=0.7, color="black", label="observations") ax1.plot(ts2, m_smooth[:, 0], label="smoothed") ax1.plot(ts, x_true, "k--", label="true") ax1.fill_between( ts2, m_smooth[:, 0] - 2 * np.sqrt(P[:, 0, 0]), m_smooth[:, 0] + 2 * np.sqrt(P[:, 0, 0]), alpha=0.3, ) ax1.set_ylabel("x(t)"), ax1.legend() # Inferred gamma (at index D_x=3) ax2.plot(ts2, m_smooth[:, D_x], label="smoothed gamma") ax2.axhline(gamma, color="k", linestyle="--", label="true gamma") ax2.fill_between( ts2, m_smooth[:, D_x] - 2 * np.sqrt(P[:, D_x, D_x]), m_smooth[:, D_x] + 2 * np.sqrt(P[:, D_x, D_x]), alpha=0.3, ) ax2.set_xlabel("t"), ax2.set_ylabel("gamma"), ax2.legend() ax2.set_ylim(0, 1.0) plt.tight_layout() plt.show()

In [8]:

# Acceleration measurements using full_state=True
# For a second-order ODE with q=2, the state vector is X = [x, dx, ddx].
# A standard Measurement observes A @ x (= A @ E0 @ X), but acceleration
# lives in the full state vector. Using full_state=True, we can define
# A @ X directly, e.g. A = E2 to observe ddx.

# True velocity and acceleration from the exact solution
dx_true = (
    -omega_d * np.sin(omega_d * ts) * np.exp(-zeta * omega * ts)
    - zeta * omega * np.cos(omega_d * ts) * np.exp(-zeta * omega * ts)
)
ddx_true = -(omega**2) * x_true - gamma * dx_true

# Subsample acceleration observations (every 5th step)
accel_noise_std = 0.05
key = jrandom.PRNGKey(123)
obs_every = 5
accel_indices = np.arange(obs_every, N + 1, obs_every)
accel_ts = ts[accel_indices]
accel_true = ddx_true[accel_indices]
accel_z = (accel_true + accel_noise_std * jrandom.normal(key, shape=accel_true.shape)).reshape(-1, 1)

plt.figure(figsize=(10, 3))
plt.scatter(accel_ts, accel_z, s=20, color="red", label="noisy acceleration obs")
plt.plot(ts, ddx_true, "k--", label="true acceleration")
plt.xlabel("t"), plt.ylabel("ddx(t)"), plt.legend()
plt.title("Acceleration measurements")
plt.show()

# Acceleration measurements using full_state=True # For a second-order ODE with q=2, the state vector is X = [x, dx, ddx]. # A standard Measurement observes A @ x (= A @ E0 @ X), but acceleration # lives in the full state vector. Using full_state=True, we can define # A @ X directly, e.g. A = E2 to observe ddx. # True velocity and acceleration from the exact solution dx_true = ( -omega_d * np.sin(omega_d * ts) * np.exp(-zeta * omega * ts) - zeta * omega * np.cos(omega_d * ts) * np.exp(-zeta * omega * ts) ) ddx_true = -(omega**2) * x_true - gamma * dx_true # Subsample acceleration observations (every 5th step) accel_noise_std = 0.05 key = jrandom.PRNGKey(123) obs_every = 5 accel_indices = np.arange(obs_every, N + 1, obs_every) accel_ts = ts[accel_indices] accel_true = ddx_true[accel_indices] accel_z = (accel_true + accel_noise_std * jrandom.normal(key, shape=accel_true.shape)).reshape(-1, 1) plt.figure(figsize=(10, 3)) plt.scatter(accel_ts, accel_z, s=20, color="red", label="noisy acceleration obs") plt.plot(ts, ddx_true, "k--", label="true acceleration") plt.xlabel("t"), plt.ylabel("ddx(t)"), plt.legend() plt.title("Acceleration measurements") plt.show()

In [9]:

# Create acceleration measurement with full_state=True
# A = E2 (shape [1, 3]) operates on full state X = [x, dx, ddx]
accel_obs = Measurement(
    A=prior.E2,
    z=accel_z,
    z_t=accel_ts,
    noise=accel_noise_std,
    full_state=True,  # A operates on full state X, not just x
)

# ODE model with acceleration measurements
measure_accel = SecondOrderODEInformation(
    vf, prior.E0, prior.E1, prior.E2, constraints=[accel_obs]
)

# Run filter and smoother (non-jitted, since Measurement uses find_index)
m_filt_a, P_filt_sqr_a, _, _, G_a, d_a, P_back_sqr_a, _, _, ll_joint = ekf1_sqr_loop(
    mu_0, Sigma_0_sqr, prior, measure_accel, tspan, N
)
m_smooth_a, P_smooth_sqr_a = rts_sqr_smoother_loop(
    m_filt_a[-1], P_filt_sqr_a[-1],
    np.array(G_a), np.array(d_a), np.array(P_back_sqr_a), N
)
m_smooth_a = np.array(m_smooth_a)
P_smooth_sqr_a = np.array(P_smooth_sqr_a)

# Create acceleration measurement with full_state=True # A = E2 (shape [1, 3]) operates on full state X = [x, dx, ddx] accel_obs = Measurement( A=prior.E2, z=accel_z, z_t=accel_ts, noise=accel_noise_std, full_state=True, # A operates on full state X, not just x ) # ODE model with acceleration measurements measure_accel = SecondOrderODEInformation( vf, prior.E0, prior.E1, prior.E2, constraints=[accel_obs] ) # Run filter and smoother (non-jitted, since Measurement uses find_index) m_filt_a, P_filt_sqr_a, _, _, G_a, d_a, P_back_sqr_a, _, _, ll_joint = ekf1_sqr_loop( mu_0, Sigma_0_sqr, prior, measure_accel, tspan, N ) m_smooth_a, P_smooth_sqr_a = rts_sqr_smoother_loop( m_filt_a[-1], P_filt_sqr_a[-1], np.array(G_a), np.array(d_a), np.array(P_back_sqr_a), N ) m_smooth_a = np.array(m_smooth_a) P_smooth_sqr_a = np.array(P_smooth_sqr_a)

In [10]:

# Compare: ODE-only vs ODE + acceleration measurements
P_a = np.einsum("ijk,ijl->ikl", P_smooth_sqr_a, P_smooth_sqr_a)

fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True)

# Position
ax1.plot(ts, m_smooth_a[:, 0], label="with accel. obs")
ax1.plot(ts, x_true, "k--", label="true")
ax1.fill_between(
    ts,
    m_smooth_a[:, 0] - 2 * np.sqrt(P_a[:, 0, 0]),
    m_smooth_a[:, 0] + 2 * np.sqrt(P_a[:, 0, 0]),
    alpha=0.3,
)
ax1.set_ylabel("x(t)"), ax1.legend()
ax1.set_title("Position estimate (ODE + acceleration measurements)")

# Acceleration (index 2 in state [x, dx, ddx])
ax2.scatter(accel_ts, accel_z, s=20, color="red", zorder=3, label="observations")
ax2.plot(ts, m_smooth_a[:, 2], label="smoothed ddx")
ax2.plot(ts, ddx_true, "k--", label="true ddx")
ax2.fill_between(
    ts,
    m_smooth_a[:, 2] - 2 * np.sqrt(P_a[:, 2, 2]),
    m_smooth_a[:, 2] + 2 * np.sqrt(P_a[:, 2, 2]),
    alpha=0.3,
)
ax2.set_xlabel("t"), ax2.set_ylabel("ddx(t)"), ax2.legend()

plt.tight_layout()
plt.show()

# Compare: ODE-only vs ODE + acceleration measurements P_a = np.einsum("ijk,ijl->ikl", P_smooth_sqr_a, P_smooth_sqr_a) fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 7), sharex=True) # Position ax1.plot(ts, m_smooth_a[:, 0], label="with accel. obs") ax1.plot(ts, x_true, "k--", label="true") ax1.fill_between( ts, m_smooth_a[:, 0] - 2 * np.sqrt(P_a[:, 0, 0]), m_smooth_a[:, 0] + 2 * np.sqrt(P_a[:, 0, 0]), alpha=0.3, ) ax1.set_ylabel("x(t)"), ax1.legend() ax1.set_title("Position estimate (ODE + acceleration measurements)") # Acceleration (index 2 in state [x, dx, ddx]) ax2.scatter(accel_ts, accel_z, s=20, color="red", zorder=3, label="observations") ax2.plot(ts, m_smooth_a[:, 2], label="smoothed ddx") ax2.plot(ts, ddx_true, "k--", label="true ddx") ax2.fill_between( ts, m_smooth_a[:, 2] - 2 * np.sqrt(P_a[:, 2, 2]), m_smooth_a[:, 2] + 2 * np.sqrt(P_a[:, 2, 2]), alpha=0.3, ) ax2.set_xlabel("t"), ax2.set_ylabel("ddx(t)"), ax2.legend() plt.tight_layout() plt.show()

Sequential update: ODE first, then measurements¶

By default, the EKF incorporates ODE information and measurement constraints in a single joint update, both linearized at the predicted state. With ekf1_sqr_loop_sequential, the update is split into two stages:

1. ODE update -- incorporate the ODE residual (linearized at m_pred)
2. Measurement update -- incorporate observations (linearized at m_ode, the ODE-updated state)

The benefit is that the measurement update sees a tighter prior (m_ode, P_ode), which can improve accuracy when measurements are informative.

Observations are passed separately via the observations keyword argument, decoupling the ODE model from the observation schedule. The function returns two log-likelihoods: one for the ODE residual and one for the observations.

In [11]:

from ode_filters.filters import ekf1_sqr_loop_sequential

# ODE-only measure (no measurement constraints bundled in)
measure_ode_only = SecondOrderODEInformation(vf, prior.E0, prior.E1, prior.E2)

# Run sequential filter -- observations passed separately
(
    m_seq_s, P_seq_sqr_s,
    _, _,
    G_s, d_s, P_back_sqr_s,
    mz_ode_s, Pz_ode_sqr_s,
    mz_obs_s, Pz_obs_sqr_s,
    ll_ode, ll_obs,
) = ekf1_sqr_loop_sequential(
    mu_0, Sigma_0_sqr, prior, measure_ode_only, tspan, N,
    observations=[accel_obs],
)

# Smooth with the same RTS smoother (backward pass is unchanged)
m_smooth_s, P_smooth_sqr_s = rts_sqr_smoother_loop(
    m_seq_s[-1], P_seq_sqr_s[-1],
    np.array(G_s), np.array(d_s), np.array(P_back_sqr_s), N
)
m_smooth_s = np.array(m_smooth_s)
P_smooth_sqr_s = np.array(P_smooth_sqr_s)

print(f"Joint update log-likelihood:        {float(ll_joint):.4f}")
print(f"Sequential ODE log-likelihood:       {float(ll_ode):.4f}")
print(f"Sequential obs log-likelihood:       {float(ll_obs):.4f}")
print(f"Sequential total log-likelihood:     {float(ll_ode) + float(ll_obs):.4f}")

from ode_filters.filters import ekf1_sqr_loop_sequential # ODE-only measure (no measurement constraints bundled in) measure_ode_only = SecondOrderODEInformation(vf, prior.E0, prior.E1, prior.E2) # Run sequential filter -- observations passed separately ( m_seq_s, P_seq_sqr_s, _, _, G_s, d_s, P_back_sqr_s, mz_ode_s, Pz_ode_sqr_s, mz_obs_s, Pz_obs_sqr_s, ll_ode, ll_obs, ) = ekf1_sqr_loop_sequential( mu_0, Sigma_0_sqr, prior, measure_ode_only, tspan, N, observations=[accel_obs], ) # Smooth with the same RTS smoother (backward pass is unchanged) m_smooth_s, P_smooth_sqr_s = rts_sqr_smoother_loop( m_seq_s[-1], P_seq_sqr_s[-1], np.array(G_s), np.array(d_s), np.array(P_back_sqr_s), N ) m_smooth_s = np.array(m_smooth_s) P_smooth_sqr_s = np.array(P_smooth_sqr_s) print(f"Joint update log-likelihood: {float(ll_joint):.4f}") print(f"Sequential ODE log-likelihood: {float(ll_ode):.4f}") print(f"Sequential obs log-likelihood: {float(ll_obs):.4f}") print(f"Sequential total log-likelihood: {float(ll_ode) + float(ll_obs):.4f}")

Joint update log-likelihood:        -28.8593
Sequential ODE log-likelihood:       -37.5454
Sequential obs log-likelihood:       8.6861
Sequential total log-likelihood:     -28.8593

In [12]:

# Compare joint vs sequential smoothed estimates
P_s = np.einsum("ijk,ijl->ikl", P_smooth_sqr_s, P_smooth_sqr_s)

fig, axes = plt.subplots(2, 2, figsize=(12, 7), sharex=True)

# Position: joint vs sequential
for ax, m_sm, P_sm, title in [
    (axes[0, 0], m_smooth_a, P_a, "Joint update"),
    (axes[0, 1], m_smooth_s, P_s, "Sequential update"),
]:
    ax.plot(ts, m_sm[:, 0], label="smoothed")
    ax.plot(ts, x_true, "k--", label="true")
    ax.fill_between(
        ts,
        m_sm[:, 0] - 2 * np.sqrt(P_sm[:, 0, 0]),
        m_sm[:, 0] + 2 * np.sqrt(P_sm[:, 0, 0]),
        alpha=0.3,
    )
    ax.set_ylabel("x(t)"), ax.set_title(title), ax.legend(fontsize=8)

# Acceleration: joint vs sequential
for ax, m_sm, P_sm in [
    (axes[1, 0], m_smooth_a, P_a),
    (axes[1, 1], m_smooth_s, P_s),
]:
    ax.scatter(accel_ts, accel_z, s=20, color="red", zorder=3, label="observations")
    ax.plot(ts, m_sm[:, 2], label="smoothed ddx")
    ax.plot(ts, ddx_true, "k--", label="true ddx")
    ax.fill_between(
        ts,
        m_sm[:, 2] - 2 * np.sqrt(P_sm[:, 2, 2]),
        m_sm[:, 2] + 2 * np.sqrt(P_sm[:, 2, 2]),
        alpha=0.3,
    )
    ax.set_xlabel("t"), ax.set_ylabel("ddx(t)"), ax.legend(fontsize=8)

plt.suptitle("Joint vs Sequential EKF update (with acceleration measurements)")
plt.tight_layout()
plt.show()

# Compare joint vs sequential smoothed estimates P_s = np.einsum("ijk,ijl->ikl", P_smooth_sqr_s, P_smooth_sqr_s) fig, axes = plt.subplots(2, 2, figsize=(12, 7), sharex=True) # Position: joint vs sequential for ax, m_sm, P_sm, title in [ (axes[0, 0], m_smooth_a, P_a, "Joint update"), (axes[0, 1], m_smooth_s, P_s, "Sequential update"), ]: ax.plot(ts, m_sm[:, 0], label="smoothed") ax.plot(ts, x_true, "k--", label="true") ax.fill_between( ts, m_sm[:, 0] - 2 * np.sqrt(P_sm[:, 0, 0]), m_sm[:, 0] + 2 * np.sqrt(P_sm[:, 0, 0]), alpha=0.3, ) ax.set_ylabel("x(t)"), ax.set_title(title), ax.legend(fontsize=8) # Acceleration: joint vs sequential for ax, m_sm, P_sm in [ (axes[1, 0], m_smooth_a, P_a), (axes[1, 1], m_smooth_s, P_s), ]: ax.scatter(accel_ts, accel_z, s=20, color="red", zorder=3, label="observations") ax.plot(ts, m_sm[:, 2], label="smoothed ddx") ax.plot(ts, ddx_true, "k--", label="true ddx") ax.fill_between( ts, m_sm[:, 2] - 2 * np.sqrt(P_sm[:, 2, 2]), m_sm[:, 2] + 2 * np.sqrt(P_sm[:, 2, 2]), alpha=0.3, ) ax.set_xlabel("t"), ax.set_ylabel("ddx(t)"), ax.legend(fontsize=8) plt.suptitle("Joint vs Sequential EKF update (with acceleration measurements)") plt.tight_layout() plt.show()

In [ ]: