Skip to content

Continuous Nonlinear Estimation

This example starts from a simple but realistic estimation problem: the plant is continuous-time, the measurements arrive on a discrete grid, the input varies inside each sample interval, and the sensor only reports part of the state.

That combination is where Contrax's continuous-to-discrete estimation bridge is supposed to matter. The model stays continuous and nonlinear, sample_system() turns it into a discrete transition map, and the UKF/UKS pipeline works on the measurement grid without application-side glue code.

Runnable script: examples/continuous_nonlinear_estimation.py

Problem Setup

The plant is a damped pendulum with state

\[ x = \begin{bmatrix} \theta \\ \dot{\theta} \end{bmatrix}, \]

input torque u(t), and dynamics

\[ \dot{\theta} = \dot{\theta}, \qquad \ddot{\theta} = -0.35\dot{\theta} - \sin(\theta) + u(t). \]

The sensor only measures angle:

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

So the estimator has to reconstruct angular rate from model structure and the history of partial observations.

Build The Continuous Model And FOH Bridge

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

import jax.numpy as jnp
import jax.random as jr

import contrax as cx

DT = 0.1
DURATION = 4.0
NUM_STEPS = int(DURATION / DT)
NOISE_KEY = jr.PRNGKey(7)
X0_TRUE = jnp.array([0.7, -0.1])
X0_EST = jnp.array([0.45, 0.2])
P0 = jnp.diag(jnp.array([0.08, 0.12]))
Q_NOISE = 6e-4 * jnp.eye(2)
R_NOISE = jnp.array([[6e-3]])
MEASUREMENT_STD = jnp.sqrt(R_NOISE[0, 0])
SAMPLE_TIMES = jnp.arange(NUM_STEPS, dtype=jnp.float64) * DT
INPUT_SAMPLES = (
    0.55 * jnp.sin(1.25 * SAMPLE_TIMES) + 0.2 * jnp.cos(0.55 * SAMPLE_TIMES + 0.3)
)[:, None]
INPUTS_FOH = cx.foh_inputs(INPUT_SAMPLES)


def pendulum_dynamics(t, x, u):
    theta, theta_dot = x
    torque = u[0]
    return jnp.array([theta_dot, -0.35 * theta_dot - jnp.sin(theta) + torque])


def angle_sensor(t, x, u):
    del t, u
    return x[:1]


SYS_CONT = cx.nonlinear_system(
    pendulum_dynamics,
    observation=angle_sensor,
    state_dim=2,
    input_dim=1,
    output_dim=1,
)
SYS_DISC = cx.sample_system(SYS_CONT, DT, input_interpolation="foh")

Two pieces matter here.

First, the plant remains a continuous NonlinearSystem. We are not rewriting it as an ad hoc discrete transition by hand.

Second, the sampled estimator model uses sample_system(..., input_interpolation="foh"). Each discrete input step is an endpoint pair (u_k, u_{k+1}), so the internal one-step solver sees a piecewise-linear torque profile instead of a zero-order hold.

Run Filtering, Smoothing, And Diagnostics

def run_example():
    x_true = sample_true_trajectory()
    ys = sample_measurements(x_true)

    filtered = cx.ukf(
        SYS_DISC,
        Q_noise=Q_NOISE,
        R_noise=R_NOISE,
        ys=ys,
        us=INPUTS_FOH,
        x0=X0_EST,
        P0=P0,
        alpha=0.45,
    )
    smoothed = cx.uks(
        SYS_DISC,
        filtered,
        Q_noise=Q_NOISE,
        us=INPUTS_FOH,
        alpha=0.45,
    )
    innovation_diag, likelihood_diag = cx.ukf_diagnostics(filtered)
    rmse = summarize_rmse(filtered, smoothed, x_true)

    assert filtered.x_hat.shape == (NUM_STEPS, 2)
    assert filtered.predicted_measurements.shape == (NUM_STEPS, 1)
    assert smoothed.x_smooth.shape == (NUM_STEPS, 2)
    assert rmse["smoothed_theta_rmse"] < rmse["filtered_theta_rmse"]
    assert rmse["smoothed_rate_rmse"] < rmse["filtered_rate_rmse"]

    return {
        "times": SAMPLE_TIMES,
        "inputs": INPUT_SAMPLES[:, 0],
        "measurements": ys[:, 0],
        "true_theta": x_true[:, 0],
        "filtered_theta": filtered.x_hat[:, 0],
        "smoothed_theta": smoothed.x_smooth[:, 0],
        "true_rate": x_true[:, 1],
        "filtered_rate": filtered.x_hat[:, 1],
        "smoothed_rate": smoothed.x_smooth[:, 1],
        **rmse,
        "mean_nis": float(innovation_diag.mean_nis),
        "max_condition_number": float(
            innovation_diag.max_innovation_cov_condition_number
        ),
        "total_log_likelihood": float(likelihood_diag.total_log_likelihood),
    }

The workflow is:

  1. simulate the continuous pendulum under a smooth torque profile
  2. sample noisy angle measurements on the estimator grid
  3. build a discrete estimation model with sample_system()
  4. run ukf() on the sampled measurements and uks() on the stored forward-pass intermediates
  5. inspect ukf_diagnostics() for innovation scale, conditioning, and likelihood summaries

The measurement samples in this script use seeded Gaussian noise, so the example stays reproducible without hard-coding a synthetic offset pattern.

Continuous nonlinear pendulum estimation example showing the torque input and the true, filtered, and smoothed angle and angular-rate trajectories
Continuous nonlinear estimation with FOH inputs: the top panel shows the sampled torque profile, the middle panel shows angle measurements against the filtered and smoothed estimates, and the bottom panel shows recovery of the unmeasured angular rate.

The important behavior is in the bottom panel. Angle is observed directly, but angular rate is latent. The smoother recovers that hidden state more cleanly than the forward filter because it can use future measurements as well as past ones.

What The Script Prints

Running examples/continuous_nonlinear_estimation.py prints a compact summary of estimation quality:

Continuous nonlinear estimation
filtered theta rmse     = 0.047379
smoothed theta rmse     = 0.029366
filtered rate rmse      = 0.076574
smoothed rate rmse      = 0.039148
mean NIS                = 0.849424
max innovation cond     = 1.000000
total log likelihood    = 38.985644

The useful checks are straightforward:

  • the smoothed angle RMSE is lower than the filtered angle RMSE
  • the smoothed angular-rate RMSE is lower than the filtered angular-rate RMSE
  • the innovation covariance stays well-conditioned across the run