ekf.rs
raw
//! Extended Kalman Filter implementation.
//!
//! Best for systems where you have a good mathematical model.
//!
//! A "good" mathematical model is mono-modal, and could be described
//! in closed-form calculus expressions. Neither of these are strict
//! requirements, but the Extended Kalman Filter works best when
//! they're true or nearly so. The Particle Filter is a better choice
//! if these restrictions are too confining.
use super::*;
use nalgebra::{Const, DefaultAllocator, DimName, OMatrix, OVector, allocator::Allocator};
/// Describes how your system changes over time.
///
/// Implement this trait by filling in `predict_next` and `process_noise_std`.
/// The filter figures out the rest automatically.
///
/// # Type Parameters
/// - `N`: Number of state variables (e.g., 3 for [x, y, z])
/// - `C`: Number of control inputs (e.g., 2 for [thrust, steering])
///
/// # Example
///
/// ```
/// use estimators::ekf::ProcessModel;
///
/// struct BoatModel;
///
/// impl ProcessModel<3, 2> for BoatModel {
/// fn predict_next(&self, state: &[f64; 3], control: &[f64; 2], dt: f64) -> [f64; 3] {
/// let [x, y, heading] = *state;
/// let [throttle, rudder] = *control;
/// let speed = throttle * 5.0; // max 5 m/s
///
/// [
/// x + speed * heading.cos() * dt,
/// y + speed * heading.sin() * dt,
/// heading + rudder * dt,
/// ]
/// }
///
/// fn process_noise_std(&self) -> [f64; 3] {
/// [0.1, 0.1, 0.05] // [x noise, y noise, heading noise]
/// }
/// }
///
/// // Test the model
/// let model = BoatModel;
/// let state = [0.0, 0.0, 0.0]; // At origin, heading east
/// let control = [1.0, 0.0]; // Full throttle, no rudder
/// let next = model.predict_next(&state, &control, 1.0);
/// assert!((next[0] - 5.0).abs() < 0.01); // Moved 5m east
/// ```
pub trait ProcessModel<const N: usize, const C: usize> {
/// Predict the next state given current state, control input, and time step.
///
/// This is the heart of your model. Just describe what happens to your
/// system over a small time step.
fn predict_next(&self, state: &[f64; N], control: &[f64; C], dt: f64) -> [f64; N];
/// How much random noise affects each state variable (per second).
///
/// Return the standard deviation for each variable. For example, if your
/// position drifts by about ±0.1 meters per second due to wind or other
/// random factors, use 0.1.
fn process_noise_std(&self) -> [f64; N];
/// Advanced: Override this if you want to compute the sensitivity matrix
/// analytically instead of using automatic numerical differentiation.
///
/// Most users can ignore this — the default implementation calculates
/// this automatically by testing how small changes in each state variable
/// affect the prediction.
fn prediction_sensitivity(
&self,
state: &[f64; N],
control: &[f64; C],
dt: f64,
) -> [[f64; N]; N] {
numerical_jacobian(|s| self.predict_next(s, control, dt), state)
}
}
/// Describes what your sensors measure.
///
/// Implement this trait by filling in `expected_measurement` and `measurement_noise_std`.
///
/// # Type Parameters
/// - `N`: Number of state variables
/// - `M`: Number of sensor readings
///
/// # Example
///
/// ```
/// use estimators::ekf::MeasurementModel;
///
/// struct GpsSensor;
///
/// // State is [x, y, heading], but GPS only measures [x, y]
/// impl MeasurementModel<3, 2> for GpsSensor {
/// fn expected_measurement(&self, state: &[f64; 3]) -> [f64; 2] {
/// [state[0], state[1]] // GPS reads x and y position
/// }
///
/// fn measurement_noise_std(&self) -> [f64; 2] {
/// [2.5, 2.5] // GPS accurate to about ±2.5 meters
/// }
/// }
///
/// // Test the sensor model
/// let sensor = GpsSensor;
/// let state = [10.0, 20.0, 1.57]; // At (10, 20), heading north
/// let expected = sensor.expected_measurement(&state);
/// assert_eq!(expected, [10.0, 20.0]);
/// ```
pub trait MeasurementModel<const N: usize, const M: usize> {
/// What measurement would you expect given a particular state?
///
/// If you knew the true state exactly, what would your sensor report?
fn expected_measurement(&self, state: &[f64; N]) -> [f64; M];
/// How noisy is each sensor reading?
///
/// Return the standard deviation for each measurement. For example,
/// if your GPS is accurate to about ±2.5 meters, use 2.5.
fn measurement_noise_std(&self) -> [f64; M];
/// Advanced: Override this if you want to compute the sensitivity matrix
/// analytically instead of using automatic numerical differentiation.
///
/// Most users can ignore this — the default implementation calculates
/// this automatically by testing how small changes in each state variable
/// affect the expected measurement.
fn measurement_sensitivity(&self, state: &[f64; N]) -> [[f64; M]; N] {
numerical_jacobian(|s| self.expected_measurement(s), state)
}
}
/// Compute a Jacobian matrix using numerical differentiation.
///
/// Uses central differences for accuracy: for each input variable,
/// we nudge it slightly up and down and measure how the output changes.
pub fn numerical_jacobian<const N: usize, const M: usize, F>(
f: F,
state: &[f64; N],
) -> [[f64; M]; N]
where
F: Fn(&[f64; N]) -> [f64; M],
{
const EPSILON: f64 = 1e-8;
let mut jacobian = [[0.0; M]; N];
for i in 0..N {
let mut state_plus = *state;
let mut state_minus = *state;
state_plus[i] += EPSILON;
state_minus[i] -= EPSILON;
let f_plus = f(&state_plus);
let f_minus = f(&state_minus);
for j in 0..M {
jacobian[i][j] = (f_plus[j] - f_minus[j]) / (2.0 * EPSILON);
}
}
jacobian
}
/// Convert array-of-arrays to nalgebra matrix (transposed to match convention).
fn jacobian_to_matrix<const ROWS: usize, const COLS: usize>(
jac: [[f64; ROWS]; COLS],
) -> OMatrix<f64, Const<ROWS>, Const<COLS>>
where
Const<ROWS>: DimName,
Const<COLS>: DimName,
DefaultAllocator: Allocator<Const<ROWS>, Const<COLS>>,
{
let mut mat = OMatrix::<f64, Const<ROWS>, Const<COLS>>::zeros();
for col in 0..COLS {
for row in 0..ROWS {
mat[(row, col)] = jac[col][row];
}
}
mat
}
/// Build a diagonal covariance matrix from standard deviations.
fn std_to_covariance<const N: usize>(std_devs: [f64; N]) -> OMatrix<f64, Const<N>, Const<N>>
where
Const<N>: DimName,
DefaultAllocator: Allocator<Const<N>, Const<N>>,
{
let mut cov = OMatrix::<f64, Const<N>, Const<N>>::zeros();
for i in 0..N {
cov[(i, i)] = std_devs[i] * std_devs[i];
}
cov
}
/// Extended Kalman Filter.
///
/// Tracks state by combining predictions from your model with noisy measurements.
///
/// # Example
///
/// ```
/// use estimators::ekf::{ProcessModel, MeasurementModel, ExtendedKalmanFilter};
/// use estimators::StateEstimator;
///
/// // Simple constant-velocity model: state is [position, velocity]
/// struct ConstantVelocity;
///
/// impl ProcessModel<2, 0> for ConstantVelocity {
/// fn predict_next(&self, state: &[f64; 2], _control: &[f64; 0], dt: f64) -> [f64; 2] {
/// [state[0] + state[1] * dt, state[1]]
/// }
/// fn process_noise_std(&self) -> [f64; 2] {
/// [0.1, 0.01]
/// }
/// }
///
/// struct PositionSensor;
///
/// impl MeasurementModel<2, 1> for PositionSensor {
/// fn expected_measurement(&self, state: &[f64; 2]) -> [f64; 1] {
/// [state[0]] // We only measure position
/// }
/// fn measurement_noise_std(&self) -> [f64; 1] {
/// [0.5]
/// }
/// }
///
/// // Create the filter
/// let mut ekf = ExtendedKalmanFilter::new(
/// [0.0, 1.0], // Start at position 0, velocity 1
/// [1.0, 0.5], // Initial uncertainty
/// ConstantVelocity,
/// PositionSensor,
/// );
///
/// // Run a few iterations
/// for i in 0..5 {
/// ekf.predict(&[], 0.1); // 100ms time step, no control input
/// let measurement = [(i as f64 + 1.0) * 0.1]; // Simulated measurement
/// ekf.update(&measurement).unwrap();
/// }
///
/// let state = ekf.estimate();
/// println!("Final estimate: pos={:.2}, vel={:.2}", state[0], state[1]);
/// ```
pub struct ExtendedKalmanFilter<const N: usize, const M: usize, const C: usize, PM, MM>
where
Const<N>: DimName,
Const<M>: DimName,
Const<C>: DimName,
DefaultAllocator: Allocator<Const<N>>
+ Allocator<Const<N>, Const<N>>
+ Allocator<Const<M>>
+ Allocator<Const<M>, Const<N>>
+ Allocator<Const<N>, Const<M>>
+ Allocator<Const<M>, Const<M>>,
{
state: OVector<f64, Const<N>>,
covariance: OMatrix<f64, Const<N>, Const<N>>,
process_model: PM,
measurement_model: MM,
}
impl<const N: usize, const M: usize, const C: usize, PM, MM> ExtendedKalmanFilter<N, M, C, PM, MM>
where
Const<N>: DimName,
Const<M>: DimName,
Const<C>: DimName,
PM: ProcessModel<N, C>,
MM: MeasurementModel<N, M>,
DefaultAllocator: Allocator<Const<N>>
+ Allocator<Const<N>, Const<N>>
+ Allocator<Const<M>>
+ Allocator<Const<M>, Const<N>>
+ Allocator<Const<N>, Const<M>>
+ Allocator<Const<M>, Const<M>>,
{
/// Create a new Extended Kalman Filter.
///
/// # Arguments
/// * `initial_state` - Your best guess at the starting state
/// * `initial_uncertainty` - How uncertain you are about each state variable
/// (standard deviation). Use larger values if you're very unsure.
/// * `process_model` - Describes how your system evolves
/// * `measurement_model` - Describes what your sensors measure
pub fn new(
initial_state: [f64; N],
initial_uncertainty: [f64; N],
process_model: PM,
measurement_model: MM,
) -> Self {
Self {
state: OVector::<f64, Const<N>>::from_column_slice(&initial_state),
covariance: std_to_covariance(initial_uncertainty),
process_model,
measurement_model,
}
}
/// Create with a custom initial covariance matrix (advanced).
pub fn with_covariance(
initial_state: [f64; N],
initial_covariance: [[f64; N]; N],
process_model: PM,
measurement_model: MM,
) -> Self {
let mut cov = OMatrix::<f64, Const<N>, Const<N>>::zeros();
for i in 0..N {
for j in 0..N {
cov[(i, j)] = initial_covariance[i][j];
}
}
Self {
state: OVector::<f64, Const<N>>::from_column_slice(&initial_state),
covariance: cov,
process_model,
measurement_model,
}
}
/// Reset the filter to a known state.
///
/// Useful if you have external information (like a GPS fix) that you
/// want to use as a fresh starting point.
pub fn reset(&mut self, state: [f64; N], uncertainty: [f64; N]) {
self.state = OVector::<f64, Const<N>>::from_column_slice(&state);
self.covariance = std_to_covariance(uncertainty);
}
}
impl<const N: usize, const M: usize, const C: usize, PM, MM> StateEstimator
for ExtendedKalmanFilter<N, M, C, PM, MM>
where
Const<N>: DimName,
Const<M>: DimName,
Const<C>: DimName,
PM: ProcessModel<N, C>,
MM: MeasurementModel<N, M>,
DefaultAllocator: Allocator<Const<N>>
+ Allocator<Const<N>, Const<N>>
+ Allocator<Const<M>>
+ Allocator<Const<M>, Const<N>>
+ Allocator<Const<N>, Const<M>>
+ Allocator<Const<M>, Const<M>>,
{
type State = [f64; N];
type Measurement = [f64; M];
type Control = [f64; C];
fn predict(&mut self, control: &[f64; C], dt: f64) {
// Get current state as array
let state_arr: [f64; N] = self.state.as_slice().try_into().unwrap();
// Predict next state
let next_state = self.process_model.predict_next(&state_arr, control, dt);
self.state = OVector::<f64, Const<N>>::from_column_slice(&next_state);
// Get sensitivity matrix and convert to nalgebra
let sensitivity = self
.process_model
.prediction_sensitivity(&state_arr, control, dt);
let f = jacobian_to_matrix::<N, N>(sensitivity);
// Build process noise covariance (scaled by dt)
let noise_std = self.process_model.process_noise_std();
let mut q = OMatrix::<f64, Const<N>, Const<N>>::zeros();
for i in 0..N {
q[(i, i)] = noise_std[i] * noise_std[i] * dt;
}
// Update covariance: P = F * P * F^T + Q
self.covariance = &f * &self.covariance * f.transpose() + q;
}
fn update(&mut self, measurement: &[f64; M]) -> Result<()> {
let state_arr: [f64; N] = self.state.as_slice().try_into().unwrap();
// What measurement would we expect?
let expected = self.measurement_model.expected_measurement(&state_arr);
let z_pred = OVector::<f64, Const<M>>::from_column_slice(&expected);
let z = OVector::<f64, Const<M>>::from_column_slice(measurement);
// Innovation (difference between actual and expected)
let innovation = z - z_pred;
// Measurement sensitivity matrix
let sensitivity = self.measurement_model.measurement_sensitivity(&state_arr);
let h = jacobian_to_matrix::<M, N>(sensitivity);
// Measurement noise covariance
let r = std_to_covariance(self.measurement_model.measurement_noise_std());
// Innovation covariance: S = H * P * H^T + R
let s = &h * &self.covariance * h.transpose() + &r;
// Kalman gain: K = P * H^T * S^(-1)
let s_inv = s
.try_inverse()
.ok_or(EstimatorError::SingularInnovationCovariance)?;
let k = &self.covariance * h.transpose() * s_inv;
// Update state
self.state = &self.state + &k * innovation;
// Update covariance (Joseph form for numerical stability)
let i = OMatrix::<f64, Const<N>, Const<N>>::identity();
let i_kh = &i - &k * &h;
self.covariance = &i_kh * &self.covariance * i_kh.transpose() + &k * r * k.transpose();
// Check for numerical instability
for i in 0..N {
if !self.state[i].is_finite() {
return Err(EstimatorError::NumericalInstability(format!(
"state[{}] = {}",
i, self.state[i]
)));
}
}
Ok(())
}
fn estimate(&self) -> [f64; N] {
self.state.as_slice().try_into().unwrap()
}
}
impl<const N: usize, const M: usize, const C: usize, PM, MM> UncertaintyEstimator
for ExtendedKalmanFilter<N, M, C, PM, MM>
where
Const<N>: DimName,
Const<M>: DimName,
Const<C>: DimName,
PM: ProcessModel<N, C>,
MM: MeasurementModel<N, M>,
DefaultAllocator: Allocator<Const<N>>
+ Allocator<Const<N>, Const<N>>
+ Allocator<Const<M>>
+ Allocator<Const<M>, Const<N>>
+ Allocator<Const<N>, Const<M>>
+ Allocator<Const<M>, Const<M>>,
{
fn uncertainty(&self) -> Vec<f64> {
(0..N).map(|i| self.covariance[(i, i)].sqrt()).collect()
}
}