Beyond Linearity: EKF, UKF, and Manifolds

Where linearization lies and by exactly how much, what sigma points fix, and why a modern filter puts its mean on a Lie group and its covariance in the tangent space.

Two different ways to be wrong about a turn

Take the Kalman filter of Chapter 6, point it at Rusty's pose $(x, y, \theta)$, and hand it a compass. Drive Rusty in a circle. Two things go wrong, and they are not the same thing.

The first is quantitative. The motion model $x_t = g(u_t, x_{t-1})$ moves the robot along an arc, and arcs are curved: propagating a covariance through the tangent to that arc systematically misplaces the mean and misreports the spread. The error is small when the robot is confident and grows without limit as it becomes uncertain. It is an approximation with a size, and the size can be computed.

The second is categorical. Halfway round the circle Rusty's heading passes $\pi$ and the compass — like every real compass — starts reporting $-\pi$. The filter computes an innovation $z_t - \mu_t$ of nearly $-2\pi$, applies a gain of about a half, and swings its estimate a hundred and eighty degrees the wrong way round the world. No amount of better linearization touches this. It is not an approximation error at all; it is the arithmetic being wrong.

The chapter deals with them in that order.

Watch the sweep for a few passes before reading on. Three things are worth noticing.

The tangent moves. The EKF does not linearize once, globally; it re-fits its tangent at the current mean at every single step. That is its great strength — a fresh, locally-valid model each update — and the source of its most dangerous failure, because the quality of that model depends on where the mean happens to be, and the mean is the thing you are unsure about.

The scouts bend with the curve. The three green points are not samples. They are chosen deterministically so that their weighted mean and covariance reproduce the blue input exactly, and they are pushed through the true function. Nothing is ever differentiated.

The error has two factors, not one. Slide the operating point into a flat region and the orange and gray curves merge no matter how wide the input is. Park it in the bend and widen the input, and the gap opens as $\sigma^2$. Neither curvature nor spread alone predicts anything; their product does.

The mathematics

The nonlinear Gaussian system

Everything below concerns a system of the form

with $\varepsilon_t \sim \Normal(0, R_t)$ and $\delta_t \sim \Normal(0, Q_t)$ independent across time, exactly as in Chapter 6 except that $g$ and $h$ are now arbitrary smooth functions instead of matrices.

That single change destroys the property the Kalman filter was built on. A Gaussian pushed through a linear map is Gaussian; pushed through anything else it is not. So $\bel(x_t)$ is non-Gaussian from the first prediction step, and every filter in this chapter is a decision about which Gaussian to pretend with.

The extended Kalman filter

The move is the smallest one that could possibly work: replace $g$ and $h$ by their first-order Taylor expansions about the best point available, then run Chapter 6 unchanged.

The choice of expansion point matters and is not arbitrary: for a Gaussian, the most likely state is the mean, so the linearization is at least correct where the belief says the state most probably is. Thrun et al. make exactly this argument in §3.3.1, and the widget above is its refutation whenever the belief is wide.

That is Thrun et al.'s Table 3.3 verbatim. Lines 1 and 7 are the only places the true nonlinear functions appear; lines 3, 5, 6 and 8 are Chapter 6 with different letters.

The size of the lie

The widget's error readout is not a heuristic. It implements the following.

The unscented transform

The EKF's problem is that it approximates the function. The unscented transform approximates the distribution instead, and then uses the function exactly.

The idea, due to Julier and Uhlmann: choose a small set of points whose weighted sample mean and covariance reproduce $(\mu, \Sigma)$ exactly, push each one through $g$, and refit. Nothing is differentiated; $g$ is a black box that may be a lookup table, a ray-caster, or a physics engine.

The unscented Kalman filter is then nothing more than the Bayes filter with the UT used as the Gaussian pushforward engine in both steps.

Line 6 is worth a pause. In the EKF the gain is $\bar\Sigma_t H_t\T S_t^{-1}$, and $\bar\Sigma_t H_t\T$ is the linearized approximation of the state–measurement cross-covariance. The UKF computes that cross-covariance directly, from samples. It is the cleanest available statement of what a Kalman gain is: how much the state co-varies with the measurement, divided by how much the measurement varies with itself.

A worked example you can check by hand

A range–bearing sensor reports polar coordinates and the filter wants Cartesian. Take

independent. Because $r$ and $\theta$ are independent and $\E[\sin\theta] = \sin(\pi/2)\, e^{-\sigma_\theta^2/2}$ for a Gaussian $\theta$, the true mean of the $y$-component is available in closed form:

The EKF answers $f(\mu_r, \mu_\theta) = 1.000000$ — off by 3.4 cm at one metre of range, from a sensor whose bearing noise is a perfectly ordinary 15°.

Now the unscented transform with $n = 2$, $\alpha_{\mathrm{UT}} = 1$, $\beta = 0$, $\kappa = 1$, so $\lambda = 1$ and $\sqrt{n+\lambda} = \sqrt3$. The covariance is diagonal, so its Cholesky factor is $\diag(0.02,\, 0.261799)$ and the five points are:

which agrees with the analytic $0.966311$ to five decimals, from five evaluations of a sine.

The variances tell the same story from the other side. The true standard deviation of $y$ is $0.050680$ and the UT gives $0.051667$, a 2% overestimate. The EKF's is $\big((\partial y/\partial r)^2\sigma_r^2 + (\partial y/\partial\theta)^2\sigma_\theta^2\big)^{1/2}$, and at $\theta = \pi/2$ the second derivative is exactly zero, so the whole bearing contribution disappears and it reports $0.020000$ — two centimetres of uncertainty where there are five. In the $x$-component, where the bearing derivative is large, the same first-order formula over-states the spread instead: $0.261799$ against a true $0.253129$. Getting the variance wrong in both directions at once, from one linearization, is not an unusual outcome.

use approx::assert_relative_eq;
use ch07_nonlinear::ukf::{unscented_transform, UtParams};
use nalgebra::{SMatrix, SVector, Vector2};
use pr_core::geom::OnManifoldGaussian;

/// The chapter's worked example, pinned to five decimals.
///
/// Note the parameters: β = 0 here, not the usual 2. β deliberately inflates
/// the output covariance and would spoil a comparison against a closed form,
/// which is exactly why the book states it rather than leaving it at a default.
#[test]
fn polar_to_cartesian_matches_the_closed_form() {
    let sigma_theta: f64 = 15.0_f64.to_radians();
    let mean = Vector2::new(1.0, std::f64::consts::FRAC_PI_2);
    let cov = SMatrix::<f64, 2, 2>::from_diagonal(&Vector2::new(
        0.02 * 0.02,
        sigma_theta * sigma_theta,
    ));

    let polar_to_cartesian =
        |x: &SVector<f64, 2>| Vector2::new(x[0] * x[1].cos(), x[0] * x[1].sin());

    let prior = OnManifoldGaussian { mean, cov };
    let params = UtParams { alpha_ut: 1.0, beta: 0.0, kappa: 1.0 };
    let ut = unscented_transform(&prior, polar_to_cartesian, params);

    // E[r sin θ] = E[r] · sin(π/2) · exp(−σ²/2) for independent r, θ.
    let truth = (-0.5 * sigma_theta * sigma_theta).exp();
    assert_relative_eq!(truth, 0.966_311, epsilon = 1e-6);
    assert_relative_eq!(ut.mean[1], truth, epsilon = 1e-5);

    // The EKF's answer, for contrast: the mean pushed straight through f.
    let ekf_mean = polar_to_cartesian(&mean);
    assert_relative_eq!(ekf_mean[1], 1.0, epsilon = 1e-12);
    assert!((ekf_mean[1] - truth).abs() > 0.033); // 3.4 cm at 1 m

    // And the spread the EKF cannot see, because ∂y/∂θ = 0 at θ = π/2.
    assert_relative_eq!(ut.cov[(1, 1)].sqrt(), 0.051_667, epsilon = 1e-5);
    assert_relative_eq!(ut.cov[(0, 0)].sqrt(), 0.252_919, epsilon = 1e-5);
}

Where 2005 stops

Everything so far has been about approximating a function well. The rest of this chapter is about a mistake that survives any amount of good approximation.

Run it until Rusty crosses the seam. The orange filter is not badly approximating anything — its Jacobian is exactly $1$, its linearization is exact, and it is still catastrophically wrong. It computed $z_t - \mu_t \approx -2\pi$ and took a fraction of that step, and both of those operations are meaningless for an angle. Subtraction is not the metric on a circle, and a weighted average of two points on a circle is not a point on the circle.

On-manifold Gaussians

Definition. An on-manifold Gaussian on a $d$-dimensional manifold $\mathcal{M}$ is a pair $(\mu, \Sigma)$ with $\mu \in \mathcal{M}$ and $\Sigma \in \R^{d \times d}$, representing the distribution of the random element

The mean lives on the manifold; the covariance lives in the tangent space at that mean. This is sometimes called a concentrated Gaussian, and the name carries the caveat: the construction is only a good model of a distribution when $\Sigma$ is small enough that the tangent space is a fair picture of the neighbourhood. A heading with a standard deviation of $60°$ is not well described by any such object, and a filter is the wrong tool for it — that is Chapter 8's territory.

Definition (retraction). $\bplus : \mathcal{M} \times \R^d \to \mathcal{M}$ and $\bminus : \mathcal{M} \times \mathcal{M} \to \R^d$ satisfy, for all $x, y \in \mathcal{M}$ and small $\delta \in \R^d$:

with both maps smooth. The first says doing nothing does nothing; the second says $\bminus$ really does recover the displacement between two points; the third says the pair is locally inverse. That is all a Kalman filter ever needed from + and −, which is why the substitution works.

For $\SEtwo$ this book fixes the right (local, body-frame) convention throughout:

with $\exp$ and $\log$ as derived in Chapter 3, and the tangent ordered translation-first, $\delta = (\delta_x, \delta_y, \delta_\theta)$. A perturbation is something you do in the robot's own frame: drive forward $\delta_x$, slide left $\delta_y$, turn $\delta_\theta$. The left convention is equally valid and gives different Jacobians; Appendix C tabulates the correspondence.

Vector spaces are the special case $x \bplus \delta = x + \delta$, $y \bminus x = y - x$, which satisfies the axioms trivially. Everything that follows therefore contains Chapter 6 as an instance rather than replacing it.

The on-manifold EKF

Take the EKF algorithm and replace every state-space + with $\bplus$ and every state-space − with $\bminus$. Two lines change; nothing else does.

The covariance equations are untouched, because covariances were always tangent-space objects; we simply never said so. What does change is the meaning of the Jacobians, which are now derivatives in tangent coordinates:

Both the input and the output perturbation are tangent vectors, so $G_t$ is an honest $d \times d$ matrix even though $\mathcal{M}$ is not a vector space.

For Rusty's odometry step $g(u, x) = x \bplus u$ with $u = (v\,\Delta t,\, 0,\, \omega\,\Delta t)$, that derivative has a closed form and it is a small gem:

The UKF becomes UKF-M by the same substitution: draw the sigma points in the tangent space and retract them, $\mathcal{X}^{[\pm i]} = \mu \bplus \big(\pm\sqrt{(n{+}\lambda)\Sigma}\big)_i$. The one genuinely new ingredient is the recombination, because a weighted sum of manifold points is not defined. Instead the mean is characterised implicitly, as the point from which the weighted tangent displacements cancel, and found by iteration:

started at $\mathcal{Y}^{[0]}$ and repeated to convergence — three or four passes at filter scale. Everything downstream, including the cross-covariance, uses $\bminus$ in place of subtraction.

Error-state: linearize the small thing

There is a second, more practical way to arrive in the same place, and it is the one production systems actually use.

Split the state into a nominal part $\mu_t \in \mathcal{M}$, which carries the full trajectory and is propagated through the exact nonlinear model with no covariance attached, and an error state $\varepsilon_t = x_t \bminus \mu_t \in \R^d$, which is the only thing the Kalman machinery ever sees. The error is zero-mean by construction and stays small by construction, because every correction is immediately injected into the nominal state and the error reset to zero.

The two panels are the argument in one picture. On the left, seven metres of travel and a heading that sweeps the whole circle; on the right, an error that never leaves a hand-sized box. A first-order expansion of the dynamics of the left-hand quantity is a fiction. A first-order expansion of the dynamics of the right-hand quantity is very nearly exact, because the neglected term is again curvature times spread, and the spread here is the error's, not the state's.

Line 7 is the piece most implementations quietly skip. After injecting $\hat\varepsilon_t$ the remaining error is measured from a different nominal state, and the two tangent spaces differ by the adjoint of the injection. For a centimetre-scale correction $J$ is the identity to within rounding, which is why skipping it is usually harmless — and why it stops being harmless in Chapter 14, where a loop closure injects a correction big enough to matter.

Invariant errors, in one fact

Step 4 of the adjoint derivation above deserves to be finished.

Define the error not as a tangent-space difference but as a group element,

(the left-invariant form; the right-invariant analogue is $x_t \mu_t^{-1}$). Under noiseless body-frame odometry, both the truth and the nominal receive the same increment, $x_t = x_{t-1}\exp(u)$ and $\mu_t = \mu_{t-1}\exp(u)$, so

The error evolves autonomously: its update law involves only the control, never the estimate $\mu_{t-1}$. That is the group-affine property, and the invariant EKF is what you get by building a filter around it. Its Jacobians are constant along a trajectory, so the linearization point cannot drift away from the truth in the way that poisons a standard EKF, and Barrau and Bonnabel prove that the resulting filter is a locally stable observer with a domain of attraction independent of the trajectory — a guarantee no ordinary EKF has ever had.

We are not deriving the general machinery here; it needs left- and right-Jacobian bookkeeping that would swamp this chapter. But keep the fact, because Chapter 14 reads the EKF-SLAM consistency disaster back through exactly this lens, and Chapter 18 uses invariant filters as production VIO front-ends.

Implementation in Rust

The design goal is that one filter implementation serve the vector case and the manifold case, with the compiler enforcing the difference. That is a trait.

use nalgebra::SVector;

/// A manifold you can do estimation on. `D` is the tangent dimension.
///
/// Introduced in Chapter 3 and reused unchanged from here to the end of the
/// book: Chapter 9 samples motion noise through it, Chapter 15 retracts
/// Gauss-Newton steps with it, Chapter 18 instantiates it at SE(3) × ℝ⁹.
///
/// The three retraction axioms are not expressible in the type system, so
/// they live in a generic test harness (`manifold_axioms::<X, D>()`) that
/// every impl in the workspace is required to pass.
pub trait Manifold<const D: usize>: Clone {
    /// x ⊞ δ — move along the manifold by a tangent vector.
    fn boxplus(&self, delta: &SVector<f64, D>) -> Self;

    /// y ⊟ x — the tangent vector taking `rhs` to `self`.
    fn boxminus(&self, rhs: &Self) -> SVector<f64, D>;
}

/// A vector space is the manifold where ⊞ is +. This impl is what makes the
/// generic Ekf below reduce *exactly* to Chapter 6's Kf on a linear model —
/// a claim pinned by a test, not by assertion.
impl<const D: usize> Manifold<D> for SVector<f64, D> {
    fn boxplus(&self, delta: &SVector<f64, D>) -> Self {
        self + delta
    }
    fn boxminus(&self, rhs: &Self) -> SVector<f64, D> {
        self - rhs
    }
}

/// SE(2) with the book's right/local convention. `SE2`, `exp`, `log` and
/// `adjoint` are Chapter 3's hand-rolled types, not a dependency.
impl Manifold<3> for SE2 {
    fn boxplus(&self, delta: &SVector<f64, 3>) -> Self {
        self * SE2::exp(delta)
    }
    fn boxminus(&self, rhs: &Self) -> SVector<f64, 3> {
        (rhs.inverse() * self).log()
    }
}

/// Mean on the manifold, covariance in the tangent space at that mean.
#[derive(Clone, Debug)]
pub struct OnManifoldGaussian<X: Manifold<D>, const D: usize> {
    pub mean: X,
    pub cov: SMatrix<f64, D, D>,
}

With that in place the EKF is written once. Note that the models are traits too, so a measurement of the wrong dimension is a compile error rather than a runtime surprise.

use nalgebra::{SMatrix, SVector};
use pr_core::geom::{Manifold, OnManifoldGaussian};

pub trait MotionModel<X: Manifold<D>, const D: usize> {
    type Control;
    /// The true nonlinear map. The *mean* goes through this, never through G.
    fn g(&self, x: &X, u: &Self::Control) -> X;
    /// G_t = ∂( g(x ⊞ ε, u) ⊟ g(x, u) ) / ∂ε |₀ — in tangent coordinates.
    fn jacobian(&self, x: &X, u: &Self::Control) -> SMatrix<f64, D, D>;
    fn noise(&self, u: &Self::Control) -> SMatrix<f64, D, D>;
}

pub trait MeasurementModel<X: Manifold<D>, const D: usize, const M: usize> {
    fn h(&self, x: &X) -> SVector<f64, M>;
    fn jacobian(&self, x: &X) -> SMatrix<f64, M, D>;
    fn noise(&self) -> SMatrix<f64, M, M>;
    /// Measurement-space residual. Override for bearings — the measurement
    /// space can be a manifold too, and usually is when the state is.
    fn residual(&self, z: &SVector<f64, M>, z_hat: &SVector<f64, M>) -> SVector<f64, M> {
        z - z_hat
    }
}

pub struct Ekf<X: Manifold<D>, G, H, const D: usize, const M: usize> {
    pub belief: OnManifoldGaussian<X, D>,
    pub motion: G,
    pub meas: H,
}

impl<X, G, H, const D: usize, const M: usize> Ekf<X, G, H, D, M>
where
    X: Manifold<D>,
    G: MotionModel<X, D>,
    H: MeasurementModel<X, D, M>,
{
    pub fn predict(&mut self, u: &G::Control) {
        let g_t = self.motion.jacobian(&self.belief.mean, u);
        // The mean goes through the true model; only the covariance is linearized.
        self.belief.mean = self.motion.g(&self.belief.mean, u);
        self.belief.cov = g_t * self.belief.cov * g_t.transpose() + self.motion.noise(u);
    }

    /// Returns the normalized innovation squared — the NIS gate of Chapter 6.
    pub fn correct(&mut self, z: &SVector<f64, M>) -> f64 {
        let h_t = self.meas.jacobian(&self.belief.mean);
        let y = self.meas.residual(z, &self.meas.h(&self.belief.mean));
        let s = h_t * self.belief.cov * h_t.transpose() + self.meas.noise();
        let s_inv = s.try_inverse().expect("innovation covariance must be invertible");
        let k = self.belief.cov * h_t.transpose() * s_inv;

        // The gain produces a *tangent vector*. On SE(2) adding it to a pose
        // would not even type-check as geometry; ⊞ is the only way back.
        self.belief.mean = self.belief.mean.boxplus(&(k * y));

        let i_kh = SMatrix::<f64, D, D>::identity() - k * h_t;
        self.belief.cov = i_kh * self.belief.cov * i_kh.transpose()
            + k * self.meas.noise() * k.transpose();

        (y.transpose() * s_inv * y)[(0, 0)]
    }
}

The unscented filter needs no Jacobians and therefore no jacobian methods — only g, h, the noises, and a Cholesky. On a manifold the sigma points are retracted and the mean is found by the iteration derived above.

use nalgebra::{SMatrix, SVector};
use pr_core::geom::{Manifold, OnManifoldGaussian};

#[derive(Clone, Copy, Debug)]
pub struct UtParams {
    /// Spread of the scouts. Subscripted `_ut` book-wide so it never reads as
    /// one of the motion-noise α's of Chapter 9.
    pub alpha_ut: f64,
    pub beta: f64,
    pub kappa: f64,
}

impl UtParams {
    /// Julier's defaults: κ = 3 − n, β = 2 (exact 4th moment for a Gaussian).
    pub fn julier(d: usize) -> Self {
        Self { alpha_ut: 1.0, beta: 2.0, kappa: 3.0 - d as f64 }
    }
}

/// 2D+1 sigma points, drawn in the tangent space and retracted onto `M`.
///
/// The weights are `Vec`s rather than `[f64; 2 * D + 1]`: arithmetic in array
/// lengths needs `generic_const_exprs`, which is still unstable. One
/// allocation per call, hoisted out of the loop by the caller in practice.
pub fn sigma_points<X: Manifold<D>, const D: usize>(
    belief: &OnManifoldGaussian<X, D>,
    p: UtParams,
) -> (Vec<X>, Vec<f64>, Vec<f64>) {
    let n = D as f64;
    let lambda = p.alpha_ut * p.alpha_ut * (n + p.kappa) - n;
    let scaled = belief.cov * (n + lambda);
    let l = scaled.cholesky().expect("covariance must be positive definite").l();

    let mut points = Vec::with_capacity(2 * D + 1);
    points.push(belief.mean.clone());
    for i in 0..D {
        let col: SVector<f64, D> = l.column(i).into_owned();
        points.push(belief.mean.boxplus(&col));
        points.push(belief.mean.boxplus(&(-col)));
    }

    let w = 1.0 / (2.0 * (n + lambda));
    let mut wm = vec![w; 2 * D + 1];
    let mut wc = wm.clone();
    wm[0] = lambda / (n + lambda);
    wc[0] = wm[0] + 1.0 - p.alpha_ut * p.alpha_ut + p.beta;
    (points, wm, wc)
}

/// The weighted mean of manifold points: the fixed point of ⊟-then-⊞.
///
/// A vector-space UKF gets this from a single weighted sum. On a curved space
/// there is no closed form, so we iterate — four passes is convergence to
/// machine precision for filter-scale covariances.
pub fn manifold_mean<X: Manifold<D>, const D: usize>(points: &[X], w: &[f64]) -> X {
    let mut mean = points[0].clone();
    for _ in 0..8 {
        let mut delta = SVector::<f64, D>::zeros();
        for (p, wi) in points.iter().zip(w) {
            delta += *wi * p.boxminus(&mean);
        }
        mean = mean.boxplus(&delta);
        if delta.norm() < 1e-12 {
            break;
        }
    }
    mean
}

/// The unscented transform on a manifold: Algorithm `Unscented_transform`,
/// with ⊞ for the points and `manifold_mean` for the recombination.
pub fn unscented_transform<X, Y, F, const D: usize, const E: usize>(
    belief: &OnManifoldGaussian<X, D>,
    f: F,
    p: UtParams,
) -> OnManifoldGaussian<Y, E>
where
    X: Manifold<D>,
    Y: Manifold<E>,
    F: Fn(&X) -> Y,
{
    let (points, wm, wc) = sigma_points(belief, p);
    let images: Vec<Y> = points.iter().map(&f).collect();
    let mean = manifold_mean(&images, &wm);

    let mut cov = SMatrix::<f64, E, E>::zeros();
    for (y, wi) in images.iter().zip(&wc) {
        let e = y.boxminus(&mean);
        cov += *wi * e * e.transpose();
    }
    OnManifoldGaussian { mean, cov }
}

A compile error worth printing

The book promised that the type system would do the bookkeeping Thrun does by hand. Here is it doing it. A compass on a planar robot is a one-dimensional measurement of a three-dimensional state, so its model is MeasurementModel<SE2, 3, 1>. Hand correct a two-vector:

let mut filter: Ekf<SE2, Unicycle, Compass, 3, 1> = Ekf::new(prior, Unicycle::default(), Compass::new(6f64.to_radians()));

// The reading came from a range–bearing sensor by mistake.
filter.correct(&SVector::<f64, 2>::new(4.10, 0.62));

error[E0308]: mismatched types
  --> demos/ch07-demo/src/bin/typecheck_demo.rs:14:20
   |
14 |     filter.correct(&SVector::<f64, 2>::new(4.10, 0.62));
   |            ------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `&Matrix<f64, Const<1>, Const<1>, ArrayStorage<f64, 1, 1>>`,
   |            |                                            found `&Matrix<f64, Const<2>, Const<1>, ArrayStorage<f64, 2, 1>>`
   |            arguments to this method are incorrect
   |
   = note: expected reference `&SVector<f64, 1>`
              found reference `&SVector<f64, 2>`

The same discipline catches the subtler mistake, a Jacobian of the wrong shape:

error[E0308]: mismatched types
  --> crates/ch07_nonlinear/src/models.rs:88:9
   |
87 |     fn jacobian(&self, _x: &SE2) -> SMatrix<f64, 1, 3> {
   |                                     ------------------ expected `Matrix<f64, Const<1>, Const<3>, _>` because of return type
88 |         SMatrix::<f64, 1, 2>::new(0.0, 0.0)
   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected `1×3` matrix, found `1×2` matrix

A dimension mismatch between a state, its Jacobian, and a measurement is the single most common source of silent filter failure in the wild — it produces a filter that runs, converges to something, and is wrong. Here it does not compile.

The error-state filter, and the test that pins it

use nalgebra::{SMatrix, SVector};
use pr_core::geom::{Manifold, SE2};

/// Error-state EKF: a nominal state on the manifold plus a tangent-space
/// error that is filtered, injected, and reset every update.
pub struct Eskf<X: Manifold<D>, const D: usize> {
    pub nominal: X,
    pub cov: SMatrix<f64, D, D>,
    /// Keep the exact reset Jacobian. Off by default in most codebases;
    /// on by default here, because Chapter 14 needs to be able to turn it off
    /// and show the difference.
    pub exact_reset: bool,
}

impl Eskf<SE2, 3> {
    /// µ̄ = µ ⊞ u,  Σ̄ = G Σ Gᵀ + R  with G = Ad_exp(u)⁻¹ — no state in sight.
    pub fn predict(&mut self, u: &SVector<f64, 3>, r: &SMatrix<f64, 3, 3>) {
        let g = SE2::exp(u).inverse().adjoint();
        self.nominal = self.nominal.boxplus(u);
        self.cov = g * self.cov * g.transpose() + r;
    }

    /// The inject-and-reset cycle. `epsilon` is returned, not stored: an
    /// error state that survives its own update is a bug, and making the
    /// caller receive it rather than read it is how the API says so.
    pub fn inject_and_reset(&mut self, epsilon: SVector<f64, 3>) -> SVector<f64, 3> {
        self.nominal = self.nominal.boxplus(&epsilon);
        if self.exact_reset {
            let j = SE2::exp(&epsilon).inverse().adjoint();
            self.cov = j * self.cov * j.transpose();
        }
        epsilon
    }
}

/// The generic filter must contain Chapter 6 exactly, not approximately.
///
/// On a linear-Gaussian model, `Ekf<SVector<f64, 4>>` and `ch06::Kf` are the
/// same algorithm with different type parameters, so they must agree to
/// floating-point noise. If this ever fails, the trait abstraction has
/// silently changed the mathematics.
#[test]
fn ekf_reduces_to_kf_on_a_linear_model() {
    let (mut ekf, mut kf) = constant_velocity_pair();
    let mut rng = Pcg64::seed_from_u64(0xC0FFEE);

    for _ in 0..200 {
        let u = SVector::<f64, 2>::new(0.1, -0.05);
        let z = simulate_measurement(&mut rng);
        ekf.predict(&u);
        kf.predict(&u);
        ekf.correct(&z);
        kf.correct(&z);

        assert_relative_eq!(ekf.belief.mean, kf.mean(), epsilon = 1e-12);
        assert_relative_eq!(ekf.belief.cov, kf.cov(), epsilon = 1e-12);
    }
}

/// And the manifold impl must obey the retraction axioms, checked on random
/// poses rather than argued about in prose.
#[test]
fn se2_satisfies_the_retraction_axioms() {
    let mut rng = Pcg64::seed_from_u64(7);
    for _ in 0..10_000 {
        let (x, y) = (random_pose(&mut rng), random_pose(&mut rng));
        let d = random_twist(&mut rng, 0.3);

        assert_relative_eq!(x.boxplus(&SVector::zeros()), x, epsilon = 1e-12);
        assert_relative_eq!(x.boxplus(&y.boxminus(&x)), y, epsilon = 1e-10);
        assert_relative_eq!(x.boxplus(&d).boxminus(&x), d, epsilon = 1e-10);
    }
}

Putting it together: the figure-eight lab

cargo run --example figure_eight -p ch07_nonlinear drives Rusty three times around a lemniscate — seven metres wide, sixteen seconds a lap, heading sweeping the whole circle and crossing the $\pm\pi$ seam six times — and runs four filters on the identical seeded log: a compass at 10 Hz with $\sigma = 6°$, a beacon fix every second with $\sigma = 25$ cm, and body-frame odometry noise.

All four use the same exact motion function, $\mu \bplus u$; they differ only in how they represent the state and compute residuals.