Skip to content

FOH Estimation: EKF with First-Order-Hold Inputs

This example filters a nonlinear oscillator with an input that varies between samples. First-order-hold (FOH) interpolation represents the input as a linear profile over each sample interval.

Runnable script: examples/foh_estimation.py

Problem Setup

The plant is a Van der Pol oscillator driven by a scalar input:

\[ \dot{x}_1 = x_2, \qquad \dot{x}_2 = \mu (1 - x_1^2) x_2 - x_1 + u(t) \]

The sensor observes the first state only:

\[ y_k = x_1(t_k) + v_k. \]

The input \(u(t)\) is sinusoidal and sampled on the same grid as the measurements. FOH follows this piecewise-linear profile within each interval; zero-order hold models the input as constant there.

Setup and Model

import jax

jax.config.update("jax_enable_x64", True)

import jax.numpy as jnp
import numpy as np

import contrax as cx

# Van der Pol oscillator: ẋ = [x2, μ(1-x1²)x2 - x1 + u]
MU = 1.0
DT = 0.1  # sample period
T_STEPS = 50  # number of steps

def vdp_dynamics(t, x, u):
    """Continuous-time Van der Pol dynamics with additive scalar input."""
    x1, x2 = x[0], x[1]
    dx1 = x2
    dx2 = MU * (1 - x1**2) * x2 - x1 + u[0]
    return jnp.array([dx1, dx2])


def vdp_observation(t, x, u):
    """Observe the first state (position) only."""
    return x[:1]


nl = cx.nonlinear_system(
    vdp_dynamics, vdp_observation, dt=None, state_dim=2, input_dim=1, output_dim=1
)

# Build a discrete transition model via FOH integration
discrete_foh = cx.sample_system(nl, DT, input_interpolation="foh")

sample_system(..., input_interpolation="foh") builds a discrete transition model that integrates the continuous dynamics under a linearly interpolated input \(u(t) = (1 - \tau) u_k + \tau u_{k+1}\) over \(\tau \in [0, 1]\).

Each step of this discrete model takes a stacked input pair u_pair = [[u_k], [u_{k+1}]] with shape (2, m). foh_inputs(us) converts a raw (T, m) sequence into the (T, 2, m) pairs the model expects.

Simulating and Filtering

def generate_trajectory(seed: int = 0):
    """Simulate ground-truth trajectory and noisy measurements."""
    key = jax.random.PRNGKey(seed)

    # Piecewise-linear input: slowly varying
    t_axis = jnp.linspace(0.0, T_STEPS * DT, T_STEPS)
    us_raw = 0.3 * jnp.sin(2 * jnp.pi * 0.5 * t_axis)[:, None]  # (T, 1)
    us_foh = cx.foh_inputs(us_raw)  # (T, 2, 1)

    # Simulate true trajectory using FOH discrete model
    x0_true = jnp.array([0.5, 0.0])
    xs_true = cx.rollout(
        lambda x, u: discrete_foh.dynamics(0.0, x, u), x0_true, us_foh
    )  # (T+1, 2)

    # Noisy measurements of first state only
    R_noise = jnp.array([[0.05]])
    key, subkey = jax.random.split(key)
    noise = jax.random.normal(subkey, (T_STEPS + 1, 1)) * jnp.sqrt(R_noise[0, 0])
    ys = xs_true[:, :1] + noise  # (T+1, 1)

    return xs_true, ys, us_raw, us_foh

def run_ekf(ys, us_foh):
    """Run an online EKF using FOH-sampled discrete model."""
    Q_noise = jnp.diag(jnp.array([1e-4, 1e-3]))
    R_noise = jnp.array([[0.05]])

    x = ys[0, :1]  # initialise from first measurement
    x = jnp.concatenate([x, jnp.zeros(1)])
    P = jnp.eye(2) * 0.5

    def step(carry, inputs):
        mean, covariance = carry
        t, u_pair, y = inputs
        prediction = cx.ekf_predict(
            discrete_foh, mean, covariance, u=u_pair, Q=Q_noise, t=t
        )
        # The measurement follows the transition, at the right input endpoint.
        measurement_u = jnp.stack([u_pair[-1], u_pair[-1]])
        update = cx.ekf_update(
            discrete_foh,
            prediction.x,
            prediction.P,
            y,
            u=measurement_u,
            R=R_noise,
            t=t + DT,
        )
        return (update.x, update.P), update.x

    times = jnp.arange(us_foh.shape[0]) * DT
    _, xs_est = jax.lax.scan(step, (x, P), (times, us_foh, ys[1:]))
    return jnp.concatenate([x[None], xs_est])  # (T+1, 2)

The EKF uses ekf_predict and ekf_update inside one jax.lax.scan, compiling the time loop once. The sampled system object supplies dynamics, observation and integration diagnostics; JAX computes the local Jacobians. Prediction uses the interval's FOH pair, and the following measurement uses its right endpoint.

What the Script Prints

Continuous-model EKF with FOH input interpolation
  steps              = 50
  position RMSE      = 0.1204
  velocity RMSE      = 0.1999
  true  final state  = [0.77012901 2.5961284 ]
  estim final state  = [0.75520658 2.62896228]

All assertions passed.