Extended Kalman Filter (EKF)
Derivation, argument reference, and worked example for the ekf block.
Overview
The EKF extends the Kalman filter to nonlinear systems by linearising the dynamics and observation functions around the current estimate at each step using user-supplied Jacobians. Everywhere the KF uses the constant matrices \(F_k\) and \(H_k\), the EKF substitutes their local first-order approximations \(F^J_k\) and \(H^J_k\).
It is the workhorse estimator for mildly nonlinear systems where Jacobians are analytically available or cheap to compute.
State-Space Model
The EKF assumes the system evolves as:
Symbol |
Space |
Role |
|---|---|---|
\(x_k\) |
\(\mathbb{R}^n\) |
True (hidden) state |
\(u_k\) |
\(\mathbb{R}\) |
Known control input |
\(y_k\) |
\(\mathbb{R}^m\) |
Observed measurement |
\(f(\cdot)\) |
\(\mathbb{R}^n \to \mathbb{R}^n\) |
Nonlinear state transition |
\(h(\cdot)\) |
\(\mathbb{R}^n \to \mathbb{R}^m\) |
Nonlinear observation function |
\(w_k\) |
\(\mathbb{R}^n\) |
Process noise, covariance \(Q_k\) |
\(v_k\) |
\(\mathbb{R}^m\) |
Measurement noise, covariance \(R_k\) |
Predict Step
Propagate the estimate through the true nonlinear function, then linearise for the covariance:
where \(F^J_k = \dfrac{\partial f}{\partial x}\Big|_{\hat{x}_{k-1}}\) is the Jacobian of \(f\) with respect to \(x\),
evaluated at the previous estimate. This must be supplied by the caller as Fjk.
Update Step
Identical in structure to the KF update, with \(F^J_k\), \(H^J_k\) in place of \(F_k\), \(H_k\), and the nonlinear \(h\) used for the innovation:
where \(H^J_k = \dfrac{\partial h}{\partial x}\Big|_{\hat{x}_{k|k-1}}\). Supplied by the caller as Hjk.
Note the innovation uses \(h(\hat{x}_{k|k-1})\) rather than \(H^J_k \hat{x}_{k|k-1}\) — the nonlinear function is used directly for the residual, only the covariance propagation is linearised.
Argument Reference
zk — EKF state tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]]
Field |
Symbol |
Shape |
Description |
|---|---|---|---|
|
\(\hat{x}_k\) |
|
Current state estimate |
|
\(P_k\) |
|
Error covariance — symmetric, positive semi-definite |
Initialisation: zk = (x0, P0). Same convention as the KF.
uk — control input float
Scalar control input at step \(k\). Passed directly to f_nl(x, uk). Pass 0.0 if uncontrolled.
yk — measurement Float[np.ndarray, "m"]
Raw sensor reading at step \(k\), shape (m,). Must match the output space of h_nl.
f_nl — nonlinear dynamics Callable
Signature: f_nl(x: ndarray, u: float) -> ndarray
The true nonlinear state transition. Used directly to propagate the mean in the predict step.
def f_nl(x, u):
th, w = x
return np.array([th + dt * w, w + dt * (-g / L * np.sin(th) + u)])
h_nl — nonlinear observation Callable
Signature: h_nl(x: ndarray) -> ndarray
The true nonlinear observation function. Used in the innovation residual \(y_k - h(\hat{x}_{k|k-1})\).
def h_nl(x):
return np.array([x[0]]) # observe only theta
Fjk — state Jacobian Float[np.ndarray, "n n"]
Shape (n, n). The Jacobian of f_nl with respect to x, evaluated at \(\hat{x}_{k-1}\):
Must be recomputed each step at the current estimate. For the pendulum:
def pend_Fj(x):
th, w = x
return np.array([[1.0, dt],
[-dt * g / L * np.cos(th), 1.0]])
Fjk = pend_Fj(zk[0]) # evaluate at current estimate
Hjk — observation Jacobian Float[np.ndarray, "m n"]
Shape (m, n). The Jacobian of h_nl with respect to x, evaluated at \(\hat{x}_{k|k-1}\):
If h_nl is linear (e.g. observe one component), Hjk is constant and can be precomputed:
Hjk = np.array([[1.0, 0.0]]) # shape (1, 2), constant for linear h
Qk — process noise covariance Float[np.ndarray, "n n"]
Shape (n, n), symmetric positive semi-definite. Same role as in the KF — encodes model uncertainty.
Rk — measurement noise covariance Float[np.ndarray, "m m"]
Shape (m, m), symmetric positive definite. Same role as in the KF — encodes sensor noise.
Implementation
ekf_f runs the full predict→update cycle using the nonlinear functions and their Jacobians.
ekf_h extracts \(\hat{x}_k\) for downstream use.
The **kwargs in ekf_h absorbs any extra arguments passed by _smart_call that h does not need.
import numpy as np
from typing import Callable
from jaxtyping import Float, jaxtyped
from beartype import beartype
from dynamicalnodes import DynamicalSystem
@jaxtyped(typechecker=beartype)
def ekf_f(
zk: tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]],
uk: float,
yk: Float[np.ndarray, "m"],
f_nl: Callable,
h_nl: Callable,
Fjk: Float[np.ndarray, "n n"],
Hjk: Float[np.ndarray, "m n"],
Qk: Float[np.ndarray, "n n"],
Rk: Float[np.ndarray, "m m"],
) -> tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]]:
"""EKF predict+update via Jacobian linearisation; returns updated state and covariance.
Args:
zk: EKF state — 2-tuple of (x̂_k, P_k), state estimate and error covariance.
uk: Scalar control input.
yk: Measurement vector (m,).
f_nl: Nonlinear state transition f(x, u) → x_next.
h_nl: Nonlinear observation function h(x) → y.
Fjk: Jacobian ∂f/∂x evaluated at x̂_{k-1} (n, n).
Hjk: Jacobian ∂h/∂x evaluated at x̂_{k|k-1} (m, n).
Qk: Process noise covariance (n, n).
Rk: Measurement noise covariance (m, m).
"""
x, P = zk
x_pred = f_nl(x, uk)
P_pred = Fjk @ P @ Fjk.T + Qk
S = Hjk @ P_pred @ Hjk.T + Rk
K = P_pred @ Hjk.T @ np.linalg.inv(S)
yk = np.atleast_1d(yk)
x_upd = x_pred + K @ (yk - h_nl(x_pred))
P_upd = (np.eye(len(x)) - K @ Hjk) @ P_pred
return (x_upd, P_upd)
@jaxtyped(typechecker=beartype)
def ekf_h(
zk: tuple[Float[np.ndarray, "n"], Float[np.ndarray, "n n"]],
**kwargs,
) -> Float[np.ndarray, "n"]:
"""Return the current state estimate from the EKF state tuple.
Args:
zk: EKF state — 2-tuple of (x̂_k, P_k), state estimate and error covariance.
"""
return zk[0]
ekf = DynamicalSystem(f=ekf_f, h=ekf_h)
Worked Example — Pendulum
State x = [theta, theta_dot], uncontrolled, observe angle only.
g, L, dt = 9.81, 1.0, 0.05
def pend_f(x, u):
th, w = x
return np.array([th + dt * w, w + dt * (-g / L * np.sin(th) + u)])
def pend_h(x):
return np.array([x[0]])
def pend_Fj(x):
th, w = x
return np.array([[1.0, dt], [-dt * g / L * np.cos(th), 1.0]])
Hjk = np.array([[1.0, 0.0]])
Qk = np.diag([1e-4, 1e-3])
Rk = np.array([[0.05]])
zk = (np.array([0.1, 0.0]), np.eye(2))
zk_next, x_est = ekf.step(
zk=zk,
uk=0.0,
yk=np.array([0.09]),
f_nl=pend_f,
h_nl=pend_h,
Fjk=pend_Fj(zk[0]),
Hjk=Hjk,
Qk=Qk,
Rk=Rk,
)
print("estimate (pre-update):", x_est)
print("updated estimate: ", zk_next[0])
print("updated covariance:\n", zk_next[1])
DynamicalSystem → ROSNode
from dynamicalnodes import ROSNode
from dynamicalnodes.ros2py_py2ros import (
ros2py_float64_multiarray,
ros2py_float64,
py2ros_float64_multiarray,
)
from std_msgs.msg import Float64MultiArray, Float64
ekf_node = ROSNode(
dynamical_system=ekf,
subscribes_to=[
{
"topic": "/measurement",
"msg_type": Float64MultiArray,
"arg": "yk",
"ros2py": ros2py_float64_multiarray,
},
{
"topic": "/control",
"msg_type": Float64,
"arg": "uk",
"ros2py": ros2py_float64,
},
],
publishes_to=[
{
"topic": "/state_estimate",
"msg_type": Float64MultiArray,
"py2ros": py2ros_float64_multiarray,
},
],
)
ROSNode.write_ROSNode_to_rclpy()
g, L, dt_ekf = 9.81, 1.0, 0.1
def pendulum_f(x, u):
theta, omega = x
alpha = -g / L * np.sin(theta) + u
return np.array([theta + omega * dt_ekf, omega + alpha * dt_ekf])
def pendulum_h(x):
return np.array([x[0]])
Fjk = np.array([[1.0, dt_ekf], [-g / L * dt_ekf, 1.0]])
Hjk = np.array([[1.0, 0.0]])
ekf_node.write_ROSNode_to_rclpy(
"ekf_node.py",
node_name="extended_kalman_filter",
deps=[pendulum_f, pendulum_h],
initial_state=(
np.array([0.0, 0.0]), # x̂_0: [θ, ω]
np.eye(2), # P_0
),
static_params={
"f_nl": pendulum_f,
"h_nl": pendulum_h,
"Fjk": Fjk,
"Hjk": Hjk,
"Qk": np.diag([1e-4, 1e-3]),
"Rk": np.array([[0.01]]),
},
initial_inputs={
"yk": np.array([0.0]),
"uk": 0.0,
},
)