diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..83080d5 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,67 @@ +name: CI + +on: + push: + branches: [master, develop] + pull_request: + +jobs: + test: + name: Tests (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest] + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + - uses: cachix/install-nix-action@v31 + + - name: Rust tests + run: nix develop .#ci --command cargo test + + - name: Rust examples + run: | + nix develop .#ci --command bash -euo pipefail -c ' + for f in examples/*.rs; do + e=$(basename "$f" .rs) + echo "== $e" + cargo run --quiet --example "$e" + done' + + - name: Build Python bindings + run: | + nix develop .#ci --command bash -euo pipefail -c ' + python -m venv --system-site-packages .venv + . .venv/bin/activate + maturin develop --release' + + - name: Python tests + run: nix develop .#ci --command bash -euo pipefail -c '. .venv/bin/activate && python -m pytest tests_py' + + - name: Python examples + run: | + nix develop .#ci --command bash -euo pipefail -c ' + . .venv/bin/activate + cd examples + for f in *.py; do + # renders a video and needs manim, which the CI shell leaves out + [ "$f" = manim_two_mass_spring_damper.py ] && continue + echo "== $f" + python "$f" + done' + + lint: + name: Lint + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: cachix/install-nix-action@v31 + + - name: Clippy + run: nix develop .#ci --command cargo clippy --all-targets + + # the code has not been run through rustfmt yet; report only for now + - name: Rustfmt + continue-on-error: true + run: nix develop .#ci --command cargo fmt --check diff --git a/README.md b/README.md index fdad1d3..0c914b0 100644 --- a/README.md +++ b/README.md @@ -3,9 +3,12 @@ A lightweight library for some basic filtering methods ## Documentation -This crate implements small, statically-sized Kalman and Extended Kalman Filter -implementations using nalgebra's `SMatrix` for predictable, allocation-free -linear algebra. +This crate implements small Kalman and Extended Kalman Filters on nalgebra +matrices. The models are generic over the dimension type, so the same +implementation serves statically sized matrices (`KalmanFilterNM`, +`LinearUpdateNM`, ... for predictable, allocation-free linear algebra) +and runtime sized ones (`KalmanFilterDyn`, `ExtendedKalmanFilterDyn`), which +the Python bindings wrap. Generate API documentation locally with: @@ -20,6 +23,15 @@ The main modules are: - `crate::measurement` — measurement/observation helpers. - `crate::models` — traits for state transition and measurement models used by EKF. +## Development environment + +A [Nix](https://nixos.org) flake provides the Rust toolchain, Python (with `numpy`, `pytest` and [Manim Community](https://www.manim.community) including its native dependencies such as cairo and pango) and `maturin`: + +```bash +nix develop +cargo test +``` + ## Python Build and install the bindings into the current virtualenv with [maturin](https://www.maturin.rs): @@ -31,7 +43,44 @@ maturin develop --release # editable-style install for development pip install . # builds and installs a wheel (needs a Rust toolchain) ``` -Then `import determined as dt`. Build a distributable wheel with `maturin build --release`. +Then `import determined as dt`. + +Six filters are available from Python: + +| Class | Models | Maths runs in | +|---|---|---| +| `dt.KalmanFilter(state, *, f, q, h, r)` | numpy matrices | Rust only | +| `dt.ExtendedKalmanFilter(state, *, f, jac_f, q, h, jac_h, r)` | Python functions | Rust, calling back into Python for the model functions | +| `dt.UnscentedKalmanFilter(state, *, f, q, h, r, alpha=1, beta=2, kappa=0)` | Python functions, no Jacobians | Rust, calling back into Python for the model functions | +| `dt.UDKalmanFilter(state, *, f, q, h, r)` | numpy matrices | Rust only, in UD factorised form | +| `dt.UDExtendedKalmanFilter(state, *, f, jac_f, q, h, jac_h, r)` | Python functions | Rust, in UD factorised form, calling back into Python for the model functions | +| `dt.CustomFilter(update)` | a Python `UpdateModel` you implement | Python (the filter only delegates `predict`/`update`) | + +```python +kf = dt.KalmanFilter(state, f=F, q=Q, h=H, r=R) +kf.predict(dt.Epoch(1)) +kf.update(dt.Observation(z, epoch=dt.Epoch(1))) +``` + +`dt.ExtendedKalmanFilter` runs the filter maths in Rust and calls back into Python for the model: + +```python +ekf = dt.ExtendedKalmanFilter(state, f=f, jac_f=jac_f, q=q, h=h, jac_h=jac_h, r=r) +ekf.predict(dt.Epoch(1)) +ekf.update(dt.Observation(z, epoch=dt.Epoch(1))) +``` + +- **Model functions** are `f(x, t)`, `jac_f(x, t)`, `h(x, t)` and `jac_h(x, t)`. `x` is the state value (a numpy array) and `t` the integer epoch: `f` advances the state from epoch `t` to `t + 1`, and `h` is evaluated at the measurement epoch. They may be time-varying, and may return anything array-like, so Jacobians can come from an autodiff library (e.g. `jax.jacfwd`, which differentiates `x` and not `t`). +- **Same implementation as Rust:** each Python filter wraps the crate's `...Dyn` filter of the same name (`KalmanFilterDyn`, `ExtendedKalmanFilterDyn`, `UnscentedKalmanFilterDyn`, `UDKalmanFilterDyn`, `UDExtendedKalmanFilterDyn`), so they behave exactly like the Rust filters. +- **UD filters:** the covariance is held as factors `P = U D Uᵀ` (`U` unit upper triangular, `D` diagonal), propagated by weighted modified Gram–Schmidt and updated one scalar measurement at a time with the modified Agee–Turner update, following Ramos et al., [*A summary on the UD Kalman Filter*](https://arxiv.org/abs/2203.06105). A non-diagonal `R` is decorrelated first. This keeps `P` symmetric positive semi-definite by construction; the results equal the standard filters up to rounding. The state still carries `P = U D Uᵀ`, and setting it from outside refactorises. +- **Unscented filter:** `f` and `h` are evaluated at 2N + 1 sigma points per step (van der Merwe scaled points), so they must not keep state between calls. The default `alpha=1, beta=2, kappa=0` keeps all sigma point weights non-negative; a very small `alpha` (e.g. `1e-3`) gives a large negative centre weight and is numerically fragile. The covariance must be strictly positive definite (and `kappa > -N`): a semi-definite covariance, such as one with a component known exactly, or a missing one makes `predict`/`update` log an error and leave the state unchanged. +- **Epochs:** both filters advance the state one epoch at a time and do nothing when asked for the current epoch. Predicting backwards uses `F⁻¹` for `KalmanFilter`; the other filters log a warning and leave the state unchanged. +- **Errors in model functions:** if `f`, `h` or a Jacobian raises (or returns the wrong shape), the exception propagates and the filter state is left as it was before the call. +- **Keyword-only arguments:** constructors take their first argument positionally and the rest by keyword, e.g. `dt.State(value, covariance=cov, epoch=epoch)`. +- **States are live views.** To avoid copying, the `State` returned by `predict`, `update` and `.state` shares memory with the filter's internal state, so it changes on the next call (and is the same object you passed in). Copy the arrays (`state.value.copy()`) if you need to keep a snapshot. +- **Logging:** messages from the Rust side go through Python's `logging` under the name `determined`, e.g. `logging.getLogger("determined").setLevel(logging.DEBUG)`. + +See `examples/ekf.py`, `examples/ukf.py` and `examples/ud.py` (Rust equivalents: `simple_ekf.rs`, `simple_ukf.rs`, `simple_ud.rs`). Build a distributable wheel with `maturin build --release`. ## Examples spring damper example @@ -46,3 +95,7 @@ To run the visualisation example using `python` and `manim` use: manim -pql examples/manim_two_mass_spring_damper.py TwoMassSpringDamperScene ``` +## Roadmap + +- **Gating:** expose the innovation `y = z - h(x)`, its covariance `S` and the normalised innovation squared `yᵀ S⁻¹ y` from `update`, so outlier measurements can be rejected against a chi-squared threshold. Not exposed yet. +- **Automatic Jacobians:** optional `jac_f` / `jac_h` computed by autodiff or finite differences. diff --git a/examples/algorithm.py b/examples/algorithm.py index a05edd0..d8a3e8c 100644 --- a/examples/algorithm.py +++ b/examples/algorithm.py @@ -7,10 +7,10 @@ value = np.array([1.0, 0.0]) epoch = dt.Epoch(0) -state = dt.State(value, np.eye(value.size), epoch) -obs = dt.Observation(np.array([0.0, 1.0]), epoch) +state = dt.State(value, covariance=np.eye(value.size), epoch=epoch) +obs = dt.Observation(np.array([0.0, 1.0]), epoch=epoch) -filter = dt.KalmanFilter(update) +filter = dt.CustomFilter(update) if __name__ == "__main__": diff --git a/examples/ekf.py b/examples/ekf.py new file mode 100644 index 0000000..de00114 --- /dev/null +++ b/examples/ekf.py @@ -0,0 +1,40 @@ +"""Extended Kalman filter tracking a static 2D target from range-only measurements. + +The model functions take the state value `x` and the integer epoch `t`, and +return plain numpy arrays, so the Jacobians could equally come from an autodiff +library, e.g. `jac_h = jax.jacfwd(h)` (which differentiates `x`, not `t`). +""" +import numpy as np + +import determined as dt + +# initial guess (2, 2) with a wide covariance; the target is really at (3, 4) +state = dt.State(np.array([2.0, 2.0]), covariance=np.eye(2) * 4.0, epoch=dt.Epoch(0)) + + +def f(x, t): + return x # static target + + +def jac_f(x, t): + return np.eye(2) + + +def h(x, t): + return np.array([np.linalg.norm(x)]) + + +def jac_h(x, t): + return (x / np.linalg.norm(x))[None, :] + + +ekf = dt.ExtendedKalmanFilter( + state, + f=f, jac_f=jac_f, q=np.zeros((2, 2)), + h=h, jac_h=jac_h, r=np.array([[0.01]]), +) + +for t in range(1, 6): + ekf.predict(dt.Epoch(t)) + estimate = ekf.update(dt.Observation(np.array([5.0]), epoch=dt.Epoch(t))) + print(f"t={t} x={estimate.value} range={np.linalg.norm(estimate.value):.4f}") diff --git a/examples/manim_two_mass_spring_damper.py b/examples/manim_two_mass_spring_damper.py index fbf6cb1..342570d 100644 --- a/examples/manim_two_mass_spring_damper.py +++ b/examples/manim_two_mass_spring_damper.py @@ -3,7 +3,7 @@ This scene runs the same simulation as the Python example but renders a textbook-style diagram: two masses as boxes and springs as zig-zag lines. -It uses the Rust `PyKalmanFilter` (PyO3) to compute the filter estimates. +The filter estimates are computed in Rust by `determined.KalmanFilter`. Run (recommended, low-quality preview): pip install manim # community edition @@ -25,10 +25,6 @@ import determined as dt -from transition import LinearTransition -from measurement import LinearMeasurement -from update import LinearUpdate - def setup_two_mass_spring_damper_params(): m = 0.1 @@ -37,7 +33,7 @@ def setup_two_mass_spring_damper_params(): b = 0.0036 x0 = np.array([1.0, 0.0, 0.0, 0.0], dtype=float) - initial_state = dt.State(x0, np.eye(x0.size), dt.Epoch(0)) + initial_state = dt.State(x0, covariance=np.eye(x0.size), epoch=dt.Epoch(0)) a = np.array([ [0.0, 1.0, 0.0, 0.0], @@ -54,9 +50,6 @@ def setup_two_mass_spring_damper_params(): sigma_q = np.diag([0.3 * _dt, 0.05 * _dt, 0.3 * _dt, 0.05 * _dt]) q = sigma_q @ sigma_q.T - transition_model = LinearTransition(f, q, initial_state) - transition = dt.TransitionModel(transition_model, initial_state) - # measurement model h = np.array([ [1.0, 0.0, 0.0, 0.0], @@ -66,15 +59,8 @@ def setup_two_mass_spring_damper_params(): sigma_r = np.diag([0.1, 0.1]) r = sigma_r @ sigma_r.T - measurement_model = LinearMeasurement(h, r) - measurement = dt.MeasurementModel(measurement_model) - - # update model - update_model = LinearUpdate(transition, measurement) - update = dt.UpdateModel(update_model, transition, measurement) - # filter - kf = dt.KalmanFilter(update) + kf = dt.KalmanFilter(initial_state, f=f, q=q, h=h, r=r) # system models process_noise=0.00, @@ -100,7 +86,7 @@ def simulate_rust(num_steps=100): x_true = model(x_true) # true system dynamics + process noise z = h(x_true) # measurement + measurement noise - _z = dt.Observation(z, dt.Epoch(step)) + _z = dt.Observation(z, epoch=dt.Epoch(step)) state_vec = kf.update(_z) traj_true[step] = x_true diff --git a/examples/measurement.py b/examples/measurement.py index e6ad575..3007050 100644 --- a/examples/measurement.py +++ b/examples/measurement.py @@ -11,12 +11,15 @@ def __init__(self, h, r): def projection(self, state: dt.State): z = self.h @ state.value - return dt.Observation(z, state.epoch) + return dt.Observation(z, epoch=state.epoch) def inverse(self, observation: dt.Observation): x = self.h_inv @ observation.value cov = np.eye(x.size) - return dt.State(x, cov, observation.epoch) + return dt.State(x, covariance=cov, epoch=observation.epoch) + + def noise(self): + return self.r def jacobian(self, state: dt.State): return self.jac @@ -24,7 +27,7 @@ def jacobian(self, state: dt.State): value = np.array([1.0, 0.0]) epoch = dt.Epoch(0) -state = dt.State(value, np.eye(value.size), epoch) +state = dt.State(value, covariance=np.eye(value.size), epoch=epoch) h = np.array([ [1.0, 0.0], @@ -39,6 +42,7 @@ def jacobian(self, state: dt.State): obs = measurement.projection(state) inv = measurement.inverse(obs) jac = measurement.jacobian(state) + r = measurement.noise() print(state) print(obs) diff --git a/examples/simple_ekf.rs b/examples/simple_ekf.rs new file mode 100644 index 0000000..067e5a8 --- /dev/null +++ b/examples/simple_ekf.rs @@ -0,0 +1,51 @@ +//! Extended Kalman filter tracking a static 2D target from range-only +//! measurements, a nonlinear sensor h(x) = |x|. + +use determined::algorithms::ExtendedKalmanFilterNM; +use determined::common::na; +use determined::epoch::Epoch; +use determined::filter::Filter; +use determined::measurement::Observation; +use determined::models::*; +use determined::state::State; + +fn main() { + // initial guess (2, 2) with a wide covariance; the target is really at (3, 4) + let state = State::>::new( + vec![2.0, 2.0], + vec![vec![4.0, 0.0], vec![0.0, 4.0]], + Epoch::new(0), + ) + .ptr(); + + // static target: x(t + 1) = x(t) + let transition = NonLinearTransitionN::<2>::new( + state.clone(), + Box::new(|state| state.value), + Box::new(|_| na::SMatrix::::identity()), + na::SMatrix::::zeros(), + ); + + // range sensor and its Jacobian + let measurement = NonLinearMeasurementNM::<2, 1>::new( + Box::new(|state| na::SMatrix::::new(state.value.norm())), + Box::new(|state| { + let range = state.value.norm(); + na::SMatrix::::new(state.value[0] / range, state.value[1] / range) + }), + na::SMatrix::::new(0.01), + ); + + let mut filter = + ExtendedKalmanFilterNM::new(NonLinearUpdateNM::<2, 1>::new(measurement, transition)); + + for t in 1..=5 { + let obs = Observation { + value: na::SMatrix::::new(5.0), + epoch: Epoch::new(t), + }; + filter.predict(&obs.epoch); + let estimate = filter.update(&obs); + println!("t={} {}", t, estimate.read().unwrap()); + } +} diff --git a/examples/simple_kalman.rs b/examples/simple_kalman.rs new file mode 100644 index 0000000..79d8643 --- /dev/null +++ b/examples/simple_kalman.rs @@ -0,0 +1,58 @@ +use determined::algorithms::KalmanFilterNM; +use determined::common::na; +use determined::epoch::Epoch; +use determined::filter::Filter; +use determined::measurement::Observation; +use determined::models::*; +use determined::state::State; + +fn main() { + spdlog::info!("Setting up initial state..."); + // initial state + let state = State::>::new( + vec![1.0, 1.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) + .ptr(); + + spdlog::info!("Setting up state dynamics..."); + // Linear Transition model + let f = na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[1.5, 2.0])); + let q = &na::SMatrix::::identity() * 0.05; + + let transition = LinearTransitionN::<2>::new(state.clone(), f, q); + + spdlog::info!("Setting up measurement model..."); + // Linear Measurement model + let h = na::SMatrix::::identity(); + let r = &na::SMatrix::::identity() * 0.01; + + let measurement = LinearMeasurementNM::<2, 2>::new( + h, // observation model + r, // observation noise + ); + + spdlog::info!("Building kalman filter..."); + // Kalman Update Model + let update = LinearUpdateNM::<2, 2>::new(measurement, transition); + + // Filter wrapper + let mut filter = KalmanFilterNM::new(update); + + spdlog::info!("Generating dummy observation..."); + let obs = Observation { + value: na::SMatrix::::from_row_slice(&[1.0, 1.0]), + epoch: Epoch::new(0), + }; + + spdlog::info!("Predicting..."); + // predict + let _ = filter.predict(&Epoch::new(1)); + + spdlog::info!("Updating..."); + // measurement update + let _ = filter.update(&obs); + + spdlog::info!("Done!"); +} diff --git a/examples/simple_ud.rs b/examples/simple_ud.rs new file mode 100644 index 0000000..439ba44 --- /dev/null +++ b/examples/simple_ud.rs @@ -0,0 +1,86 @@ +//! UD Kalman filter tracking a target moving at constant velocity from +//! correlated position measurements, alongside the standard Kalman filter. +//! The UD filter propagates the factors P = U D U^T, keeping P symmetric +//! positive semi-definite by construction, and gives the same estimates. + +use determined::algorithms::{KalmanFilterNM, UDKalmanFilterNM}; +use determined::common::na; +use determined::epoch::Epoch; +use determined::filter::Filter; +use determined::measurement::Observation; +use determined::models::*; +use determined::state::State; + +fn main() { + // state x = (px, py, vx, vy), one epoch per second, position measured + let f = na::SMatrix::::new( + 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, + ); + let q = na::SMatrix::::from_diagonal(&na::SVector::::new( + 0.01, 0.01, 0.001, 0.001, + )); + let h = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0); + // correlated measurement noise: decorrelated internally by the UD filter + let r = na::SMatrix::::new(0.04, 0.03, 0.03, 0.09); + + let initial = || { + State::>::new( + vec![0.0; 4], + vec![ + vec![10.0, 0.0, 0.0, 0.0], + vec![0.0, 10.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 0.0], + vec![0.0, 0.0, 0.0, 1.0], + ], + Epoch::new(0), + ) + .ptr() + }; + + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<4, 2>::new( + LinearMeasurementNM::<4, 2>::new(h, r), + LinearTransitionN::<4>::new(initial(), f, q), + )); + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<4, 2>::new( + LinearMeasurementNM::<4, 2>::new(h, r), + LinearTransitionN::<4>::new(initial(), f, q), + )); + + let mut truth = na::SVector::::new(0.0, 0.0, 1.0, 0.5); + for t in 1..=30 { + truth = f * truth; + // deterministic stand-in for sensor noise + let noise = + na::SVector::::new(0.2 * (1.3 * t as f64).sin(), 0.3 * (0.7 * t as f64).cos()); + let obs = Observation { + value: h * truth + noise, + epoch: Epoch::new(t), + }; + + let a = ud.update(&obs); + let b = kf.update(&obs); + if t % 10 == 0 { + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + println!( + "t={:2} estimate={:.3?} truth={:.3?}", + t, + a.value.as_slice(), + truth.as_slice() + ); + println!( + " |UD - KF| state={:.1e} covariance={:.1e}", + (a.value - b.value).abs().max(), + (a.covariance() - b.covariance()).abs().max() + ); + } + } + + let d: Vec = ud + .update + .factors + .d + .iter() + .map(|d| format!("{:.3e}", d)) + .collect(); + println!("UD factors: D = [{}]", d.join(", ")); +} diff --git a/examples/simple_ukf.rs b/examples/simple_ukf.rs new file mode 100644 index 0000000..922588d --- /dev/null +++ b/examples/simple_ukf.rs @@ -0,0 +1,72 @@ +//! Unscented Kalman filter tracking a target moving at constant velocity with +//! a range-and-bearing sensor (a radar). No Jacobians are needed: sigma points +//! are pushed through f and h instead. + +use determined::algorithms::UnscentedKalmanFilterNM; +use determined::common::na; +use determined::epoch::Epoch; +use determined::filter::Filter; +use determined::measurement::Observation; +use determined::models::*; +use determined::state::State; + +fn main() { + // state x = (px, py, vx, vy), one epoch per second + let f = na::SMatrix::::new( + 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, + ); + let q = na::SMatrix::::from_diagonal(&na::SVector::::new( + 0.01, 0.01, 0.001, 0.001, + )); + let r = na::SMatrix::::from_diagonal(&na::SVector::::new( + 0.1 * 0.1, + 0.01 * 0.01, + )); + let radar = + |x: &na::SVector| na::SVector::::new(x[0].hypot(x[1]), x[1].atan2(x[0])); + + // the target starts at (10, 5) moving at (-0.5, 0.3); the initial guess is poor + let mut truth = na::SVector::::new(10.0, 5.0, -0.5, 0.3); + let state = State::>::new( + vec![9.0, 6.0, 0.0, 0.0], + vec![ + vec![4.0, 0.0, 0.0, 0.0], + vec![0.0, 4.0, 0.0, 0.0], + vec![0.0, 0.0, 1.0, 0.0], + vec![0.0, 0.0, 0.0, 1.0], + ], + Epoch::new(0), + ) + .ptr(); + + let transition = UnscentedTransitionN::<4>::new(state, Box::new(move |s| f * s.value), q); + let measurement = UnscentedMeasurementNM::<4, 2>::new(Box::new(move |s| radar(&s.value)), r); + let mut filter = + UnscentedKalmanFilterNM::new(UnscentedUpdateNM::<4, 2>::new(measurement, transition)); + + for t in 1..=20 { + truth = f * truth; + // deterministic stand-in for sensor noise + let noise = + na::SVector::::new(0.1 * (1.7 * t as f64).sin(), 0.01 * (2.3 * t as f64).cos()); + let obs = Observation { + value: radar(&truth) + noise, + epoch: Epoch::new(t), + }; + + filter.predict(&obs.epoch); + let estimate = filter.update(&obs); + if t % 5 == 0 { + let x = estimate.read().unwrap().value; + println!( + "t={:2} position=({:.3}, {:.3}) velocity=({:.3}, {:.3}) position error={:.3}", + t, + x[0], + x[1], + x[2], + x[3], + ((x[0] - truth[0]).powi(2) + (x[1] - truth[1]).powi(2)).sqrt() + ); + } + } +} diff --git a/examples/state.py b/examples/state.py index 68113ae..c7b1cd2 100644 --- a/examples/state.py +++ b/examples/state.py @@ -9,6 +9,6 @@ ]) epoch = dt.Epoch(0) -state = dt.State(value, cov, epoch) +state = dt.State(value, covariance=cov, epoch=epoch) print(state) \ No newline at end of file diff --git a/examples/transition.py b/examples/transition.py index 70a8212..b8a23ae 100644 --- a/examples/transition.py +++ b/examples/transition.py @@ -15,7 +15,7 @@ def state(self, epoch: dt.Epoch): # state / covariance update value = self.f @ self._state.value cov = self.f @ self._state.covariance @ self.f.T + self.q - self._state = dt.State(value, cov, epoch) + self._state = dt.State(value, covariance=cov, epoch=epoch) return self._state @@ -26,13 +26,13 @@ def jacobian(self, state: dt.State): cov = np.eye(value.size) epoch = dt.Epoch(0) -state = dt.State(value, cov, epoch) +state = dt.State(value, covariance=cov, epoch=epoch) f = np.diag(np.array([1.5, 2.0])) q = np.eye(f.shape[0]) model = LinearTransition(f, q, state) -transition = dt.TransitionModel(model, state) +transition = dt.TransitionModel(model, state=state) if __name__ == "__main__": for n in range(2): diff --git a/examples/two_mass_spring_damper.rs b/examples/two_mass_spring_damper.rs index 9389b3f..641a89d 100644 --- a/examples/two_mass_spring_damper.rs +++ b/examples/two_mass_spring_damper.rs @@ -1,55 +1,66 @@ +use rand_distr::Distribution; /// Example: Simulate the two-mass-spring-damper Kalman filter. -/// +/// /// This example demonstrates the Kalman filter applied to the two-mass-spring-damper /// system from: https://solmaz.eng.uci.edu/Teaching/MAE195_code/Kalman_filter_Two_Mass_spring_damper.html -/// +/// /// Run with: cargo run --example two_mass_spring_damper --release - use statrs::distribution::MultivariateNormal; -use rand; -use rand_distr::Distribution; -use determined::algorithms::{ KalmanFilterNM }; -use determined::models::{ LinearTransition, LinearMeasurement, LinearUpdate }; +use determined::algorithms::KalmanFilterNM; +use determined::common::na; +use determined::epoch::Epoch; use determined::filter::Filter; -use determined::state::State; use determined::measurement::Observation; -use determined::epoch::Epoch; -use determined::common::na; +use determined::models::{LinearMeasurementNM, LinearTransitionN, LinearUpdateNM}; +use determined::state::State; +// returns the filter with the true-system and measurement closures; `impl Fn` +// types cannot be named in a type alias +#[allow(clippy::type_complexity)] fn setup_two_mass_spring_damper_filter() -> ( KalmanFilterNM<4, 2>, impl Fn(&State>) -> State>, - impl Fn(&State>) -> Observation>) { + impl Fn(&State>) -> Observation>, +) { // Physical parameters from UCI example - let m = 0.1; // mass 2 (kg) - let m_large = 1.0; // mass 1 (kg) - let k = 0.091; // spring constant (N/m) - let b = 0.0036; // damping coefficient (N·s/m) - + let m = 0.1; // mass 2 (kg) + let m_large = 1.0; // mass 1 (kg) + let k = 0.091; // spring constant (N/m) + let b = 0.0036; // damping coefficient (N·s/m) + // Continuous-time state transition matrix A let a = na::SMatrix::::new( - 0.0, 1.0, 0.0, 0.0, - -k/m, -b/m, k/m, b/m, - 0.0, 0.0, 0.0, 1.0, - k/m_large, b/m_large, -k/m_large, -b/m_large, + 0.0, + 1.0, + 0.0, + 0.0, + -k / m, + -b / m, + k / m, + b / m, + 0.0, + 0.0, + 0.0, + 1.0, + k / m_large, + b / m_large, + -k / m_large, + -b / m_large, ); // Discretization parameters let delta_t = 0.1; let identity = na::SMatrix::::identity(); - + // Discrete state transition matrix: F = I + Δt·A let f = identity + delta_t * a; - let f_model = f.clone(); // for model closure + let f_model = f; // for model closure // Measurement matrix H: observe only positions - let h = na::SMatrix::::new( - 1.0, 0.0, 0.0, 0.0, - 0.0, 0.0, 1.0, 0.0, - ); - let h_model = h.clone(); // for measurement closure - + let h = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0); + let h_model = h; // for measurement closure + // Process noise covariance Q let sigma_q = na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[ 0.3 * delta_t, @@ -57,49 +68,36 @@ fn setup_two_mass_spring_damper_filter() -> ( 0.3 * delta_t, 0.05 * delta_t, ])); - let q = &sigma_q * sigma_q.transpose(); - + let q = sigma_q * sigma_q.transpose(); + // Measurement noise covariance R - let sigma_r = na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[0.1, 0.1])); - let r = &sigma_r * sigma_r.transpose(); + let sigma_r = + na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[0.1, 0.1])); + let r = sigma_r * sigma_r.transpose(); let x_hat_init = na::SVector::::new(-1.0, 0.0, -0.5, 0.0); - let x = State{ - value: x_hat_init.into(), - covariance: Some(identity * 2.0), - epoch: Epoch::new(0) - }.ptr(); + let x = State { + value: x_hat_init, + covariance: Some(identity * 2.0), + epoch: Epoch::new(0), + } + .ptr(); - let transition = LinearTransition::<4>::new( - x.clone(), - f, - q - ); + let transition = LinearTransitionN::<4>::new(x.clone(), f, q); - let measurement = LinearMeasurement::<4, 2>::new( - h, - r - ); + let measurement = LinearMeasurementNM::<4, 2>::new(h, r); // build algorithm update model - let update = LinearUpdate::<4, 2>::new( - x.clone(), - measurement, - transition - ); - + let update = LinearUpdateNM::<4, 2>::new(measurement, transition); + // build kalman filter with give update model - let kf = - KalmanFilterNM::new( - x, - update, - ); + let kf = KalmanFilterNM::new(update); // function closure for system dynamics - let model = move | state: &State> | -> State> { + let model = move |state: &State>| -> State> { let new_value = f_model * state.value; - State{ + State { value: new_value, covariance: None, epoch: Epoch::new(state.epoch.value + 1), @@ -107,15 +105,15 @@ fn setup_two_mass_spring_damper_filter() -> ( }; // function closure for measurement model - let h = move | state: &State> | -> Observation> { - let measurement_value = h_model * &state.value; - Observation{ + let h = move |state: &State>| -> Observation> { + let measurement_value = h_model * state.value; + Observation { value: measurement_value, epoch: Epoch::new(state.epoch.value), } }; - ( kf, model, h) + (kf, model, h) } fn main() { @@ -127,19 +125,19 @@ fn main() { epoch: Epoch { value: 0 }, }; - let num_steps = 100 as i64; + let num_steps = 100_i64; let epoch = Epoch { value: 0 }; let noise_std_dev = 0.05; - + let initial_cov_trace = kf.state().read().unwrap().covariance().trace(); let mut max_error_x1: f64 = 0.0; let mut max_error_x2: f64 = 0.0; let mean = na::SMatrix::::from_column_slice(&[0.0, 0.0]); - let cov = na::SMatrix::::from_diagonal( - &na::SVector::from_row_slice( - &[noise_std_dev * noise_std_dev, noise_std_dev * noise_std_dev] - )); + let cov = na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[ + noise_std_dev * noise_std_dev, + noise_std_dev * noise_std_dev, + ])); let mvn = MultivariateNormal::new_from_nalgebra(mean, cov).unwrap(); let rng = &mut rand::thread_rng(); @@ -148,32 +146,36 @@ fn main() { for step in 0..num_steps { // Predict kf.predict(&epoch); - + // Propagate true system x_true = model(&x_true); - + // Simulate measurement with noise let noise_matrix = sampler.next().unwrap(); let measurement_value = h(&x_true).value + noise_matrix; - let observation = Observation{ + let observation = Observation { value: measurement_value, - epoch: Epoch::new(step as i64), + epoch: Epoch::new(step), }; // Update filter kf.update(&observation); - + // Track errors let x_est = kf.state().read().unwrap().value; max_error_x1 = max_error_x1.max((x_est[0] - x_true.value[0]).abs()); max_error_x2 = max_error_x2.max((x_est[2] - x_true.value[2]).abs()); } - + let final_cov_trace = kf.state().read().unwrap().covariance().trace(); let uncertainty_reduction = (1.0 - final_cov_trace / initial_cov_trace) * 100.0; - + println!("=== Two-Mass-Spring-Damper Kalman Filter ==="); - println!("Simulation: {} steps ({:.1} seconds)", num_steps, (num_steps - 1) as f64 * 0.1); + println!( + "Simulation: {} steps ({:.1} seconds)", + num_steps, + (num_steps - 1) as f64 * 0.1 + ); println!("\nResults:"); println!(" Max error (x1): {:.6} m", max_error_x1); println!(" Max error (x2): {:.6} m", max_error_x2); diff --git a/examples/ud.py b/examples/ud.py new file mode 100644 index 0000000..6cf46ac --- /dev/null +++ b/examples/ud.py @@ -0,0 +1,55 @@ +"""UD Kalman filter tracking a target moving at constant velocity from noisy, +correlated position measurements, alongside the standard Kalman filter. + +The UD filter holds the covariance as factors P = U D U^T and propagates and +updates the factors directly, so P stays symmetric positive semi-definite by +construction, which matters for long runs and badly scaled problems. It gives +the same estimates as `KalmanFilter`, as this example checks. +""" +import numpy as np + +import determined as dt + +# state x = (px, py, vx, vy), one epoch per second, position measured +F = np.array([ + [1.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, 1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], +]) +Q = np.diag([0.01, 0.01, 0.001, 0.001]) +H = np.array([ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], +]) +R = np.array([ + [0.04, 0.03], # correlated noise: decorrelated internally + [0.03, 0.09], +]) + + +def initial_state(): + return dt.State(np.zeros(4), covariance=np.diag([10.0, 10.0, 1.0, 1.0]), epoch=dt.Epoch(0)) + + +ud = dt.UDKalmanFilter(initial_state(), f=F, q=Q, h=H, r=R) +kf = dt.KalmanFilter(initial_state(), f=F, q=Q, h=H, r=R) + +truth = np.array([0.0, 0.0, 1.0, 0.5]) +rng = np.random.default_rng(0) +noise = np.linalg.cholesky(R) + +for t in range(1, 31): + truth = F @ truth + z = H @ truth + noise @ rng.normal(size=2) + obs = dt.Observation(z, epoch=dt.Epoch(t)) + + a, b = ud.update(obs), kf.update(obs) + + if t % 10 == 0: + print(f"t={t:2d} estimate={a.value.round(3)} truth={truth.round(3)}") + print(f" |UD - KF| state={np.abs(a.value - b.value).max():.1e} " + f"covariance={np.abs(a.covariance - b.covariance).max():.1e}") + +p = ud.state.covariance +print(f"P symmetric: {np.array_equal(p, p.T)}, smallest eigenvalue: {np.linalg.eigvalsh(p).min():.2e}") diff --git a/examples/ukf.py b/examples/ukf.py new file mode 100644 index 0000000..ad46707 --- /dev/null +++ b/examples/ukf.py @@ -0,0 +1,49 @@ +"""Unscented Kalman filter tracking a target moving at constant velocity with a +range-and-bearing sensor (a radar). + +No Jacobians are needed: the filter pushes sigma points through `f` and `h` +instead of linearising them. Both are evaluated at 2N + 1 sigma points per +step, so they must not keep state between calls. +""" +import numpy as np + +import determined as dt + +# state x = (px, py, vx, vy), one epoch per second +F = np.array([ + [1.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, 1.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], +]) +Q = np.diag([0.01, 0.01, 0.001, 0.001]) +R = np.diag([0.1 ** 2, 0.01 ** 2]) # range (m) and bearing (rad) noise + + +def f(x, t): + return F @ x # constant velocity + + +def h(x, t): + return np.array([np.hypot(x[0], x[1]), np.arctan2(x[1], x[0])]) + + +# the target starts at (10, 5) moving at (-0.5, 0.3); the initial guess is poor +truth = np.array([10.0, 5.0, -0.5, 0.3]) +state = dt.State(np.array([9.0, 6.0, 0.0, 0.0]), covariance=np.diag([4.0, 4.0, 1.0, 1.0]), + epoch=dt.Epoch(0)) + +ukf = dt.UnscentedKalmanFilter(state, f=f, q=Q, h=h, r=R) + +rng = np.random.default_rng(0) +for t in range(1, 21): + truth = F @ truth + z = h(truth, t) + rng.normal(0.0, np.sqrt(np.diag(R))) + + ukf.predict(dt.Epoch(t)) + estimate = ukf.update(dt.Observation(z, epoch=dt.Epoch(t))) + + if t % 5 == 0: + error = np.linalg.norm(estimate.value[:2] - truth[:2]) + print(f"t={t:2d} position={estimate.value[:2].round(3)} velocity={estimate.value[2:].round(3)} " + f"position error={error:.3f}") diff --git a/examples/update.py b/examples/update.py index a4596eb..88bf962 100644 --- a/examples/update.py +++ b/examples/update.py @@ -24,7 +24,7 @@ def apply(self, observation: dt.Observation): # compute gains h = self.m.model.h h_t = h.transpose() - r = self.m.model.r + r = self.m.noise() s = h @ _cov @ h_t + r s_inv = np.linalg.pinv(s) @@ -35,29 +35,23 @@ def apply(self, observation: dt.Observation): value = _state.value + k @ y cov = (np.eye(value.size) - k @ h) @ _cov - return dt.State(value, cov, epoch) - - def jacobian(self, state: dt.State): - jac = self.t.jacobian(state) - return jac + return dt.State(value, covariance=cov, epoch=epoch) value = np.array([1.0, 0.0]) epoch = dt.Epoch(0) -state = dt.State(value, np.eye(value.size), epoch) -obs = dt.Observation(np.array([0.0, 1.0]), epoch) +state = dt.State(value, covariance=np.eye(value.size), epoch=epoch) +obs = dt.Observation(np.array([0.0, 1.0]), epoch=epoch) model = LinearUpdate(transition, measurement) -update = dt.UpdateModel(model, transition, measurement) +update = dt.UpdateModel(model, transition=transition, measurement=measurement) if __name__ == "__main__": future_state = update.state(epoch) new_state = update.apply(obs) - jac = update.jacobian(state) print(future_state) - print(new_state) - print(jac) \ No newline at end of file + print(new_state) \ No newline at end of file diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000..b0b97ab --- /dev/null +++ b/flake.lock @@ -0,0 +1,61 @@ +{ + "nodes": { + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1791914995, + "narHash": "sha256-Zf/xfQEAasnwnmEfHqW/4hjDE/s44yn2hD4RuVIZRjA=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "b7a79d3fea1e214f6d668d521cecec0c01e26544", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1792502884, + "narHash": "sha256-IoNisKcGbUjdd45qCqywGMCYmZ5EnkxezvLII0FWwNA=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "02c087f3fe673c9e317f206a627cb6b910336dbe", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..9b60bb7 --- /dev/null +++ b/flake.nix @@ -0,0 +1,46 @@ +{ + description = "determined: Rust filtering library with Python bindings"; + + inputs = { + nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; + flake-utils.url = "github:numtide/flake-utils"; + }; + + outputs = { self, nixpkgs, flake-utils }: + flake-utils.lib.eachDefaultSystem (system: + let + pkgs = import nixpkgs { inherit system; }; + + pythonPackages = ps: [ ps.numpy ps.pytest ]; + python = pkgs.python3.withPackages (ps: pythonPackages ps ++ [ + # Manim Community, for examples/manim_two_mass_spring_damper.py. The + # package brings its native dependencies (cairo, pango, ffmpeg...). + # Its own test suite is skipped: there is no binary cache build on + # darwin, and one upstream CLI-text test fails when building it. + (ps.manim.overridePythonAttrs (_: { doCheck = false; })) + ]); + + shell = python: pkgs.mkShell { + packages = with pkgs; [ + # Rust + cargo + rustc + clippy + rustfmt + rust-analyzer + + # Python bindings + python + maturin + ]; + + env.RUST_SRC_PATH = "${pkgs.rustPlatform.rustLibSrc}"; + }; + in { + devShells.default = shell python; + + # CI: the same without manim, which would otherwise be built from + # source on every run (the override above changes its hash) + devShells.ci = shell (pkgs.python3.withPackages pythonPackages); + }); +} diff --git a/src/algorithms.rs b/src/algorithms.rs index 8890640..719d9ea 100644 --- a/src/algorithms.rs +++ b/src/algorithms.rs @@ -1,65 +1,88 @@ //! Algorithms module //! -//! Contains concrete filter algorithm implementations. The implementations are -//! intentionally written using nalgebra's statically sized `SMatrix` types so -//! they can be used in `no_std`/embedded contexts with predictable memory -//! usage. +//! Contains concrete filter algorithm implementations. They are generic over +//! the dimension type: use the `...NM` aliases for nalgebra's statically sized +//! matrices (predictable, allocation-free memory usage for `no_std`/embedded +//! contexts), or the `...Dyn` aliases for sizes only known at runtime (as used +//! by the Python bindings). -use crate::common::na; -use crate::state::StatePtr; -use crate::measurement::Observation; +use std::marker::PhantomData; + +use crate::common::{FilterAllocator, na}; use crate::epoch::Epoch; use crate::filter::Filter; +use crate::measurement::Observation; use crate::models::{ - UpdateModel, LinearUpdate, - DefaultFromState + DefaultFromState, LinearUDUpdate, LinearUpdate, NonLinearUDUpdate, NonLinearUpdate, + UnscentedUpdate, UpdateModel, }; +use crate::state::StatePtr; - -/// Generic Kalman Filter implementation for statically sized -/// state transition and measurement models. -pub struct KalmanFilterT +/// Generic Kalman Filter implementation for a state of dimension `N` and +/// measurements of dimension `M`, built from an update model `T`. +pub struct KalmanFilterT where - T: UpdateModel, na::Const> + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, + T: UpdateModel, { - pub state: StatePtr>, + pub state: StatePtr, pub update: T, + _measurement: PhantomData M>, } -/// Implementation for constructors -impl KalmanFilterT +/// Impementation for Defaults +impl DefaultFromState for KalmanFilterT where - T: UpdateModel, na::Const> - + DefaultFromState>, DefaultType = T>, + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, + T: UpdateModel + DefaultFromState, DefaultType = T>, { - pub fn new( - state: StatePtr>, - update: T, - ) -> Self { + type DefaultType = Self; + type StateType = StatePtr; + + fn default_from_state(state: Self::StateType) -> Self::DefaultType { KalmanFilterT { - state, - update, + state: state.clone(), + update: T::default_from_state(state), + _measurement: PhantomData, } } +} - pub fn default_from_state(state: StatePtr>) -> Self { +/// Implementation for constructors +impl KalmanFilterT +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, + T: UpdateModel, +{ + pub fn new(mut update: T) -> Self { + let state = update.state(None); KalmanFilterT { state: state.clone(), - update: T::default_from_state(state), + update, + _measurement: PhantomData, } } } /// Filter trait implementation -impl Filter for KalmanFilterT +impl Filter for KalmanFilterT where - T: UpdateModel, na::Const>, + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, + T: UpdateModel, { - type StateType = StatePtr>; - type ObservationType = Observation>; + type StateType = StatePtr; + type ObservationType = Observation; fn predict(&mut self, _epoch: &Epoch) -> Self::StateType { - self.update.state(_epoch).clone() + self.update.state(Some(_epoch)).clone() } fn update(&mut self, observation: &Self::ObservationType) -> Self::StateType { @@ -72,14 +95,46 @@ where fn reset(&mut self) { let mut state = self.state.write().unwrap(); + let (n, _) = state.value.shape_generic(); state.value.fill(0.0); - state.covariance = Some(na::SMatrix::::identity()); + state.covariance = Some(na::OMatrix::::identity_generic(n, n)); } } /// KalmanFilter pub type KalmanFilterNM = - KalmanFilterT< - N, M, - LinearUpdate - >; \ No newline at end of file + KalmanFilterT, na::Const, LinearUpdate, na::Const>>; + +/// KalmanFilter with runtime sized state and measurement +pub type KalmanFilterDyn = KalmanFilterT>; + +/// Extended Kalman Filter +pub type ExtendedKalmanFilterNM = + KalmanFilterT, na::Const, NonLinearUpdate, na::Const>>; + +/// Extended Kalman Filter with runtime sized state and measurement +pub type ExtendedKalmanFilterDyn = + KalmanFilterT>; + +/// Unscented Kalman Filter +pub type UnscentedKalmanFilterNM = + KalmanFilterT, na::Const, UnscentedUpdate, na::Const>>; + +/// Unscented Kalman Filter with runtime sized state and measurement +pub type UnscentedKalmanFilterDyn = + KalmanFilterT>; + +/// UD factorised Kalman Filter +pub type UDKalmanFilterNM = + KalmanFilterT, na::Const, LinearUDUpdate, na::Const>>; + +/// UD factorised Kalman Filter with runtime sized state and measurement +pub type UDKalmanFilterDyn = KalmanFilterT>; + +/// UD factorised Extended Kalman Filter +pub type UDExtendedKalmanFilterNM = + KalmanFilterT, na::Const, NonLinearUDUpdate, na::Const>>; + +/// UD factorised Extended Kalman Filter with runtime sized state and measurement +pub type UDExtendedKalmanFilterDyn = + KalmanFilterT>; diff --git a/src/common/composition.rs b/src/common/composition.rs index 7832cc4..7cb1b57 100644 --- a/src/common/composition.rs +++ b/src/common/composition.rs @@ -1,4 +1,4 @@ -pub trait Composable { - type Output; - fn add(self, other: Rhs) -> Self::Output; -} \ No newline at end of file +pub trait Composable { + type Output; + fn add(self, other: Rhs) -> Self::Output; +} diff --git a/src/common/linalg.rs b/src/common/linalg.rs index f6fb5d0..771410e 100644 --- a/src/common/linalg.rs +++ b/src/common/linalg.rs @@ -1,14 +1,66 @@ +use na::allocator::Allocator; use nalgebra as na; -pub fn moore_penrose_right_inverse( - a: &na::SMatrix, -) -> Option> +/// Moore-Penrose pseudo-inverse of a full-rank matrix. +/// +/// Uses the left inverse `(A^T A)^-1 A^T` for a tall or square `A` (full +/// column rank) and the right inverse `A^T (A A^T)^-1` for a wide `A` (full +/// row rank). The choice is made by shape: for a wide `A`, `A^T A` is singular +/// but can still be inverted numerically because of rounding, giving garbage. +/// Returns `None` if `A` is rank deficient. +pub fn moore_penrose_inverse(a: &na::OMatrix) -> Option> where - T: nalgebra::RealField, + T: na::RealField, + R: na::Dim, + C: na::Dim, + na::DefaultAllocator: Allocator + Allocator + Allocator + Allocator, { - // Moore-Penrose right inverse: A^+ = (A^T A)^-1 A^T let a_t = a.transpose(); - let ata = &a_t * a; - let ata_inv = ata.try_inverse()?; - Some(&ata_inv * &a_t) -} \ No newline at end of file + if a.nrows() >= a.ncols() { + let ata_inv = (&a_t * a).try_inverse()?; + Some(ata_inv * a_t) + } else { + let aat_inv = (a * &a_t).try_inverse()?; + Some(a_t * aat_inv) + } +} + +/// UD decomposition `M = U D U^T` of a symmetric positive semi-definite +/// matrix, with `U` unit upper triangular and `D` diagonal (returned as a +/// vector). Ramos et al., "A summary on the UD Kalman Filter" (arXiv:2203.06105), +/// Appendix A. +/// +/// A zero pivot (a semi-definite `M`) gives zero off-diagonal entries in that +/// column of `U`. A negative entry in `D` means `M` was not positive +/// semi-definite. +pub fn udu_decomposition( + m: &na::OMatrix, +) -> (na::OMatrix, na::OVector) +where + D: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + let (nd, _) = m.shape_generic(); + let n = m.nrows(); + let mut u = na::OMatrix::::identity_generic(nd, nd); + let mut d = na::OVector::::zeros_generic(nd, na::Const::<1>); + + // pivots below this are treated as zero + let scale = m.diagonal().iter().fold(0.0f64, |acc, x| acc.max(x.abs())); + let tiny = f64::EPSILON * scale.max(f64::MIN_POSITIVE); + + for j in (0..n).rev() { + for i in (0..=j).rev() { + let mut sigma = m[(i, j)]; + for k in j + 1..n { + sigma -= u[(i, k)] * d[k] * u[(j, k)]; + } + if i == j { + d[j] = sigma; + } else if d[j].abs() > tiny { + u[(i, j)] = sigma / d[j]; + } + } + } + (u, d) +} diff --git a/src/common/mod.rs b/src/common/mod.rs index 8c72a91..7e14fdc 100644 --- a/src/common/mod.rs +++ b/src/common/mod.rs @@ -1,6 +1,29 @@ pub use nalgebra as na; -pub mod linalg; pub mod composition; +pub mod linalg; + +pub use composition::Composable; + +use na::allocator::Allocator; + +/// nalgebra allocator bounds needed by the filter models, for a state of +/// dimension `N` and a measurement of dimension `M`. +/// +/// Lets the models be written once for both statically sized (`na::Const`) +/// and dynamically sized (`na::Dyn`) matrices. Use as +/// `where na::DefaultAllocator: FilterAllocator`. +pub trait FilterAllocator: + Allocator + Allocator + Allocator + Allocator + Allocator + Allocator +{ +} -pub use composition::Composable; \ No newline at end of file +impl FilterAllocator for A where + A: Allocator + + Allocator + + Allocator + + Allocator + + Allocator + + Allocator +{ +} diff --git a/src/epoch.rs b/src/epoch.rs index 4e22779..f9bdaa9 100644 --- a/src/epoch.rs +++ b/src/epoch.rs @@ -5,9 +5,8 @@ use std::ops::{Add, Sub}; -pub use nalgebra as na; pub use crate::common::Composable; - +pub use nalgebra as na; #[derive(Default, Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] pub struct Epoch { @@ -17,7 +16,7 @@ pub struct Epoch { impl Epoch { pub fn new(value: i64) -> Self { Epoch { value } - } + } } impl Add for Epoch { @@ -58,4 +57,4 @@ impl Sub for &Epoch { value: self.value - other.value, } } -} \ No newline at end of file +} diff --git a/src/filter.rs b/src/filter.rs index f70fdd7..584b942 100644 --- a/src/filter.rs +++ b/src/filter.rs @@ -1,24 +1,24 @@ -//! `Filter` is a trait for basic algorithms. - -use crate::epoch::Epoch; - -/// Core Filter trait shared across algorithm implementations. -/// -/// Implementations expose their concrete `StateType` and `ObservationType` as -/// associated types. The trait intentionally returns references from -/// `predict`/`update` to avoid unnecessary cloning of large state vectors. -/// -/// Minimal contract: -/// - `predict(epoch)` advances internal state (e.g. x, P) and returns a -/// reference to the predicted state. -/// - `update(obs)` fuses an observation into the internal state and returns -/// a reference to the posterior state. -pub trait Filter { - type StateType; - type ObservationType; - - fn predict(&mut self, epoch: &Epoch) -> Self::StateType; - fn update(&mut self, observation: &Self::ObservationType) -> Self::StateType; - fn state(&self) -> Self::StateType; - fn reset(&mut self); -} +//! `Filter` is a trait for basic algorithms. + +use crate::epoch::Epoch; + +/// Core Filter trait shared across algorithm implementations. +/// +/// Implementations expose their concrete `StateType` and `ObservationType` as +/// associated types. The trait intentionally returns references from +/// `predict`/`update` to avoid unnecessary cloning of large state vectors. +/// +/// Minimal contract: +/// - `predict(epoch)` advances internal state (e.g. x, P) and returns a +/// reference to the predicted state. +/// - `update(obs)` fuses an observation into the internal state and returns +/// a reference to the posterior state. +pub trait Filter { + type StateType; + type ObservationType; + + fn predict(&mut self, epoch: &Epoch) -> Self::StateType; + fn update(&mut self, observation: &Self::ObservationType) -> Self::StateType; + fn state(&self) -> Self::StateType; + fn reset(&mut self); +} diff --git a/src/lib.rs b/src/lib.rs index 57110f3..7380b34 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,26 +1,28 @@ -//! Determined — small filtering primitives -//! -//! This crate contains compact implementations of static-size Kalman and -//! Extended Kalman filters using nalgebra's `SMatrix`/`SVector`. It's designed -//! for clarity and small embedded-friendly usage. -//! -//! Run `cargo doc --open` to generate API documentation locally. -//! -//! Main modules: -//! - `common` — core traits and aliases (e.g. `Algorithm`). -//! - `algorithms` — Kalman and EKF implementations. -//! - `state` — `State` wrapper around nalgebra matrices. -//! - `measurement` — measurement/observation helpers. -//! - `models` — traits for `StateTransition` and `MeasurementModel` used by EKF. - -pub mod common; -pub mod algorithms; -pub mod measurement; -pub mod state; -pub mod epoch; -pub mod models; -pub mod filter; -pub mod python; - -#[cfg(test)] -mod tests; \ No newline at end of file +//! Determined — small filtering primitives +//! +//! This crate contains compact implementations of Kalman and Extended Kalman +//! filters on nalgebra matrices. The models are generic over the dimension +//! type: statically sized (`na::Const`, via the `...N`/`...NM` aliases such +//! as `KalmanFilterNM`) for allocation-free embedded use, or runtime sized +//! (`na::Dyn`, e.g. `KalmanFilterDyn`), which the Python bindings wrap. +//! +//! Run `cargo doc --open` to generate API documentation locally. +//! +//! Main modules: +//! - `common` — core traits and aliases (e.g. `Algorithm`). +//! - `algorithms` — Kalman and EKF implementations. +//! - `state` — `State` wrapper around nalgebra matrices. +//! - `measurement` — measurement/observation helpers. +//! - `models` — traits for `StateTransition` and `MeasurementModel` used by EKF. + +pub mod algorithms; +pub mod common; +pub mod epoch; +pub mod filter; +pub mod measurement; +pub mod models; +pub mod python; +pub mod state; + +#[cfg(test)] +mod tests; diff --git a/src/measurement.rs b/src/measurement.rs index 28ade71..ea0d23a 100644 --- a/src/measurement.rs +++ b/src/measurement.rs @@ -1,36 +1,36 @@ -//! Measurement helpers and observation types. -//! -//! `Measurement` is a small container that groups a state, observation and a -//! joint covariance / jacobian. It provides convenience constructors and a -//! minimal update/predict placeholder useful for tests and examples. - -use std::fmt; - -use crate::common::na; -use crate::epoch::Epoch; - -#[derive(Clone, Debug)] -pub struct Observation -where - T: na::Dim, - na::DefaultAllocator: na::base::allocator::Allocator - + na::allocator::Allocator, -{ - pub value: na::Matrix>, - pub epoch: Epoch, -} - -impl fmt::Display for Observation -where - T: na::Dim, - na::DefaultAllocator: na::base::allocator::Allocator - + na::allocator::Allocator, -{ - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Format `value`, `covariance`, and `epoch` as desired - writeln!(f, "Observation:")?; - writeln!(f, "‾‾‾‾‾‾‾‾‾‾‾‾")?; - write!(f, " Value: {:0.6}", self.value)?; - writeln!(f, " Epoch: {}", self.epoch.value) - } -} \ No newline at end of file +//! Measurement helpers and observation types. +//! +//! `Measurement` is a small container that groups a state, observation and a +//! joint covariance / jacobian. It provides convenience constructors and a +//! minimal update/predict placeholder useful for tests and examples. + +use std::fmt; + +use crate::common::na; +use crate::epoch::Epoch; + +#[derive(Clone, Debug)] +pub struct Observation +where + T: na::Dim, + na::DefaultAllocator: + na::base::allocator::Allocator + na::allocator::Allocator, +{ + pub value: na::Matrix>, + pub epoch: Epoch, +} + +impl fmt::Display for Observation +where + T: na::Dim, + na::DefaultAllocator: + na::base::allocator::Allocator + na::allocator::Allocator, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Format `value`, `covariance`, and `epoch` as desired + writeln!(f, "Observation:")?; + writeln!(f, "‾‾‾‾‾‾‾‾‾‾‾‾")?; + write!(f, " Value: {:0.6}", self.value)?; + writeln!(f, " Epoch: {}", self.epoch.value) + } +} diff --git a/src/models/linear.rs b/src/models/linear.rs index 5c82229..e47d8e6 100644 --- a/src/models/linear.rs +++ b/src/models/linear.rs @@ -1,71 +1,84 @@ -//! Model traits used by the Extended Kalman Filter implementation. +//! Linear models used by the Kalman Filter implementation. //! -//! These traits allow users to supply nonlinear state transition and -//! measurement models with accompanying Jacobians so the EKF can operate. -//! +//! These models take linear state transition and measurement matrices +//! defining a standard Kalman Filter. They are generic over the dimension +//! type, so work with statically sized (`na::Const`, see the `...N` / +//! `...NM` aliases) and dynamically sized (`na::Dyn`) matrices. -use crate::common::na as na; +use crate::common::linalg::moore_penrose_inverse; +use crate::common::{FilterAllocator, na}; use crate::epoch::Epoch; use crate::measurement::Observation; -use crate::state::{ State, StatePtr }; +use crate::state::{State, StatePtr}; -use crate::common::linalg::{ moore_penrose_right_inverse }; +use na::allocator::Allocator; -use crate::models::traits::{ - TransitionModel, MeasurementModel, UpdateModel, DefaultFromState +use crate::models::traits::{ + DefaultFromState, InverseMeasurement, MeasurementJacobian, MeasurementModel, + TransitionJacobian, TransitionModel, UpdateModel, }; +use crate::models::update::kalman_update; // Transition Models // Linear (Discrete Time) #[derive(Clone, Debug)] -pub struct LinearTransition { - pub state: StatePtr>, // state vector (N x 1) - pub f: na::SMatrix, // state transition (N x N) - pub q: na::SMatrix, // process noise covariance (N x N) - f_inv: Option>, // inverse of F +pub struct LinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + pub state: StatePtr, // state vector (N x 1) + pub f: na::OMatrix, // state transition (N x N) + pub q: na::OMatrix, // process noise covariance (N x N) + f_inv: Option>, // inverse of F } -impl LinearTransition { - pub fn new( - state: StatePtr>, - f: na::SMatrix, - q: na::SMatrix ) -> Self { - - let f_inv = match f.try_inverse() { - None => { - spdlog::warn!("State transition matrix F is non-invertible; - backward propagation will do nothing."); - None - }, - Some(f_inv) => Some(f_inv), - }; +/// Statically sized `LinearTransition`. +pub type LinearTransitionN = LinearTransition>; + +impl LinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + pub fn new(state: StatePtr, f: na::OMatrix, q: na::OMatrix) -> Self { + let f_inv = f.clone().try_inverse(); + if f_inv.is_none() { + spdlog::warn!( + "State transition matrix F is non-invertible; backward propagation will do nothing." + ); + } LinearTransition { state, f, q, f_inv } } +} - fn f_inv(&self) -> na::SMatrix { - // getter for F inverse with default zero matrix - self.f_inv.unwrap_or(na::SMatrix::::zeros()) - } -} - -impl DefaultFromState for LinearTransition { - +impl DefaultFromState for LinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ type DefaultType = LinearTransition; - type StateType = StatePtr>; + type StateType = StatePtr; fn default_from_state(state: Self::StateType) -> Self::DefaultType - where Self: Sized { + where + Self: Sized, + { + let n = N::from_usize(state.read().unwrap().dim()); LinearTransition::new( state, - na::SMatrix::::identity(), - na::SMatrix::::zeros(), + na::OMatrix::::identity_generic(n, n), + na::OMatrix::::zeros_generic(n, n), ) } } -impl TransitionModel> for LinearTransition { - - fn state(&mut self, epoch: &Epoch) -> &StatePtr> { +impl TransitionModel for LinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + fn state(&mut self, epoch: &Epoch) -> &StatePtr { // Determine timesteps to propagate let mut state = self.state.write().unwrap(); @@ -74,177 +87,207 @@ impl TransitionModel> for LinearTransition { return &self.state; } - let cov = &state.covariance(); - if delta > 0 { // propagate forward for _ in 0..delta { - state.value = self.f * state.value; - state.covariance = Some(self.f * cov * self.f.transpose() + self.q); + let cov = state.covariance(); + state.value = &self.f * &state.value; + state.covariance = Some(&self.f * cov * self.f.transpose() + &self.q); } } else { // propagate backward - let f_inv = self.f_inv(); + let Some(f_inv) = &self.f_inv else { + spdlog::warn!( + "State transition matrix F is non-invertible; backward propagation will do nothing." + ); + return &self.state; + }; for _ in 0..delta.abs() { - state.value = f_inv * state.value; - state.covariance = Some(f_inv * (cov - self.q) * f_inv.transpose()); + let cov = state.covariance(); + state.value = f_inv * &state.value; + state.covariance = Some(f_inv * (cov - &self.q) * f_inv.transpose()); } }; + state.epoch = *epoch; &self.state } +} - fn jacobian(&self, _state: &State>) -> na::SMatrix { - self.f +impl TransitionJacobian for LinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + /// The Jacobian of a linear transition is F. + fn jacobian(&self, _state: &State) -> na::OMatrix { + self.f.clone() } } /// Measurement Models /// Linear #[derive(Clone, Debug)] -pub struct LinearMeasurement { - pub h: na::SMatrix, // measurement matrix - pub r: na::SMatrix, // measurement noise covariance(M x M) - h_t: na::SMatrix, // transpose of H - h_inv: na::SMatrix, // inverse of H +pub struct LinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub h: na::OMatrix, // measurement matrix + pub r: na::OMatrix, // measurement noise covariance(M x M) + h_inv: Option>, // pseudo-inverse of H, if full rank } -impl LinearMeasurement { - pub fn new(h: na::SMatrix, r: na::SMatrix) -> Self { - let h_inv = match moore_penrose_right_inverse(&h) { - None => { - spdlog::warn!("Measurement matrix H is non-invertible; inverse projection will return zero."); - na::SMatrix::::zeros() - }, - Some(h_inv) => h_inv, - }; - LinearMeasurement { - h, - r, - h_t: h.transpose(), - h_inv, - } +/// Statically sized `LinearMeasurement`. +pub type LinearMeasurementNM = + LinearMeasurement, na::Const>; + +impl LinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new(h: na::OMatrix, r: na::OMatrix) -> Self { + let h_inv = moore_penrose_inverse(&h); + LinearMeasurement { h, r, h_inv } } } - -impl Default for LinearMeasurement { +impl Default for LinearMeasurement +where + N: na::DimName, + M: na::DimName, + na::DefaultAllocator: FilterAllocator, +{ fn default() -> Self { - let _h = na::SMatrix::::identity(); - LinearMeasurement { - h: _h, - r: na::SMatrix::::identity(), - h_t: _h.transpose(), - h_inv: moore_penrose_right_inverse(&_h).unwrap_or(na::SMatrix::::zeros()), - } + LinearMeasurement::new( + na::OMatrix::::identity(), + na::OMatrix::::identity(), + ) } } -impl MeasurementModel, na::Const> for LinearMeasurement { - fn projection(&self, state: &State>) -> Observation> { +impl MeasurementModel for LinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn projection(&self, state: &State) -> Observation { Observation { - value: self.h * state.value, - epoch: state.epoch + value: &self.h * &state.value, + epoch: state.epoch, } } - fn inverse(&self, obs: &Observation>) -> State> { - State { - value: self.h_inv * obs.value, + fn noise(&self) -> na::OMatrix { + self.r.clone() + } +} + +impl MeasurementJacobian for LinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + /// The Jacobian of a linear measurement is H. + fn jacobian(&self, _state: &State) -> na::OMatrix { + self.h.clone() + } +} + +impl InverseMeasurement for LinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + /// Pseudo-inverse projection; `None` if H is rank deficient. + fn inverse(&self, obs: &Observation) -> Option> { + self.h_inv.as_ref().map(|h_inv| State { + value: h_inv * &obs.value, covariance: None, epoch: obs.epoch, - } - } // TODO: probably dont need this since its could often be ill-defined - - fn jacobian(&self, _state: &State>) -> na::SMatrix { - self.h // TODO: code proper jacobian + }) } } // Update Models // Linear -pub struct LinearUpdate { - pub state: StatePtr>, // mutable ref. to state - pub measurement: LinearMeasurement, // measurement model - pub transition: LinearTransition, // transition model - identity: na::SMatrix, // identity matrix +pub struct LinearUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub state: StatePtr, // mutable ref. to state + pub measurement: LinearMeasurement, // measurement model + pub transition: LinearTransition, // transition model } -impl LinearUpdate { - pub fn new( - state: StatePtr>, - measurement: LinearMeasurement, - transition: LinearTransition, - ) -> Self { +/// Statically sized `LinearUpdate`. +pub type LinearUpdateNM = LinearUpdate, na::Const>; + +impl LinearUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new(measurement: LinearMeasurement, transition: LinearTransition) -> Self { + let _state = transition.state.clone(); LinearUpdate { - state, + state: _state, measurement, transition, - identity: na::SMatrix::::identity(), } } } -impl DefaultFromState for LinearUpdate { - +impl DefaultFromState for LinearUpdate +where + N: na::DimName, + M: na::DimName, + na::DefaultAllocator: FilterAllocator, +{ type DefaultType = LinearUpdate; - type StateType = StatePtr>; + type StateType = StatePtr; fn default_from_state(state: Self::StateType) -> Self::DefaultType { LinearUpdate::new( - state.clone(), LinearMeasurement::::default(), - LinearTransition::::default_from_state(state) + LinearTransition::::default_from_state(state), ) } } -impl UpdateModel, na::Const> for LinearUpdate { - - fn state(&mut self, epoch: &Epoch) -> &StatePtr> { - self.transition.state(epoch) +impl UpdateModel for LinearUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr { + match epoch { + Some(epoch) => self.transition.state(epoch), + _ => &self.transition.state, + } } - - fn apply(&mut self, observation: &Observation>) -> &StatePtr> { + fn apply(&mut self, observation: &Observation) -> &StatePtr { // propagate state to observation epoch _ = self.transition.state(&observation.epoch); let mut state = self.state.write().unwrap(); // project state to measurement domain - let z_x = &self.measurement.projection(&state).value; - let z = &observation.value; - - let cov = &state.covariance(); - - // compute gain matrix - let h = &self.measurement.h; - let h_t = &self.measurement.h_t; - let s = h * cov * h_t + self.measurement.r; - let s_inv = match s.try_inverse(){ - None => { - spdlog::error!( - "Invalid update: Innovation covariance S is singular; state will not be updated." - ); - return &self.state; - } - Some(s_inv) => s_inv, - }; - - // Compute gain and innovation - let k = cov * h_t * s_inv; - let y = z - z_x; - - // Update state - state.value = state.value + k * y; + let z_x = self.measurement.projection(&state).value; - // Update Covariance - state.covariance = Some((self.identity - k * h) * cov); + let h = self.measurement.jacobian(&state); + let r = self.measurement.noise(); + kalman_update(&mut state, &observation.value, &z_x, &h, &r); &self.state - - } - - fn jacobian(&self, _x: &State>) -> na::SMatrix { - self.measurement.h } -} \ No newline at end of file +} diff --git a/src/models/mod.rs b/src/models/mod.rs index 14dfb6e..66ef707 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,6 +1,23 @@ pub mod linear; +pub mod nonlinear; pub mod traits; +pub mod ud; +pub mod unscented; +pub mod update; +pub use linear::{ + LinearMeasurement, LinearMeasurementNM, LinearTransition, LinearTransitionN, LinearUpdate, + LinearUpdateNM, +}; +pub use nonlinear::{ + NonLinearMeasurement, NonLinearMeasurementNM, NonLinearTransition, NonLinearTransitionN, + NonLinearUpdate, NonLinearUpdateNM, +}; pub use traits::DefaultFromState; -pub use traits::{ TransitionModel, MeasurementModel, UpdateModel }; -pub use linear::{ LinearTransition, LinearMeasurement, LinearUpdate }; \ No newline at end of file +pub use traits::{InverseMeasurement, MeasurementJacobian, TransitionJacobian}; +pub use traits::{MeasurementModel, TransitionModel, UpdateModel}; +pub use ud::{LinearUDUpdate, LinearUDUpdateNM, NonLinearUDUpdate, NonLinearUDUpdateNM, UDFactors}; +pub use unscented::{ + UnscentedMeasurement, UnscentedMeasurementNM, UnscentedParams, UnscentedTransition, + UnscentedTransitionN, UnscentedUpdate, UnscentedUpdateNM, +}; diff --git a/src/models/nonlinear.rs b/src/models/nonlinear.rs new file mode 100644 index 0000000..028e669 --- /dev/null +++ b/src/models/nonlinear.rs @@ -0,0 +1,275 @@ +//! Nonlinear models used by the Extended Kalman Filter implementation. +//! +//! These models are built from user supplied closures: a propagation function +//! `f` and measurement function `h`, each with an accompanying Jacobian, so the +//! EKF can linearise about the current state estimate. They are generic over +//! the dimension type, so work with statically sized (`na::Const`, see the +//! `...N` / `...NM` aliases) and dynamically sized (`na::Dyn`) matrices. +//! +//! The closures are run on a copy of the state, so no lock is held while user +//! code runs. + +use crate::common::{FilterAllocator, na}; +use crate::epoch::Epoch; +use crate::measurement::Observation; +use crate::state::{State, StatePtr}; + +use na::allocator::Allocator; + +use crate::models::traits::{ + InverseMeasurement, MeasurementJacobian, MeasurementModel, TransitionJacobian, TransitionModel, + UpdateModel, +}; +use crate::models::update::kalman_update; + +/// One-step propagation: x(t) -> x(t + 1) +pub type F = Box) -> na::OVector + Send + Sync>; +/// Jacobian of `F` with respect to the state (N x N) +pub type J = Box) -> na::OMatrix + Send + Sync>; + +/// Measurement function: h(x) -> z +pub type H = Box) -> na::OVector + Send + Sync>; +/// Jacobian of `H` with respect to the state (M x N) +pub type JH = Box) -> na::OMatrix + Send + Sync>; +/// Optional inverse measurement: h^-1(z) -> x +pub type HInv = Box) -> State + Send + Sync>; + +// Transition Models +// NonLinear (Discrete Time) +pub struct NonLinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + pub state: StatePtr, + pub f: F, + pub j: J, + pub q: na::OMatrix, +} + +/// Statically sized `NonLinearTransition`. +pub type NonLinearTransitionN = NonLinearTransition>; + +impl NonLinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + /// `f` advances the state by a single epoch; it is called once per epoch + /// step, and `j` is evaluated at the state *before* each step. + pub fn new(state: StatePtr, f: F, j: J, q: na::OMatrix) -> Self { + NonLinearTransition { state, f, j, q } + } +} + +impl TransitionModel for NonLinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + fn state(&mut self, epoch: &Epoch) -> &StatePtr { + let mut state = self.state.read().unwrap().clone(); + + let delta = { *epoch - state.epoch }.value; + if delta == 0 { + return &self.state; + } + if delta < 0 { + // a general nonlinear f has no inverse + spdlog::warn!( + "Backward propagation is not supported for nonlinear transitions; state is unchanged." + ); + return &self.state; + } + + for _ in 0..delta { + let jac = (self.j)(&state); + let value = (self.f)(&state); + let cov = &jac * state.covariance() * jac.transpose() + &self.q; + + state.value = value; + state.covariance = Some(cov); + state.epoch = state.epoch + Epoch::new(1); + } + + *self.state.write().unwrap() = state; + &self.state + } +} + +impl TransitionJacobian for NonLinearTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + fn jacobian(&self, _state: &State) -> na::OMatrix { + (self.j)(_state) + } +} + +// Measurement Models +// NonLinear +pub struct NonLinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub h: H, + pub j: JH, + pub r: na::OMatrix, + h_inv: Option>, +} + +/// Statically sized `NonLinearMeasurement`. +pub type NonLinearMeasurementNM = + NonLinearMeasurement, na::Const>; + +impl NonLinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new(h: H, j: JH, r: na::OMatrix) -> Self { + NonLinearMeasurement { + h, + j, + r, + h_inv: None, + } + } + + /// Supply an inverse measurement function h^-1(z) -> x. + /// Without one, `inverse` returns `None`. + pub fn with_inverse(mut self, h_inv: HInv) -> Self { + self.h_inv = Some(h_inv); + self + } + + /// Linear sensor `z = H x` expressed as a nonlinear model, so a linear + /// measurement can be paired with a nonlinear transition. + pub fn linear(h: na::OMatrix, r: na::OMatrix) -> Self + where + na::OMatrix: Send + Sync + 'static, + { + let jac = h.clone(); + Self::new( + Box::new(move |state| &h * &state.value), + Box::new(move |_| jac.clone()), + r, + ) + } +} + +impl MeasurementModel for NonLinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn projection(&self, state: &State) -> Observation { + Observation { + value: (self.h)(state), + epoch: state.epoch, + } + } + + fn noise(&self) -> na::OMatrix { + self.r.clone() + } +} + +impl MeasurementJacobian for NonLinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn jacobian(&self, state: &State) -> na::OMatrix { + (self.j)(state) + } +} + +impl InverseMeasurement for NonLinearMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + /// Uses the function supplied with `with_inverse`, if any. + fn inverse(&self, obs: &Observation) -> Option> { + self.h_inv.as_ref().map(|h_inv| h_inv(obs)) + } +} + +// Update +pub struct NonLinearUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub state: StatePtr, // mutable ref. to state + pub measurement: NonLinearMeasurement, // measurement model + pub transition: NonLinearTransition, // transition model +} + +/// Statically sized `NonLinearUpdate`. +pub type NonLinearUpdateNM = + NonLinearUpdate, na::Const>; + +impl NonLinearUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new( + measurement: NonLinearMeasurement, + transition: NonLinearTransition, + ) -> Self { + let _state = transition.state.clone(); + NonLinearUpdate { + state: _state, + measurement, + transition, + } + } +} + +impl UpdateModel for NonLinearUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr { + match epoch { + Some(epoch) => self.transition.state(epoch), + _ => &self.transition.state, + } + } + + fn apply(&mut self, observation: &Observation) -> &StatePtr { + // propagate state to observation epoch + _ = self.transition.state(&observation.epoch); + let mut state = self.state.read().unwrap().clone(); + + // project state to measurement domain, linearised about the current estimate + let z_x = self.measurement.projection(&state).value; + let h = self.measurement.jacobian(&state); + + if kalman_update( + &mut state, + &observation.value, + &z_x, + &h, + &self.measurement.noise(), + ) { + *self.state.write().unwrap() = state; + } + + &self.state + } +} diff --git a/src/models/traits.rs b/src/models/traits.rs index 21d55a1..eb1e586 100644 --- a/src/models/traits.rs +++ b/src/models/traits.rs @@ -1,16 +1,20 @@ -//! Model traits used by the Extended Kalman Filter implementation. +//! Model traits used by the filter implementations. +//! +//! The core traits hold only what every filter needs: propagating the state +//! (`TransitionModel`), projecting it into the measurement domain with the +//! measurement noise (`MeasurementModel`), and the filter's update step +//! (`UpdateModel`). Capabilities that only some filters or models have are +//! separate traits: Jacobians (`TransitionJacobian`, `MeasurementJacobian`, +//! needed by the Extended Kalman Filter) and the inverse measurement model +//! (`InverseMeasurement`). //! -//! These traits allow users to supply nonlinear state transition and -//! measurement models with accompanying Jacobians so the EKF can operate. -//! -use crate::common::na as na; +use crate::common::na; use crate::epoch::Epoch; use crate::measurement::Observation; -use crate::state::{ State, StatePtr }; - -pub trait DefaultFromState{ +use crate::state::{State, StatePtr}; +pub trait DefaultFromState { type DefaultType; type StateType; @@ -18,15 +22,21 @@ pub trait DefaultFromState{ } /// State transition model trait. -pub trait TransitionModel +pub trait TransitionModel where N: na::Dim, - na::DefaultAllocator: na::allocator::Allocator - + na::allocator::Allocator, + na::DefaultAllocator: na::allocator::Allocator + na::allocator::Allocator, { /// x(t) -> x(t') fn state(&mut self, epoch: &Epoch) -> &StatePtr; +} +/// Jacobian of a state transition. +pub trait TransitionJacobian +where + N: na::Dim, + na::DefaultAllocator: na::allocator::Allocator + na::allocator::Allocator, +{ /// Jacobian df/dx (N x N) fn jacobian(&self, state: &State) -> na::OMatrix; } @@ -40,18 +50,45 @@ where + na::allocator::Allocator + na::allocator::Allocator + na::allocator::Allocator - + na::allocator::Allocator + + na::allocator::Allocator, { /// measurement projection h(x) -> z fn projection(&self, state: &State) -> Observation; - /// inverse measurement model h^-1(z) -> x - fn inverse(&self, obs: &Observation) -> State; + /// measurement noise covariance R (M x M) + fn noise(&self) -> na::OMatrix; +} +/// Jacobian of a measurement model. +pub trait MeasurementJacobian +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: na::allocator::Allocator + + na::allocator::Allocator + + na::allocator::Allocator + + na::allocator::Allocator + + na::allocator::Allocator, +{ /// Jacobian dh/dx (M x N) fn jacobian(&self, state: &State) -> na::OMatrix; } +/// Inverse measurement model. +pub trait InverseMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: na::allocator::Allocator + + na::allocator::Allocator + + na::allocator::Allocator + + na::allocator::Allocator, +{ + /// inverse measurement model h^-1(z) -> x, or `None` if this model has + /// no inverse (e.g. a rank deficient measurement matrix) + fn inverse(&self, obs: &Observation) -> Option>; +} + /// Update model trait. pub trait UpdateModel where @@ -61,14 +98,11 @@ where + na::allocator::Allocator + na::allocator::Allocator + na::allocator::Allocator - + na::allocator::Allocator -{ - /// State transition model - fn state(&mut self, epoch: &Epoch) -> &StatePtr; + + na::allocator::Allocator, +{ + /// State accessor + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr; /// Compute updated gains fn apply(&mut self, observation: &Observation) -> &StatePtr; - - /// Jacobian dh/dx (M x N) - fn jacobian(&self, x: &State) -> na::OMatrix; } diff --git a/src/models/ud.rs b/src/models/ud.rs new file mode 100644 index 0000000..a0e51cc --- /dev/null +++ b/src/models/ud.rs @@ -0,0 +1,474 @@ +//! UD factorised Kalman filters. +//! +//! The covariance is held as factors `P = U D U^T` (`U` unit upper triangular, +//! `D` diagonal) and the filter propagates and updates the factors directly, +//! never forming `P`. This keeps the covariance symmetric and positive +//! semi-definite by construction. Following Ramos et al., "A summary on the +//! UD Kalman Filter" (arXiv:2203.06105): +//! +//! - initialise: `[U, D] = udu(P)` +//! - time update: weighted modified Gram-Schmidt on `W = [F U, G]`, +//! `D̆ = diag(D, D_Q)`, where `Q = G D_Q G^T` is factorised the same way +//! - measurement update: scalar measurements processed sequentially with the +//! modified Agee-Turner rank-one update, after decorrelating a non-diagonal +//! `R = U_R D_R U_R^T` +//! +//! `LinearUDUpdate` and `NonLinearUDUpdate` use the existing linear and +//! nonlinear models; they differ only in how the state value is stepped and +//! how the transition Jacobian `F` is obtained. +//! +//! The shared `State` still carries `P = U D U^T` after every step, so it can +//! be read like any other filter's state, but the factors are what the filter +//! uses. If the state's covariance is changed from outside, the factors are +//! recomputed from it. + +use crate::common::linalg::udu_decomposition; +use crate::common::{FilterAllocator, na}; +use crate::epoch::Epoch; +use crate::measurement::Observation; +use crate::state::{State, StatePtr}; + +use na::allocator::Allocator; + +use crate::models::linear::{LinearMeasurement, LinearTransition}; +use crate::models::nonlinear::{NonLinearMeasurement, NonLinearTransition}; +use crate::models::traits::{ + MeasurementJacobian, MeasurementModel, TransitionJacobian, UpdateModel, +}; + +/// UD factors of the process noise, `Q = G diag(dq) G^T`, with the `Q` they +/// were computed from. +#[derive(Clone, Debug)] +struct ProcessNoiseFactors +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + q: na::OMatrix, + g: na::OMatrix, + dq: na::OVector, +} + +/// UD factors of a covariance, `P = U diag(d) U^T`. +#[derive(Clone, Debug)] +pub struct UDFactors +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + pub u: na::OMatrix, + pub d: na::OVector, + /// covariance last written to the state, to detect outside changes + written: Option>, + /// factors of the process noise, reused until Q changes + process: Option>, +} + +impl UDFactors +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + pub fn from_covariance(p: &na::OMatrix) -> Self { + let (u, d) = udu_decomposition(p); + if d.iter().any(|x| *x < 0.0) { + spdlog::error!( + "Covariance is not positive semi-definite; its UD factors have negative entries." + ); + } + UDFactors { + u, + d, + written: None, + process: None, + } + } + + /// `P = U D U^T` + pub fn covariance(&self) -> na::OMatrix { + let p = &self.u * na::OMatrix::::from_diagonal(&self.d) * self.u.transpose(); + // remove rounding asymmetry from the product + (&p + p.transpose()) * 0.5 + } + + /// Recompute the factors if the state's covariance is not the one these + /// factors last wrote, i.e. it was set from outside the filter. + fn sync(&mut self, state: &State) { + let p = state.covariance(); + if self.written.as_ref() != Some(&p) { + // the process noise did not change, keep its factors + let process = self.process.take(); + *self = Self::from_covariance(&p); + self.process = process; + } + } + + /// Write `P = U D U^T` into the state. + fn write(&mut self, state: &mut State) { + let p = self.covariance(); + state.covariance = Some(p.clone()); + self.written = Some(p); + } + + /// Time update: replace the factors with those of `F P F^T + Q`, by + /// weighted modified Gram-Schmidt on `W = [F U, G]` with weights + /// `diag(D, D_Q)` (§II.A). `W` is kept as its two blocks `F U` and `G`. + /// + /// `Q = G D_Q G^T` is factorised once and reused; it is factorised again + /// whenever `q` differs from the `Q` the factors came from, so a `Q` + /// changed between steps (non-stationary process noise) is picked up. + fn predict(&mut self, f: &na::OMatrix, q: &na::OMatrix) { + let n = self.d.len(); + if self.process.as_ref().is_none_or(|process| &process.q != q) { + let (g, dq) = udu_decomposition(q); + self.process = Some(ProcessNoiseFactors { + q: q.clone(), + g, + dq, + }); + } + let process = self.process.as_ref().unwrap(); + let dq = &process.dq; + let mut a = f * &self.u; // rows of W, first block + let mut b = process.g.clone(); // rows of W, second block + + let (nd, _) = self.u.shape_generic(); + let mut u = na::OMatrix::::identity_generic(nd, nd); + let mut d = na::OVector::::zeros_generic(nd, na::Const::<1>); + + for k in (0..n).rev() { + // weighted norm of the (already orthogonalised) row k + let dk: f64 = (0..n) + .map(|j| a[(k, j)] * self.d[j] * a[(k, j)] + b[(k, j)] * dq[j] * b[(k, j)]) + .sum(); + // a weighted sum of squares: negative only through rounding + d[k] = dk.max(0.0); + if dk <= 0.0 { + continue; + } + // orthogonalise the rows above against row k + for i in 0..k { + let dot: f64 = (0..n) + .map(|j| a[(i, j)] * self.d[j] * a[(k, j)] + b[(i, j)] * dq[j] * b[(k, j)]) + .sum(); + let uik = dot / dk; + u[(i, k)] = uik; + for j in 0..n { + a[(i, j)] -= uik * a[(k, j)]; + b[(i, j)] -= uik * b[(k, j)]; + } + } + } + self.u = u; + self.d = d; + } + + /// Measurement update for innovation `y = z - h(x)`, measurement Jacobian + /// `h` at `x`, and noise `r`. Measurements are processed one at a time + /// with the modified Agee-Turner update (§III, Appendix B), after + /// decorrelating `r` if it is not diagonal (Appendix D). Returns the state + /// correction, or `None` (logged, factors unchanged) if an innovation + /// variance is not positive or the factors are not positive semi-definite. + fn update( + &mut self, + h: &na::OMatrix, + y: &na::OVector, + r: &na::OMatrix, + ) -> Option> + where + M: na::Dim, + na::DefaultAllocator: FilterAllocator, + { + let m = r.nrows(); + let n = self.d.len(); + + let diagonal = (0..m).all(|i| (0..m).all(|j| i == j || r[(i, j)] == 0.0)); + let (h, y, r) = if diagonal { + (h.clone(), y.clone(), r.diagonal()) + } else { + // R = U_R D_R U_R^T: filter on U_R^-1 H, U_R^-1 y with noise D_R + let (ur, dr) = udu_decomposition(r); + ( + ur.solve_upper_triangular(h)?, + ur.solve_upper_triangular(y)?, + dr, + ) + }; + + let mut u = self.u.clone(); + let mut d = self.d.clone(); + let (nd, _) = u.shape_generic(); + let mut dx = na::OVector::::zeros_generic(nd, na::Const::<1>); + + for i in 0..m { + // innovation of this measurement, given the corrections so far + let e = y[i] - (0..n).map(|j| h[(i, j)] * dx[j]).sum::(); + + // w = U^T h_i, v = D w + let w: Vec = (0..n) + .map(|j| (0..n).map(|k| u[(k, j)] * h[(i, k)]).sum()) + .collect(); + let v: Vec = (0..n).map(|j| d[j] * w[j]).collect(); + + let mut k = na::OVector::::zeros_generic(nd, na::Const::<1>); + k[0] = v[0]; + + // alpha_j = r + sum_{l <= j} d_l w_l^2 is the innovation variance + // accumulated over the first j states. It is zero for a perfect + // measurement (r = 0) of states not reached yet; the updates take + // the limit for that case. A negative alpha means r or the factors + // are not positive semi-definite. + let invalid = || { + spdlog::error!( + "Invalid update: measurement noise or covariance factors are not positive semi-definite; state will not be updated." + ); + None + }; + let mut alpha = r[i] + v[0] * w[0]; + if alpha < 0.0 { + return invalid(); + } + if alpha > 0.0 { + d[0] = r[i] * d[0] / alpha; + } + + for j in 1..n { + let alpha_prev = alpha; + alpha = alpha_prev + v[j] * w[j]; + if alpha < 0.0 { + return invalid(); + } + if alpha > 0.0 { + d[j] = d[j] * alpha_prev / alpha; + } + // alpha_prev = 0 only when the gain so far (k) is zero, so the + // value of lambda does not matter then + let lambda = if alpha_prev > 0.0 { + -w[j] / alpha_prev + } else { + 0.0 + }; + for row in 0..n { + let u_old = u[(row, j)]; + u[(row, j)] = u_old + lambda * k[row]; + k[row] += v[j] * u_old; + } + } + + if alpha <= 0.0 { + spdlog::error!( + "Invalid update: innovation variance is not positive; state will not be updated." + ); + return None; + } + + // gain K = k / alpha + dx += k * (e / alpha); + } + + self.u = u; + self.d = d; + Some(dx) + } +} + +/// Propagate `state_ptr` to `epoch` one step at a time. `step` returns the +/// next state value and the transition Jacobian `F`, both at the state before +/// the step. +fn propagate( + state_ptr: &StatePtr, + factors: &mut UDFactors, + q: &na::OMatrix, + epoch: &Epoch, + mut step: impl FnMut(&State) -> (na::OVector, na::OMatrix), +) where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + let mut state = state_ptr.read().unwrap().clone(); + + let delta = { *epoch - state.epoch }.value; + if delta == 0 { + return; + } + if delta < 0 { + spdlog::warn!( + "Backward propagation is not supported by the UD filters; state is unchanged." + ); + return; + } + + factors.sync(&state); + for _ in 0..delta { + let (value, f) = step(&state); + factors.predict(&f, q); + state.value = value; + state.epoch = state.epoch + Epoch::new(1); + } + factors.write(&mut state); + *state_ptr.write().unwrap() = state; +} + +/// Measurement update of `state_ptr`, linearised about the current estimate. +fn apply( + state_ptr: &StatePtr, + factors: &mut UDFactors, + measurement: &Z, + observation: &Observation, +) where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, + Z: MeasurementModel + MeasurementJacobian, +{ + let mut state = state_ptr.read().unwrap().clone(); + factors.sync(&state); + + let z_x = measurement.projection(&state).value; + let h = measurement.jacobian(&state); + let y = &observation.value - z_x; + + if let Some(dx) = factors.update(&h, &y, &measurement.noise()) { + state.value += dx; + factors.write(&mut state); + *state_ptr.write().unwrap() = state; + } +} + +// Update Models +// Linear UD +pub struct LinearUDUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub state: StatePtr, // mutable ref. to state + pub measurement: LinearMeasurement, // measurement model + pub transition: LinearTransition, // transition model + pub factors: UDFactors, // covariance factors +} + +/// Statically sized `LinearUDUpdate`. +pub type LinearUDUpdateNM = + LinearUDUpdate, na::Const>; + +impl LinearUDUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new(measurement: LinearMeasurement, transition: LinearTransition) -> Self { + let _state = transition.state.clone(); + let factors = UDFactors::from_covariance(&_state.read().unwrap().covariance()); + LinearUDUpdate { + state: _state, + measurement, + transition, + factors, + } + } +} + +impl UpdateModel for LinearUDUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr { + if let Some(epoch) = epoch { + let f = &self.transition.f; + propagate( + &self.state, + &mut self.factors, + &self.transition.q, + epoch, + |state| (f * &state.value, f.clone()), + ); + } + &self.state + } + + fn apply(&mut self, observation: &Observation) -> &StatePtr { + // propagate state to observation epoch + self.state(Some(&observation.epoch)); + apply( + &self.state, + &mut self.factors, + &self.measurement, + observation, + ); + &self.state + } +} + +// NonLinear UD (UD Extended Kalman Filter) +pub struct NonLinearUDUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub state: StatePtr, // mutable ref. to state + pub measurement: NonLinearMeasurement, // measurement model + pub transition: NonLinearTransition, // transition model + pub factors: UDFactors, // covariance factors +} + +/// Statically sized `NonLinearUDUpdate`. +pub type NonLinearUDUpdateNM = + NonLinearUDUpdate, na::Const>; + +impl NonLinearUDUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new( + measurement: NonLinearMeasurement, + transition: NonLinearTransition, + ) -> Self { + let _state = transition.state.clone(); + let factors = UDFactors::from_covariance(&_state.read().unwrap().covariance()); + NonLinearUDUpdate { + state: _state, + measurement, + transition, + factors, + } + } +} + +impl UpdateModel for NonLinearUDUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr { + if let Some(epoch) = epoch { + let transition = &mut self.transition; + let q = transition.q.clone(); + propagate(&self.state, &mut self.factors, &q, epoch, |state| { + let f = transition.jacobian(state); + ((transition.f)(state), f) + }); + } + &self.state + } + + fn apply(&mut self, observation: &Observation) -> &StatePtr { + // propagate state to observation epoch + self.state(Some(&observation.epoch)); + apply( + &self.state, + &mut self.factors, + &self.measurement, + observation, + ); + &self.state + } +} diff --git a/src/models/unscented.rs b/src/models/unscented.rs new file mode 100644 index 0000000..accc4fa --- /dev/null +++ b/src/models/unscented.rs @@ -0,0 +1,393 @@ +//! Unscented models used by the Unscented Kalman Filter implementation. +//! +//! Like the nonlinear models, these are built from user supplied closures: a +//! propagation function `f` and a measurement function `h`. No Jacobians are +//! needed; instead a set of sigma points is pushed through `f` / `h` and the +//! mean and covariance are recovered from them (the unscented transform). +//! They are generic over the dimension type, so work with statically sized +//! (`na::Const`, see the `...N` / `...NM` aliases) and dynamically sized +//! (`na::Dyn`) matrices. +//! +//! The closures are run on copies of the state, so no lock is held while user +//! code runs. +//! +//! The sigma points need a Cholesky factor of `(n + lambda) P`, so `P` must be +//! strictly positive definite and `n + lambda = alpha^2 (n + kappa)` positive. +//! A semi-definite `P` (e.g. a component known exactly), a missing covariance +//! or `kappa <= -n` makes predict and update log an error and leave the state +//! unchanged. + +use crate::common::{FilterAllocator, na}; +use crate::epoch::Epoch; +use crate::measurement::Observation; +use crate::state::{State, StatePtr}; + +use na::allocator::Allocator; + +use crate::models::traits::{InverseMeasurement, MeasurementModel, TransitionModel, UpdateModel}; +// same closure types as the nonlinear models for h and the inverse measurement +pub use crate::models::nonlinear::{H, HInv}; + +/// One-step propagation: x(t) -> x(t + 1). Called once per sigma point, so +/// it must not keep state between calls (hence `Fn`, not `FnMut`). +pub type F = Box) -> na::OVector + Send + Sync>; + +/// Scaled sigma point parameters (van der Merwe). +/// +/// - `alpha`: spread of the sigma points around the mean. +/// - `beta`: prior knowledge of the distribution; 2 is optimal for Gaussians. +/// - `kappa`: secondary scaling. +/// +/// The default `alpha = 1, beta = 2, kappa = 0` keeps every sigma point +/// weight non-negative, so the predicted covariance stays positive +/// semi-definite. Small `alpha` (e.g. 1e-3) gives a large negative centre +/// weight, which is numerically fragile. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct UnscentedParams { + pub alpha: f64, + pub beta: f64, + pub kappa: f64, +} + +impl Default for UnscentedParams { + fn default() -> Self { + UnscentedParams { + alpha: 1.0, + beta: 2.0, + kappa: 0.0, + } + } +} + +/// Sigma points for a state, with their mean and covariance weights. +struct SigmaPoints +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + points: Vec>, + wm: Vec, + wc: Vec, +} + +impl SigmaPoints +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + /// 2n + 1 sigma points for `state`, or `None` (logged) if its covariance + /// is not positive definite. + fn new(state: &State, params: &UnscentedParams) -> Option { + let n = state.dim(); + let nf = n as f64; + let lambda = params.alpha * params.alpha * (nf + params.kappa) - nf; + + let sqrt = match na::Cholesky::new(state.covariance() * (nf + lambda)) { + Some(chol) => chol.unpack(), + None => { + spdlog::error!( + "Invalid unscented transform: covariance is not positive definite; state will not be updated." + ); + return None; + } + }; + + let at = |value| State { + value, + covariance: None, + epoch: state.epoch, + }; + let mut points = Vec::with_capacity(2 * n + 1); + points.push(at(state.value.clone())); + for i in 0..n { + points.push(at(&state.value + sqrt.column(i))); + } + for i in 0..n { + points.push(at(&state.value - sqrt.column(i))); + } + + let w = 1.0 / (2.0 * (nf + lambda)); + let wm0 = lambda / (nf + lambda); + let wc0 = wm0 + (1.0 - params.alpha * params.alpha + params.beta); + let mut wm = vec![w; 2 * n + 1]; + let mut wc = vec![w; 2 * n + 1]; + wm[0] = wm0; + wc[0] = wc0; + + Some(SigmaPoints { points, wm, wc }) + } +} + +/// Weighted mean of transformed sigma points. +fn weighted_mean(values: &[na::OVector], wm: &[f64]) -> na::OVector +where + D: na::Dim, + na::DefaultAllocator: Allocator, +{ + let (d, _) = values[0].shape_generic(); + let mut mean = na::OVector::::zeros_generic(d, na::Const::<1>); + for (v, w) in values.iter().zip(wm) { + mean.axpy(*w, v, 1.0); + } + mean +} + +/// Weighted cross covariance sum_i wc_i (a_i - a_mean)(b_i - b_mean)^T. +fn weighted_covariance( + a: &[na::OVector], + a_mean: &na::OVector, + b: &[na::OVector], + b_mean: &na::OVector, + wc: &[f64], +) -> na::OMatrix +where + A: na::Dim, + B: na::Dim, + na::DefaultAllocator: Allocator + Allocator + Allocator, +{ + let (ra, _) = a_mean.shape_generic(); + let (rb, _) = b_mean.shape_generic(); + let mut cov = na::OMatrix::::zeros_generic(ra, rb); + for ((ai, bi), w) in a.iter().zip(b).zip(wc) { + cov.ger(*w, &(ai - a_mean), &(bi - b_mean), 1.0); + } + cov +} + +// Transition Models +// Unscented (Discrete Time) +pub struct UnscentedTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + pub state: StatePtr, + pub f: F, + pub q: na::OMatrix, + pub params: UnscentedParams, +} + +/// Statically sized `UnscentedTransition`. +pub type UnscentedTransitionN = UnscentedTransition>; + +impl UnscentedTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + /// `f` advances the state by a single epoch; it is applied to every sigma + /// point once per epoch step. Uses the default `UnscentedParams`. + pub fn new(state: StatePtr, f: F, q: na::OMatrix) -> Self { + UnscentedTransition { + state, + f, + q, + params: UnscentedParams::default(), + } + } + + /// Use different sigma point parameters. + pub fn with_params(mut self, params: UnscentedParams) -> Self { + self.params = params; + self + } +} + +impl TransitionModel for UnscentedTransition +where + N: na::Dim, + na::DefaultAllocator: Allocator + Allocator, +{ + fn state(&mut self, epoch: &Epoch) -> &StatePtr { + let mut state = self.state.read().unwrap().clone(); + + let delta = { *epoch - state.epoch }.value; + if delta == 0 { + return &self.state; + } + if delta < 0 { + // a general nonlinear f has no inverse + spdlog::warn!( + "Backward propagation is not supported for unscented transitions; state is unchanged." + ); + return &self.state; + } + + for _ in 0..delta { + let sigma = match SigmaPoints::new(&state, &self.params) { + Some(sigma) => sigma, + None => return &self.state, + }; + let propagated: Vec<_> = sigma.points.iter().map(|p| (self.f)(p)).collect(); + + let mean = weighted_mean(&propagated, &sigma.wm); + let cov = + weighted_covariance(&propagated, &mean, &propagated, &mean, &sigma.wc) + &self.q; + + state.value = mean; + state.covariance = Some(cov); + state.epoch = state.epoch + Epoch::new(1); + } + + *self.state.write().unwrap() = state; + &self.state + } +} + +// Measurement Models +// Unscented +pub struct UnscentedMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub h: H, + pub r: na::OMatrix, + h_inv: Option>, +} + +/// Statically sized `UnscentedMeasurement`. +pub type UnscentedMeasurementNM = + UnscentedMeasurement, na::Const>; + +impl UnscentedMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new(h: H, r: na::OMatrix) -> Self { + UnscentedMeasurement { h, r, h_inv: None } + } + + /// Supply an inverse measurement function h^-1(z) -> x. + /// Without one, `inverse` returns `None`. + pub fn with_inverse(mut self, h_inv: HInv) -> Self { + self.h_inv = Some(h_inv); + self + } +} + +impl MeasurementModel for UnscentedMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn projection(&self, state: &State) -> Observation { + Observation { + value: (self.h)(state), + epoch: state.epoch, + } + } + + fn noise(&self) -> na::OMatrix { + self.r.clone() + } +} + +impl InverseMeasurement for UnscentedMeasurement +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + /// Uses the function supplied with `with_inverse`, if any. + fn inverse(&self, obs: &Observation) -> Option> { + self.h_inv.as_ref().map(|h_inv| h_inv(obs)) + } +} + +// Update +pub struct UnscentedUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub state: StatePtr, // mutable ref. to state + pub measurement: UnscentedMeasurement, // measurement model + pub transition: UnscentedTransition, // transition model (and sigma point parameters) +} + +/// Statically sized `UnscentedUpdate`. +pub type UnscentedUpdateNM = + UnscentedUpdate, na::Const>; + +impl UnscentedUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + pub fn new( + measurement: UnscentedMeasurement, + transition: UnscentedTransition, + ) -> Self { + let _state = transition.state.clone(); + UnscentedUpdate { + state: _state, + measurement, + transition, + } + } +} + +impl UpdateModel for UnscentedUpdate +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr { + match epoch { + Some(epoch) => self.transition.state(epoch), + _ => &self.transition.state, + } + } + + fn apply(&mut self, observation: &Observation) -> &StatePtr { + // propagate state to observation epoch + _ = self.transition.state(&observation.epoch); + let mut state = self.state.read().unwrap().clone(); + + // push sigma points through the measurement function + let sigma = match SigmaPoints::new(&state, &self.transition.params) { + Some(sigma) => sigma, + None => return &self.state, + }; + let values: Vec<_> = sigma.points.iter().map(|p| p.value.clone()).collect(); + let projected: Vec<_> = sigma + .points + .iter() + .map(|p| (self.measurement.h)(p)) + .collect(); + + // predicted measurement, innovation covariance and cross covariance + let z_mean = weighted_mean(&projected, &sigma.wm); + let s = weighted_covariance(&projected, &z_mean, &projected, &z_mean, &sigma.wc) + + self.measurement.noise(); + let pxz = weighted_covariance(&values, &state.value, &projected, &z_mean, &sigma.wc); + + let s_inv = match s.clone().try_inverse() { + Some(s_inv) => s_inv, + None => { + spdlog::error!( + "Invalid update: Innovation covariance S is singular; state will not be updated." + ); + return &self.state; + } + }; + + // gain, state and covariance update + let k = &pxz * s_inv; + state.value = &state.value + &k * (&observation.value - &z_mean); + let cov = state.covariance() - &k * &s * k.transpose(); + // remove rounding asymmetry + state.covariance = Some((&cov + cov.transpose()) * 0.5); + + *self.state.write().unwrap() = state; + &self.state + } +} diff --git a/src/models/update.rs b/src/models/update.rs new file mode 100644 index 0000000..c27be58 --- /dev/null +++ b/src/models/update.rs @@ -0,0 +1,50 @@ +//! Kalman measurement update shared by the linear and nonlinear update models. + +use crate::common::{FilterAllocator, na}; +use crate::state::State; + +/// Apply the Kalman measurement update to `state` in place. +/// +/// `z_x` is the predicted measurement h(x) and `h` its Jacobian at `x` (the +/// measurement matrix for a linear model). The covariance uses the Joseph form +/// `(I - KH) P (I - KH)^T + K R K^T`, which stays symmetric positive +/// semi-definite under rounding. Returns `false`, leaving `state` unchanged, +/// if the innovation covariance is singular. +pub fn kalman_update( + state: &mut State, + z: &na::OVector, + z_x: &na::OVector, + h: &na::OMatrix, + r: &na::OMatrix, +) -> bool +where + N: na::Dim, + M: na::Dim, + na::DefaultAllocator: FilterAllocator, +{ + let cov = state.covariance(); + + // innovation covariance and gain + let h_t = h.transpose(); + let s = h * &cov * &h_t + r; + let s_inv = match s.try_inverse() { + Some(s_inv) => s_inv, + None => { + spdlog::error!( + "Invalid update: Innovation covariance S is singular; state will not be updated." + ); + return false; + } + }; + let k = &cov * &h_t * s_inv; + + // update state with the innovation + state.value = &state.value + &k * (z - z_x); + + // Joseph form covariance + let (n, _) = cov.shape_generic(); + let i_kh = na::OMatrix::::identity_generic(n, n) - &k * h; + state.covariance = Some(&i_kh * &cov * i_kh.transpose() + &k * r * k.transpose()); + + true +} diff --git a/src/python/algorithms.rs b/src/python/algorithms.rs index 2e3bfec..21c6de0 100644 --- a/src/python/algorithms.rs +++ b/src/python/algorithms.rs @@ -1,115 +1,369 @@ #![allow(unsafe_op_in_unsafe_fn)] -use spdlog; +//! Python wrappers around the Rust filters in `crate::algorithms`: +//! +//! | Rust | Python | +//! |------------------------------------------|------------------------| +//! | `KalmanFilterDyn` | `KalmanFilter` | +//! | `ExtendedKalmanFilterDyn` | `ExtendedKalmanFilter` | +//! | `UnscentedKalmanFilterDyn` | `UnscentedKalmanFilter`| +//! | `UDKalmanFilterDyn` | `UDKalmanFilter` | +//! | `UDExtendedKalmanFilterDyn` | `UDExtendedKalmanFilter` | +//! | `KalmanFilterT` | `CustomFilter` | +//! +//! `KalmanFilter` and `UDKalmanFilter` are computed entirely in Rust from numpy +//! matrices. The extended, UD extended and unscented filters do the maths in Rust and +//! call back into Python only for the model functions (see `callbacks`). `CustomFilter` runs +//! whatever update your Python `UpdateModel` implements. +//! +//! Each class holds a `FilterCore` (see `filter`) around its Rust filter and +//! forwards `predict`/`update`/`state`/`reset` to it. +//! +//! States returned by `predict`/`update`/`state` are live, zero-copy views of +//! the filter's internal state, not snapshots: they change on the next call. +//! Copy the arrays (`state.value.copy()`) if you need to keep one. +use numpy::PyReadonlyArray2; use pyo3::prelude::*; -use crate::common::na as na; -use crate::filter::Filter; -use crate::epoch::Epoch; -use crate::measurement::Observation; -use crate::state::StatePtr; - -use crate::models::traits::UpdateModel; +use crate::algorithms::{ + ExtendedKalmanFilterDyn, KalmanFilterDyn, KalmanFilterT, UDExtendedKalmanFilterDyn, + UDKalmanFilterDyn, UnscentedKalmanFilterDyn, +}; +use crate::common::na; +use crate::models::{ + LinearMeasurement, LinearTransition, LinearUDUpdate, LinearUpdate, NonLinearMeasurement, + NonLinearTransition, NonLinearUDUpdate, NonLinearUpdate, UnscentedMeasurement, UnscentedParams, + UnscentedTransition, UnscentedUpdate, +}; +use crate::python::arrays::{check_shape, check_square, to_dmatrix}; +use crate::python::callbacks::{ErrorSlot, check_callable, matrix_model, vector_model}; use crate::python::epoch::PyEpoch; -use crate::python::state::PyState; +use crate::python::filter::FilterCore; use crate::python::measurement::PyObservation; -use crate::python::update::PyUpdateModel; +use crate::python::models::PyUpdateModel; +use crate::python::state::PyState; + +/// The `#[pymethods]` of a Python filter class: the constructor given, plus +/// the methods every filter shares, forwarded to the class's +/// `core: FilterCore<_>`. +macro_rules! py_filter { + ($py:ident { $($constructor:tt)* }) => { + #[pymethods] + impl $py { + $($constructor)* + + fn predict(&mut self, epoch: &PyEpoch) -> PyResult { + self.core.predict(epoch) + } + + fn update(&mut self, observation: &PyObservation) -> PyResult { + self.core.update(observation) + } + + #[getter] + fn get_state(&self) -> PyState { + self.core.get_state() + } + #[setter] + fn set_state(&self, state: PyState) -> PyResult<()> { + self.core.set_state(&state) + } -#[pyclass(name="KalmanFilter")] + /// Zero the state and reset the covariance to identity. + fn reset(&mut self) { + self.core.reset(); + } + } + }; +} + +// Kalman Filter +/// Linear Kalman filter, computed in Rust. +/// +/// f, q: (N, N) transition and process noise; h: (M, N) measurement matrix; +/// r: (M, M) measurement noise. Returned states are live views of the filter state. +#[pyclass(name = "KalmanFilter")] pub struct PyKalmanFilter { - update: PyUpdateModel, - state: PyState, + core: FilterCore, } -#[pymethods] -impl PyKalmanFilter { +py_filter!(PyKalmanFilter { #[new] - pub fn new(update: Py) -> Self { - Python::attach(| py| { - let _update: PyUpdateModel = update.extract(py).unwrap(); - let _state: PyState = update.borrow(py).state.clone(); - - PyKalmanFilter { - update: _update, - state: _state - } - }) + #[pyo3(signature = (state, *, f, q, h, r))] + fn new( + state: PyState, + f: PyReadonlyArray2, + q: PyReadonlyArray2, + h: PyReadonlyArray2, + r: PyReadonlyArray2, + ) -> PyResult { + let (f, q) = (to_dmatrix(f.as_array()), to_dmatrix(q.as_array())); + let (h, r) = (to_dmatrix(h.as_array()), to_dmatrix(r.as_array())); - } + let n = state.inner.read().unwrap().dim(); + let m = h.nrows(); + check_square(&f, n, "f")?; + check_square(&q, n, "q")?; + check_shape(&h, (m, n), "h")?; + check_square(&r, m, "r")?; - #[pyo3(name="predict")] - fn predict_state(&mut self, epoch: &PyEpoch) -> PyState { - let state = self.predict(&epoch.inner).clone(); - PyState { inner: state } + let filter = KalmanFilterDyn::new(LinearUpdate::new( + LinearMeasurement::new(h, r), + LinearTransition::new(state.inner.clone(), f, q), + )); + Ok(PyKalmanFilter { core: FilterCore::new(filter, state, None, Some(m)) }) } +}); - #[pyo3(name="update")] - fn update_state(&mut self, observation: &PyObservation) -> PyState { - let state = self.update(&observation.inner).clone(); - PyState { inner: state } +// Extended Kalman Filter +/// Extended Kalman filter: the maths in Rust, the model functions in Python. +/// +/// f(x, t), h(x, t) and their Jacobians jac_f (N, N), jac_h (M, N); q (N, N) and +/// r (M, M) noise; t is the integer epoch. Returned states are live views of the +/// filter state. +#[pyclass(name = "ExtendedKalmanFilter")] +pub struct PyExtendedKalmanFilter { + core: FilterCore, +} + +py_filter!(PyExtendedKalmanFilter { + // the parameters are the Python keyword arguments + #[allow(clippy::too_many_arguments)] + #[new] + #[pyo3(signature = (state, *, f, jac_f, q, h, jac_h, r))] + fn new( + py: Python<'_>, + state: PyState, + f: Py, + jac_f: Py, + q: PyReadonlyArray2, + h: Py, + jac_h: Py, + r: PyReadonlyArray2, + ) -> PyResult { + for (func, name) in [(&f, "f"), (&jac_f, "jac_f"), (&h, "h"), (&jac_h, "jac_h")] { + check_callable(func, py, name)?; + } + + let q = to_dmatrix(q.as_array()); + let r = to_dmatrix(r.as_array()); + let n = state.inner.read().unwrap().dim(); + let m = r.nrows(); + check_square(&q, n, "q")?; + check_square(&r, m, "r")?; + + let error = ErrorSlot::default(); + let transition = NonLinearTransition::new( + state.inner.clone(), + Box::new(vector_model(f, "f", n, error.clone())), + Box::new(matrix_model(jac_f, "jac_f", n, n, error.clone())), + q, + ); + let measurement = NonLinearMeasurement::new( + Box::new(vector_model(h, "h", m, error.clone())), + Box::new(matrix_model(jac_h, "jac_h", m, n, error.clone())), + r, + ); + let filter = ExtendedKalmanFilterDyn::new(NonLinearUpdate::new(measurement, transition)); + + Ok(PyExtendedKalmanFilter { core: FilterCore::new(filter, state, Some(error), Some(m)) }) } +}); - #[getter] //TODO: conflict with Filter.state() - fn get_state(&self) -> PyState { - let _state = self.state(); - PyState { inner: _state } +// Unscented Kalman Filter +/// Unscented Kalman filter: no Jacobians, f(x, t) and h(x, t) are evaluated at +/// 2N + 1 sigma points. +/// +/// q (N, N) and r (M, M) noise; alpha, beta, kappa: sigma point parameters. Needs a +/// positive definite covariance. Returned states are live views of the filter state. +#[pyclass(name = "UnscentedKalmanFilter")] +pub struct PyUnscentedKalmanFilter { + core: FilterCore, +} + +py_filter!(PyUnscentedKalmanFilter { + // the parameters are the Python keyword arguments + #[allow(clippy::too_many_arguments)] + #[new] + #[pyo3(signature = (state, *, f, q, h, r, alpha=1.0, beta=2.0, kappa=0.0))] + fn new( + py: Python<'_>, + state: PyState, + f: Py, + q: PyReadonlyArray2, + h: Py, + r: PyReadonlyArray2, + alpha: f64, + beta: f64, + kappa: f64, + ) -> PyResult { + for (func, name) in [(&f, "f"), (&h, "h")] { + check_callable(func, py, name)?; + } + + let q = to_dmatrix(q.as_array()); + let r = to_dmatrix(r.as_array()); + let n = state.inner.read().unwrap().dim(); + let m = r.nrows(); + check_square(&q, n, "q")?; + check_square(&r, m, "r")?; + + let error = ErrorSlot::default(); + let transition = UnscentedTransition::new( + state.inner.clone(), + Box::new(vector_model(f, "f", n, error.clone())), + q, + ).with_params(UnscentedParams { alpha, beta, kappa }); + let measurement = UnscentedMeasurement::new( + Box::new(vector_model(h, "h", m, error.clone())), + r, + ); + let filter = UnscentedKalmanFilterDyn::new(UnscentedUpdate::new(measurement, transition)); + + Ok(PyUnscentedKalmanFilter { core: FilterCore::new(filter, state, Some(error), Some(m)) }) } +}); - #[setter] - fn set_state(&mut self, state: PyState) { +// UD Kalman Filter +/// Kalman filter in UD factorised form (P = U D U^T), with the same models as +/// KalmanFilter. Returned states are live views of the filter state. +#[pyclass(name = "UDKalmanFilter")] +pub struct PyUDKalmanFilter { + core: FilterCore, +} - let py_state = state.inner.read().unwrap(); - let mut inner_state = self.state.inner.write().unwrap(); +py_filter!(PyUDKalmanFilter { + #[new] + #[pyo3(signature = (state, *, f, q, h, r))] + fn new( + state: PyState, + f: PyReadonlyArray2, + q: PyReadonlyArray2, + h: PyReadonlyArray2, + r: PyReadonlyArray2, + ) -> PyResult { + let (f, q) = (to_dmatrix(f.as_array()), to_dmatrix(q.as_array())); + let (h, r) = (to_dmatrix(h.as_array()), to_dmatrix(r.as_array())); - inner_state.value = py_state.value.clone(); - inner_state.covariance = py_state.covariance.clone(); - inner_state.epoch = py_state.epoch; + let n = state.inner.read().unwrap().dim(); + let m = h.nrows(); + check_square(&f, n, "f")?; + check_square(&q, n, "q")?; + check_shape(&h, (m, n), "h")?; + check_square(&r, m, "r")?; + let filter = UDKalmanFilterDyn::new(LinearUDUpdate::new( + LinearMeasurement::new(h, r), + LinearTransition::new(state.inner.clone(), f, q), + )); + Ok(PyUDKalmanFilter { core: FilterCore::new(filter, state, None, Some(m)) }) } +}); + +// UD Extended Kalman Filter +/// Extended Kalman filter in UD factorised form (P = U D U^T), with the same models +/// as ExtendedKalmanFilter. Returned states are live views of the filter state. +#[pyclass(name = "UDExtendedKalmanFilter")] +pub struct PyUDExtendedKalmanFilter { + core: FilterCore, +} - #[pyo3(name="reset")] - fn reset_state(&mut self) { - self.reset(); +py_filter!(PyUDExtendedKalmanFilter { + // the parameters are the Python keyword arguments + #[allow(clippy::too_many_arguments)] + #[new] + #[pyo3(signature = (state, *, f, jac_f, q, h, jac_h, r))] + fn new( + py: Python<'_>, + state: PyState, + f: Py, + jac_f: Py, + q: PyReadonlyArray2, + h: Py, + jac_h: Py, + r: PyReadonlyArray2, + ) -> PyResult { + for (func, name) in [(&f, "f"), (&jac_f, "jac_f"), (&h, "h"), (&jac_h, "jac_h")] { + check_callable(func, py, name)?; } + + let q = to_dmatrix(q.as_array()); + let r = to_dmatrix(r.as_array()); + let n = state.inner.read().unwrap().dim(); + let m = r.nrows(); + check_square(&q, n, "q")?; + check_square(&r, m, "r")?; + + let error = ErrorSlot::default(); + let transition = NonLinearTransition::new( + state.inner.clone(), + Box::new(vector_model(f, "f", n, error.clone())), + Box::new(matrix_model(jac_f, "jac_f", n, n, error.clone())), + q, + ); + let measurement = NonLinearMeasurement::new( + Box::new(vector_model(h, "h", m, error.clone())), + Box::new(matrix_model(jac_h, "jac_h", m, n, error.clone())), + r, + ); + let filter = UDExtendedKalmanFilterDyn::new(NonLinearUDUpdate::new(measurement, transition)); + + Ok(PyUDExtendedKalmanFilter { core: FilterCore::new(filter, state, Some(error), Some(m)) }) + } +}); + +// Custom Filter +/// Filter built from a Python `UpdateModel`: the Python equivalent of a Rust +/// `KalmanFilterT` with your own update model. The update maths is whatever +/// your model implements. See `KalmanFilter` for the Rust-computed filter. +#[pyclass(name = "CustomFilter")] +pub struct PyCustomFilter { + core: FilterCore>, } -impl Filter for PyKalmanFilter { +#[pymethods] +impl PyCustomFilter { + #[new] + pub fn new(update: PyRef<'_, PyUpdateModel>) -> Self { + let update = update.clone(); + let state = update.state.clone(); + // the measurement is a Python object, so its size is not known here + PyCustomFilter { + core: FilterCore::new(KalmanFilterT::new(update), state, None, None), + } + } - type ObservationType = Observation; - type StateType = StatePtr; + fn predict(&mut self, epoch: &PyEpoch) -> PyResult { + self.core.predict(epoch) + } - fn predict(&mut self, epoch: &Epoch) -> Self::StateType { - // return a reference counted ptr - self.update.state(epoch).clone() + fn update(&mut self, observation: &PyObservation) -> PyResult { + self.core.update(observation) } - fn update(&mut self, observation: &Observation) -> Self::StateType { - self.update.apply(observation).clone() + #[getter] + fn get_state(&self) -> PyState { + self.core.get_state() } - fn state(&self) -> Self::StateType { - // return a reference counted ptr - self.state.inner.clone() + #[setter] + fn set_state(&self, state: PyState) -> PyResult<()> { + self.core.set_state(&state) } - fn reset(&mut self) { - // make this call optional such that the expects dont fail shoudl the method not exists - Python::attach(| py| { - let py_obj = self.update.model.bind(py); - let result = py_obj - .call_method("reset", (), None); - - let _ = match result { - Ok(res) => res, - Err(_) => { - spdlog::warn!("No update.reset() method found."); - return; - } - }; - }); + /// Calls the update model's `reset()`, if it has one; errors it raises + /// propagate. The state belongs to your models, so the filter does not + /// reset it itself. + fn reset(&mut self, py: Python<'_>) -> PyResult<()> { + let model = self.core.filter.update.model.bind(py); + if !model.hasattr("reset")? { + spdlog::warn!("No update.reset() method found."); + return Ok(()); + } + model.call_method0("reset")?; + Ok(()) } } - diff --git a/src/python/arrays.rs b/src/python/arrays.rs index cbb5b96..11fad28 100644 --- a/src/python/arrays.rs +++ b/src/python/arrays.rs @@ -1,5 +1,55 @@ +use numpy::ndarray::ArrayView2; +use numpy::{PyArray1, PyArray2, PyArrayMethods, ToPyArray}; use pyo3::prelude::*; -use numpy::{ PyArray1, PyArray2 }; + +use crate::common::na; pub type PyVector<'py> = Bound<'py, PyArray1>; -pub type PyMatrix<'py> = Bound<'py, PyArray2>; \ No newline at end of file +pub type PyMatrix<'py> = Bound<'py, PyArray2>; + +/// numpy (row-major view) -> nalgebra matrix, preserving the element layout. +/// +/// nalgebra stores column-major, so the elements must be read row by row. +pub fn to_dmatrix(array: ArrayView2<'_, f64>) -> na::DMatrix { + let (nrows, ncols) = array.dim(); + na::DMatrix::::from_row_iterator(nrows, ncols, array.iter().copied()) +} + +/// nalgebra matrix -> numpy (C-order) array of the same shape. +pub fn to_pymatrix<'py>(py: Python<'py>, matrix: &na::DMatrix) -> PyMatrix<'py> { + // transpose so the column-major slice reads out in row-major order + matrix + .transpose() + .as_slice() + .to_pyarray(py) + .reshape([matrix.nrows(), matrix.ncols()]) + .unwrap() +} + +/// Check a matrix has the given shape, for validating Python inputs before +/// they reach nalgebra (which panics on mismatched dimensions). +pub fn check_shape(m: &na::DMatrix, shape: (usize, usize), name: &str) -> PyResult<()> { + if m.shape() == shape { + Ok(()) + } else { + Err(pyo3::exceptions::PyValueError::new_err(format!( + "{name} must have shape {shape:?}, got {:?}", + m.shape() + ))) + } +} + +pub fn check_square(m: &na::DMatrix, dim: usize, name: &str) -> PyResult<()> { + check_shape(m, (dim, dim), name) +} + +pub fn check_length(v: &na::DVector, len: usize, name: &str) -> PyResult<()> { + if v.len() == len { + Ok(()) + } else { + Err(pyo3::exceptions::PyValueError::new_err(format!( + "{name} has {} elements, expected {len}", + v.len() + ))) + } +} diff --git a/src/python/callbacks.rs b/src/python/callbacks.rs new file mode 100644 index 0000000..9971102 --- /dev/null +++ b/src/python/callbacks.rs @@ -0,0 +1,148 @@ +//! Python model functions as Rust model closures. +//! +//! Used by `ExtendedKalmanFilter` to pass Python callables `f(x, t)` / +//! `h(x, t)` (and their Jacobians) to the Rust `NonLinear*` models. The +//! callables take the state value as a numpy array and the epoch as an int, +//! and may return anything array-like (so autodiff libraries such as JAX work). + +use std::sync::{Arc, Mutex}; + +use numpy::{PyReadonlyArray2, PyReadonlyArrayDyn, ToPyArray}; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyAny, PyDict}; + +use crate::common::na; +use crate::state::{State, StatePtr}; + +use crate::python::arrays::to_dmatrix; + +/// First error raised by a Python model function during a filter call. +/// +/// The Rust model closures cannot fail, so on error they record it here and +/// return a placeholder; the wrapper then restores the state and raises. +pub type ErrorSlot = Arc>>; + +/// Coerce anything array-like (numpy, jax, lists...) to a float64 ndarray. +fn as_f64_array<'py>(py: Python<'py>, obj: Bound<'py, PyAny>) -> PyResult> { + let np = py.import("numpy")?; + let kwargs = PyDict::new(py); + kwargs.set_item("dtype", np.getattr("float64")?)?; + np.call_method("asarray", (obj,), Some(&kwargs)) +} + +pub fn check_callable(func: &Py, py: Python<'_>, name: &str) -> PyResult<()> { + if func.bind(py).is_callable() { + Ok(()) + } else { + Err(PyTypeError::new_err(format!("{name} must be callable"))) + } +} + +/// Call `func(x, t)` and read the result as a vector of `len` elements. +fn call_vector( + py: Python<'_>, + func: &Py, + name: &str, + state: &State, + len: usize, +) -> PyResult> { + let x = state.value.as_slice().to_pyarray(py); + let result = as_f64_array(py, func.bind(py).call1((x, state.epoch.value))?)?; + let array: PyReadonlyArrayDyn = result.extract()?; + let view = array.as_array(); + if view.len() != len { + return Err(PyValueError::new_err(format!( + "{name}() returned {} elements, expected {len}", + view.len() + ))); + } + Ok(na::DVector::::from_iterator(len, view.iter().copied())) +} + +/// Call `func(x, t)` and read the result as an `nrows` x `ncols` matrix. +fn call_matrix( + py: Python<'_>, + func: &Py, + name: &str, + state: &State, + nrows: usize, + ncols: usize, +) -> PyResult> { + let x = state.value.as_slice().to_pyarray(py); + let result = as_f64_array(py, func.bind(py).call1((x, state.epoch.value))?)?; + let array: PyReadonlyArray2 = result + .extract() + .map_err(|_| PyValueError::new_err(format!("{name}() must return a 2D array")))?; + if array.as_array().dim() != (nrows, ncols) { + return Err(PyValueError::new_err(format!( + "{name}() returned shape {:?}, expected ({nrows}, {ncols})", + array.as_array().dim() + ))); + } + Ok(to_dmatrix(array.as_array())) +} + +/// Wrap a Python model function as a Rust model closure returning a vector. +pub fn vector_model( + func: Py, + name: &'static str, + len: usize, + error: ErrorSlot, +) -> impl Fn(&State) -> na::DVector + Send + Sync { + move |state| { + let placeholder = || na::DVector::::zeros(len); + if error.lock().unwrap().is_some() { + return placeholder(); + } + Python::attach(|py| call_vector(py, &func, name, state, len)).unwrap_or_else(|err| { + *error.lock().unwrap() = Some(err); + placeholder() + }) + } +} + +/// Wrap a Python model function as a Rust model closure returning a matrix. +pub fn matrix_model( + func: Py, + name: &'static str, + nrows: usize, + ncols: usize, + error: ErrorSlot, +) -> impl Fn(&State) -> na::DMatrix + Send + Sync { + move |state| { + let placeholder = || na::DMatrix::::zeros(nrows, ncols); + if error.lock().unwrap().is_some() { + return placeholder(); + } + Python::attach(|py| call_matrix(py, &func, name, state, nrows, ncols)).unwrap_or_else( + |err| { + *error.lock().unwrap() = Some(err); + placeholder() + }, + ) + } +} + +/// Run a filter operation whose models call back into Python. If a Python +/// model function raised (or returned the wrong shape), restore `state` to +/// what it was before the call and re-raise that error. +pub fn guarded( + filter: &mut F, + state: &StatePtr, + error: &ErrorSlot, + op: impl FnOnce(&mut F) -> T, +) -> PyResult { + let backup = state.read().unwrap().clone(); + *error.lock().unwrap() = None; + + let result = op(filter); + + match error.lock().unwrap().take() { + Some(err) => { + *state.write().unwrap() = backup; + Err(err) + } + None => Ok(result), + } +} diff --git a/src/python/epoch.rs b/src/python/epoch.rs index df9742f..7e55f22 100644 --- a/src/python/epoch.rs +++ b/src/python/epoch.rs @@ -1,9 +1,9 @@ #![allow(unsafe_op_in_unsafe_fn)] -use pyo3::prelude::*; use crate::epoch::Epoch; +use pyo3::prelude::*; -#[pyclass(name="Epoch")] +#[pyclass(name = "Epoch")] #[derive(Clone, Debug)] pub struct PyEpoch { pub inner: Epoch, @@ -13,10 +13,10 @@ pub struct PyEpoch { impl PyEpoch { #[new] fn new(value: i64) -> PyEpoch { - PyEpoch{ - inner: Epoch::new(value) + PyEpoch { + inner: Epoch::new(value), } - } + } #[getter] fn value(&self) -> i64 { @@ -26,4 +26,4 @@ impl PyEpoch { fn __repr__(&self) -> PyResult { Ok(format!("{:?}", self.inner)) } -} \ No newline at end of file +} diff --git a/src/python/filter.rs b/src/python/filter.rs new file mode 100644 index 0000000..1928f6f --- /dev/null +++ b/src/python/filter.rs @@ -0,0 +1,85 @@ +//! The part every Python filter class shares. +//! +//! PyO3 classes cannot be generic, so each Python filter class in +//! `algorithms` is its own struct. Each holds a `FilterCore` around its Rust +//! filter, and its `predict`/`update`/`state`/`reset` methods forward to it, +//! so that behaviour is written once. + +use pyo3::prelude::*; + +use crate::common::na; +use crate::filter::Filter; +use crate::measurement::Observation; +use crate::state::StatePtr; + +use crate::python::arrays::check_length; +use crate::python::callbacks::{ErrorSlot, guarded}; +use crate::python::epoch::PyEpoch; +use crate::python::measurement::PyObservation; +use crate::python::state::PyState; + +/// A Rust filter, the state it shares with Python, and what is needed to run +/// it safely from Python. +pub struct FilterCore { + pub filter: F, + state: PyState, + /// error slot of the filter's Python model functions, if it has any + error: Option, + /// number of elements an observation must have, if known + measurement_dim: Option, +} + +impl FilterCore +where + F: Filter, ObservationType = Observation>, +{ + pub fn new( + filter: F, + state: PyState, + error: Option, + measurement_dim: Option, + ) -> Self { + FilterCore { + filter, + state, + error, + measurement_dim, + } + } + + /// Run `op` on the filter; if it has Python model functions and one of + /// them raised, restore the state and re-raise. + fn run(&mut self, op: impl FnOnce(&mut F) -> StatePtr) -> PyResult { + let inner = match &self.error { + Some(error) => guarded(&mut self.filter, &self.state.inner, error, op)?, + None => op(&mut self.filter), + }; + Ok(PyState { inner }) + } + + pub fn predict(&mut self, epoch: &PyEpoch) -> PyResult { + self.run(|filter| filter.predict(&epoch.inner)) + } + + pub fn update(&mut self, observation: &PyObservation) -> PyResult { + if let Some(m) = self.measurement_dim { + check_length(&observation.inner.value, m, "observation")?; + } + self.run(|filter| filter.update(&observation.inner)) + } + + pub fn get_state(&self) -> PyState { + PyState { + inner: self.filter.state(), + } + } + + pub fn set_state(&self, state: &PyState) -> PyResult<()> { + self.state.assign(state) + } + + /// Zero the state and reset the covariance to identity. + pub fn reset(&mut self) { + self.filter.reset(); + } +} diff --git a/src/python/logging.rs b/src/python/logging.rs new file mode 100644 index 0000000..c67bfcf --- /dev/null +++ b/src/python/logging.rs @@ -0,0 +1,71 @@ +//! Routes Rust-side `spdlog` messages into Python's standard `logging` module, +//! under the logger name `determined`, so they can be filtered, captured and +//! redirected like any other Python log output. + +use std::sync::Arc; + +use pyo3::prelude::*; +use spdlog::formatter::Formatter; +use spdlog::sink::Sink; +use spdlog::{ErrorHandler, Level, LevelFilter, Logger, Record}; + +const LOGGER_NAME: &str = "determined"; + +/// Python `logging` level numbers. +fn py_level(level: Level) -> u32 { + match level { + Level::Critical => 50, + Level::Error => 40, + Level::Warn => 30, + Level::Info => 20, + Level::Debug => 10, + Level::Trace => 5, + } +} + +/// spdlog sink that forwards every record to `logging.getLogger("determined")`. +/// +/// Filtering and formatting are left to Python, so the level filter, +/// formatter and error handler are intentionally ignored. +struct PyLogSink; + +impl Sink for PyLogSink { + fn log(&self, record: &Record) -> spdlog::Result<()> { + Python::attach(|py| { + // never fail the caller because of a logging problem + let _ = py + .import("logging") + .and_then(|logging| logging.call_method1("getLogger", (LOGGER_NAME,))) + .and_then(|logger| { + logger.call_method1("log", (py_level(record.level()), record.payload())) + }); + }); + Ok(()) + } + + fn flush(&self) -> spdlog::Result<()> { + Ok(()) + } + + fn level_filter(&self) -> LevelFilter { + LevelFilter::All + } + + fn set_level_filter(&self, _level_filter: LevelFilter) {} + + fn set_formatter(&self, _formatter: Box) {} + + fn set_error_handler(&self, _handler: Option) {} +} + +/// Make `spdlog` use Python's `logging`. Called once on module import. +pub fn install() -> PyResult<()> { + let logger = Logger::builder() + .name(LOGGER_NAME) + .level_filter(LevelFilter::All) + .sink(Arc::new(PyLogSink)) + .build() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + spdlog::set_default_logger(Arc::new(logger)); + Ok(()) +} diff --git a/src/python/measurement.rs b/src/python/measurement.rs index 66935f2..e63a934 100644 --- a/src/python/measurement.rs +++ b/src/python/measurement.rs @@ -1,197 +1,53 @@ #![allow(unsafe_op_in_unsafe_fn)] +use numpy::PyReadonlyArray1; +use numpy::ToPyArray; use pyo3::prelude::*; -use pyo3::types::PyAny; -use numpy::{ PyReadonlyArray1 }; -use numpy::{ ToPyArray ,PyArrayMethods }; -use crate::common::na as na; +use crate::common::na; use crate::measurement::Observation; -use crate::models::MeasurementModel; -use crate::state::State; +use crate::python::arrays::PyVector; use crate::python::epoch::PyEpoch; -use crate::python::state::PyState; -use crate::python::arrays::{ PyVector, PyMatrix }; - -#[pyclass(name="Observation")] +#[pyclass(name = "Observation")] #[derive(Clone, Debug)] pub struct PyObservation { pub inner: Observation, -} +} #[pymethods] impl PyObservation { #[new] - fn new( - value: PyReadonlyArray1, - epoch: &PyEpoch - ) -> PyResult { + #[pyo3(signature = (value, *, epoch))] + fn new(value: PyReadonlyArray1, epoch: &PyEpoch) -> PyResult { // numpy -> ndarray (zero-copy view) -> nalgebra let array = value.as_array(); - + // TODO: handle ndarray with more than 1 dimension! - let vector = na::DVector::::from_iterator( - array.len(), - array.iter().copied() - ); - + let vector = na::DVector::::from_iterator(array.len(), array.iter().copied()); + Ok(PyObservation { inner: Observation { value: vector, epoch: epoch.inner, - } + }, }) } #[getter] fn value<'py>(&self, py: Python<'py>) -> PyVector<'py> { - self.inner.value.as_slice() - .to_pyarray(py) + self.inner.value.as_slice().to_pyarray(py) } #[getter] fn epoch(&self) -> PyEpoch { - PyEpoch { inner: self.inner.epoch } + PyEpoch { + inner: self.inner.epoch, + } } fn __repr__(&self) -> PyResult { Ok(format!("{:0.6}", self.inner)) } } - -// Wrapper for Python-defined Measurement models -#[pyclass(name="MeasurementModel")] -pub struct PyMeasurementModel { - pub model: Py, -} - -impl Clone for PyMeasurementModel { - fn clone(&self) -> Self { - Python::attach(| py| { - let _ref = self.model.clone_ref(py); - Self { - model: _ref, - } - }) - } -} - -#[pymethods] -impl PyMeasurementModel { - - #[new] - pub fn new(model: Py) -> Self { - PyMeasurementModel { model: model } - } - - #[getter] - fn get_model(&self) -> &Py { - &self.model - } - - #[pyo3(name="projection")] - pub fn project_state(&mut self, state: &PyState) -> PyObservation { - let state_data = state.inner.read().unwrap(); - let proj = self.projection(&state_data); - PyObservation { inner: proj } - } - - #[pyo3(name="inverse")] - pub fn project_observation(&mut self, observation: PyObservation) -> PyState { - let proj = self.inverse(&observation.inner); - PyState { inner: proj.ptr() } - } - - #[pyo3(name="jacobian")] - pub fn jacobian_matrix<'py>(&mut self, py: Python<'py>, state: &PyState) -> PyMatrix<'py> { - let state_data = state.inner.read().unwrap(); - let jac = self.jacobian(&state_data); - jac.as_slice() - .to_pyarray(py) - .reshape([jac.nrows(), jac.ncols()]) - .unwrap() - } -} - -impl MeasurementModel for PyMeasurementModel { - fn projection(&self, state: &State) -> Observation { - let py_state = PyState { inner: state.ptr() }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - - let result = py_obj - .call_method( - "projection", - (py_state,), - None - ) - .expect("Failed to call projection()"); - - let obs: PyObservation = result - .extract() - .expect("projection() must return "); - - obs.inner - }) - } - - fn inverse(&self, obs: &Observation) -> State { - - let py_obs = PyObservation { - inner: Observation { - value: obs.value.clone(), - epoch: obs.epoch, - } - }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - - let result = py_obj - .call_method( - "inverse", - (py_obs,), - None - ) - .expect("Failed to call inverse()"); - - let state: PyState = result - .extract() - .expect("inverse() must return "); - - state.inner.read().unwrap().clone() - }) - } - - fn jacobian(&self, state: &State) -> na::OMatrix { - let py_state = PyState{ inner: state.ptr() }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - let result = py_obj - .call_method( - "jacobian", - (py_state,), - None) - .expect("Failed to call jacobian()"); - - let result_vec: PyMatrix<'_> = result - .extract() - .expect("jacobian() must return "); - - // numpy -> ndarray (zero copy view) - let array = unsafe { result_vec.as_array() }; - let (nrows, ncols) = array.dim(); - - // ndarray -> nalgebra (single copy, row-major) - na::DMatrix::::from_iterator( - nrows, - ncols, - array.iter().copied() - ) - }) - } -} diff --git a/src/python/mod.rs b/src/python/mod.rs index 9b5a1fa..208a0dd 100644 --- a/src/python/mod.rs +++ b/src/python/mod.rs @@ -1,30 +1,51 @@ #![allow(unsafe_op_in_unsafe_fn)] +//! Python bindings. The modules mirror the crate: +//! +//! - `epoch`, `state`, `measurement`: `Epoch`, `State`, `Observation` +//! - `models`: Python-implemented `TransitionModel`, `MeasurementModel`, `UpdateModel` +//! - `algorithms`: `KalmanFilter`, `ExtendedKalmanFilter`, `UnscentedKalmanFilter`, +//! `UDKalmanFilter`, `UDExtendedKalmanFilter`, `CustomFilter` +//! +//! plus binding helpers: `arrays` (numpy <-> nalgebra), `callbacks` (Python +//! functions as Rust model closures, for the EKF and UKF), `filter` (method +//! bodies shared by the filter classes) and `logging` (spdlog -> Python logging). + +pub mod algorithms; pub mod epoch; -pub mod state; -pub mod transition; pub mod measurement; -pub mod update; -pub mod algorithms; +pub mod models; +pub mod state; + pub mod arrays; +pub mod callbacks; +pub mod filter; +pub mod logging; use pyo3::prelude::*; -///TODO: Update Model, KalmanFilter #[pymodule] fn determined(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { + // Route Rust log messages to Python's `logging` + logging::install()?; + // Structs m.add_class::()?; m.add_class::()?; m.add_class::()?; // Models - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + // Algorithms m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) -} \ No newline at end of file +} diff --git a/src/python/models.rs b/src/python/models.rs new file mode 100644 index 0000000..a24996b --- /dev/null +++ b/src/python/models.rs @@ -0,0 +1,357 @@ +#![allow(unsafe_op_in_unsafe_fn)] + +//! Python-implemented models: wrappers that let a Python object implement the +//! Rust model traits from `crate::models`, so it can be used with +//! `CustomFilter`. They mirror the Rust traits: +//! +//! | Wrapper | Rust traits | Python methods | +//! |--------------------|--------------------------------------------------------------|---------------------------------------------| +//! | `TransitionModel` | `TransitionModel`, `TransitionJacobian` | `state(epoch)`, `jacobian(state)` | +//! | `MeasurementModel` | `MeasurementModel`, `MeasurementJacobian`, `InverseMeasurement` | `projection(state)`, `noise()`, `jacobian(state)`, `inverse(observation)` (may return `None`) | +//! | `UpdateModel` | `UpdateModel` | `apply(observation)` | +//! +//! Each Python method is only called when the matching Rust method is. + +use numpy::PyArrayMethods; +use pyo3::prelude::*; +use pyo3::types::PyAny; + +use crate::common::na; +use crate::epoch::Epoch; +use crate::measurement::Observation; +use crate::models::{ + InverseMeasurement, MeasurementJacobian, MeasurementModel, TransitionJacobian, TransitionModel, + UpdateModel, +}; +use crate::state::{State, StatePtr}; + +use crate::python::arrays::{PyMatrix, to_dmatrix, to_pymatrix}; +use crate::python::epoch::PyEpoch; +use crate::python::measurement::PyObservation; +use crate::python::state::PyState; + +// Wrapper for Python-defined state transition - DYNAMIC only +#[pyclass(name = "TransitionModel")] +pub struct PyTransitionModel { + pub model: Py, + pub state: PyState, +} + +impl Clone for PyTransitionModel { + fn clone(&self) -> Self { + Python::attach(|py| { + let _ref = self.model.clone_ref(py); + Self { + state: self.state.clone(), + model: _ref, + } + }) + } +} + +#[pymethods] +impl PyTransitionModel { + #[new] + #[pyo3(signature = (model, *, state))] + pub fn new(model: Py, state: PyState) -> Self { + PyTransitionModel { model, state } + } + + #[getter] + fn get_model(&self) -> &Py { + &self.model + } + + #[pyo3(name = "state")] + fn state_transition(&mut self, epoch: &PyEpoch) -> PyState { + let state = self.state(&epoch.inner).clone(); + PyState { inner: state } + } + + #[pyo3(name = "jacobian")] + fn jacobian_matrix<'py>(&mut self, py: Python<'py>, state: &'py PyState) -> PyMatrix<'py> { + let state_data = state.inner.read().unwrap(); + let jac = self.jacobian(&state_data); + to_pymatrix(py, &jac) + } +} + +impl TransitionModel for PyTransitionModel { + fn state(&mut self, epoch: &Epoch) -> &StatePtr { + let py_epoch = PyEpoch { inner: *epoch }; + + Python::attach(|py| { + let py_obj = self.model.bind(py); + + // Call Python method with state and epoch + let result = py_obj + .call_method("state", (py_epoch,), None) + .expect("Failed to call state()"); + + let result_state: PyState = result.extract().expect("state() must return "); + + // copy out first: the returned state may share this state's lock + let _state = result_state.inner.read().unwrap().clone(); + + { + // lock and update state + let mut state = self.state.inner.write().unwrap(); + state.value = _state.value; + state.covariance = _state.covariance; + state.epoch = *epoch; + } + }); + &self.state.inner + } +} + +impl TransitionJacobian for PyTransitionModel { + fn jacobian(&self, state: &State) -> na::OMatrix { + let py_state = PyState { inner: state.ptr() }; + + Python::attach(|py| { + let py_obj = self.model.bind(py); + let result = py_obj + .call_method("jacobian", (py_state,), None) + .expect("Failed to call jacobian()"); + + let result_vec: PyMatrix<'_> = + result.extract().expect("jacobian() must return ndarray"); + + // numpy -> ndarray (zero copy view) -> nalgebra (single copy) + to_dmatrix(unsafe { result_vec.as_array() }) + }) + } +} + +// Wrapper for Python-defined Measurement models +#[pyclass(name = "MeasurementModel")] +pub struct PyMeasurementModel { + pub model: Py, +} + +impl Clone for PyMeasurementModel { + fn clone(&self) -> Self { + Python::attach(|py| { + let _ref = self.model.clone_ref(py); + Self { model: _ref } + }) + } +} + +#[pymethods] +impl PyMeasurementModel { + #[new] + pub fn new(model: Py) -> Self { + PyMeasurementModel { model } + } + + #[getter] + fn get_model(&self) -> &Py { + &self.model + } + + #[pyo3(name = "projection")] + pub fn project_state(&mut self, state: &PyState) -> PyObservation { + let state_data = state.inner.read().unwrap(); + let proj = self.projection(&state_data); + PyObservation { inner: proj } + } + + #[pyo3(name = "inverse")] + pub fn project_observation(&mut self, observation: PyObservation) -> Option { + self.inverse(&observation.inner) + .map(|proj| PyState { inner: proj.ptr() }) + } + + #[pyo3(name = "jacobian")] + pub fn jacobian_matrix<'py>(&mut self, py: Python<'py>, state: &PyState) -> PyMatrix<'py> { + let state_data = state.inner.read().unwrap(); + let jac = self.jacobian(&state_data); + to_pymatrix(py, &jac) + } + + #[pyo3(name = "noise")] + pub fn noise_matrix<'py>(&mut self, py: Python<'py>) -> PyMatrix<'py> { + to_pymatrix(py, &self.noise()) + } +} + +impl MeasurementModel for PyMeasurementModel { + fn projection(&self, state: &State) -> Observation { + let py_state = PyState { inner: state.ptr() }; + + Python::attach(|py| { + let py_obj = self.model.bind(py); + + let result = py_obj + .call_method("projection", (py_state,), None) + .expect("Failed to call projection()"); + + let obs: PyObservation = result + .extract() + .expect("projection() must return "); + + obs.inner + }) + } + + fn noise(&self) -> na::OMatrix { + Python::attach(|py| { + let py_obj = self.model.bind(py); + let result = py_obj + .call_method0("noise") + .expect("Failed to call noise()"); + + let result_mat: PyMatrix<'_> = result.extract().expect("noise() must return ndarray"); + + // numpy -> ndarray (zero copy view) -> nalgebra (single copy) + to_dmatrix(unsafe { result_mat.as_array() }) + }) + } +} + +impl MeasurementJacobian for PyMeasurementModel { + fn jacobian(&self, state: &State) -> na::OMatrix { + let py_state = PyState { inner: state.ptr() }; + + Python::attach(|py| { + let py_obj = self.model.bind(py); + let result = py_obj + .call_method("jacobian", (py_state,), None) + .expect("Failed to call jacobian()"); + + let result_vec: PyMatrix<'_> = + result.extract().expect("jacobian() must return "); + + // numpy -> ndarray (zero copy view) -> nalgebra (single copy) + to_dmatrix(unsafe { result_vec.as_array() }) + }) + } +} + +impl InverseMeasurement for PyMeasurementModel { + fn inverse(&self, obs: &Observation) -> Option> { + let py_obs = PyObservation { + inner: Observation { + value: obs.value.clone(), + epoch: obs.epoch, + }, + }; + + Python::attach(|py| { + let py_obj = self.model.bind(py); + + let result = py_obj + .call_method("inverse", (py_obs,), None) + .expect("Failed to call inverse()"); + + let state: Option = result + .extract() + .expect("inverse() must return or None"); + + state.map(|state| state.inner.read().unwrap().clone()) + }) + } +} + +// Wrapper for Python-defined update models +#[pyclass(name = "UpdateModel")] +pub struct PyUpdateModel { + pub model: Py, + pub state: PyState, + pub transition: PyTransitionModel, + pub measurement: PyMeasurementModel, +} + +impl Clone for PyUpdateModel { + fn clone(&self) -> Self { + Python::attach(|py| { + let _ref = self.model.clone_ref(py); + Self { + model: _ref, + state: self.state.clone(), + transition: self.transition.clone(), + measurement: self.measurement.clone(), + } + }) + } +} + +#[pymethods] +impl PyUpdateModel { + #[new] + #[pyo3(signature = (model, *, transition, measurement))] + pub fn new( + model: Py, + transition: Py, + measurement: Py, + ) -> Self { + Python::attach(|py| { + let _transition: PyTransitionModel = transition.extract(py).unwrap(); + let _measurement: PyMeasurementModel = measurement.extract(py).unwrap(); + + PyUpdateModel { + model, + state: _transition.state.clone(), + transition: _transition, + measurement: _measurement, + } + }) + } + + #[getter] + fn get_model(&self) -> &Py { + &self.model + } + + #[pyo3(name = "state")] + fn state_transition(&mut self, epoch: &PyEpoch) -> PyState { + let state = self.state(Some(&epoch.inner)); + PyState { + inner: state.clone(), + } + } + + #[pyo3(name = "apply")] + fn apply_update(&mut self, observation: &PyObservation) -> PyState { + let obs = &observation.inner; + let state = self.apply(obs).clone(); + PyState { inner: state } + } +} + +impl UpdateModel for PyUpdateModel { + fn state(&mut self, epoch: Option<&Epoch>) -> &StatePtr { + match epoch { + Some(epoch) => self.transition.state(epoch), + _ => &self.transition.state.inner, + } + } + + fn apply(&mut self, observation: &Observation) -> &StatePtr { + let py_obs = PyObservation { + inner: observation.clone(), + }; + + Python::attach(|py| { + let py_obj = self.model.bind(py); + + // Call Python method with state and epoch + let result = py_obj + .call_method("apply", (py_obs,), None) + .expect("Failed to call apply()"); + + let result_state: PyState = result.extract().expect("apply() must return "); + + // copy out first: the returned state may share this state's lock + let _state = result_state.inner.read().unwrap().clone(); + + { + // lock and update state + *self.state.inner.write().unwrap() = _state; + } + }); + &self.state.inner + } +} diff --git a/src/python/state.rs b/src/python/state.rs index 626ad12..ad1bf6a 100644 --- a/src/python/state.rs +++ b/src/python/state.rs @@ -1,16 +1,15 @@ #![allow(unsafe_op_in_unsafe_fn)] +use numpy::ToPyArray; +use numpy::{PyReadonlyArray1, PyReadonlyArray2}; use pyo3::prelude::*; -use numpy::{ PyReadonlyArray1, PyReadonlyArray2 }; -use numpy::{ ToPyArray, PyArrayMethods }; -use crate::common::na as na; +use crate::common::na; +use crate::python::arrays::{PyMatrix, PyVector, check_square, to_dmatrix, to_pymatrix}; use crate::python::epoch::PyEpoch; use crate::state::{State, StatePtr}; -use crate::python::arrays::{ PyVector, PyMatrix }; - -#[pyclass(name="State")] +#[pyclass(name = "State")] #[derive(Clone, Debug)] pub struct PyState { pub inner: StatePtr, @@ -18,40 +17,33 @@ pub struct PyState { #[pymethods] impl PyState { - #[new] + #[pyo3(signature = (value, *, covariance, epoch))] fn new( value: PyReadonlyArray1, covariance: Option>, - epoch: &PyEpoch + epoch: &PyEpoch, ) -> PyResult { - - // numpy -> ndarray (zero-copy view) -> nalgebra - let array = value.as_array(); - let vector = na::DVector::::from_iterator( - array.len(), - array.iter().copied() - ); - - // Same for covariance - let cov = covariance.map(|cov| { - let nd_cov = cov.as_array(); - let (nrows, ncols) = nd_cov.dim(); - - // ndarray -> nalgebra (single copy, row-major) - na::DMatrix::::from_iterator( - nrows, - ncols, - nd_cov.iter().copied() - ) - }); - + // numpy -> ndarray (zero-copy view) -> nalgebra + let array = value.as_array(); + let vector = na::DVector::::from_iterator(array.len(), array.iter().copied()); + + // Same for covariance + let cov = covariance.map(|cov| { + // ndarray -> nalgebra (single copy) + to_dmatrix(cov.as_array()) + }); + if let Some(cov) = &cov { + check_square(cov, vector.len(), "covariance")?; + } + Ok(PyState { inner: State { value: vector, covariance: cov, epoch: epoch.inner, - }.ptr() + } + .ptr(), }) } @@ -71,14 +63,28 @@ impl PyState { #[getter] fn covariance<'py>(&self, py: Python<'py>) -> Option> { let state = self.inner.read().unwrap(); - state.covariance.as_ref().map(|cov| { - cov.as_slice().to_pyarray(py) - .reshape([cov.nrows(), cov.ncols()]) - .unwrap() - }) + state.covariance.as_ref().map(|cov| to_pymatrix(py, cov)) } fn __repr__(&self) -> PyResult { Ok(format!("{:0.6}", self.inner.read().unwrap())) } -} \ No newline at end of file +} + +impl PyState { + /// Copy `other` into this state's storage, keeping this state's identity, + /// so everything sharing it (e.g. a filter) sees the new values. + pub fn assign(&self, other: &PyState) -> PyResult<()> { + // copy out first: `other` may share this state's lock + let incoming = other.inner.read().unwrap().clone(); + let dim = self.inner.read().unwrap().dim(); + if incoming.dim() != dim { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "state has dimension {}, expected {dim}", + incoming.dim() + ))); + } + *self.inner.write().unwrap() = incoming; + Ok(()) + } +} diff --git a/src/python/transition.rs b/src/python/transition.rs deleted file mode 100644 index d8068bd..0000000 --- a/src/python/transition.rs +++ /dev/null @@ -1,126 +0,0 @@ -#![allow(unsafe_op_in_unsafe_fn)] - -use pyo3::prelude::*; -use pyo3::types::PyAny; -use numpy::{ ToPyArray, PyArrayMethods }; - -use crate::common::na as na; -use crate::epoch::Epoch; -use crate::models::TransitionModel; -use crate::state::{State, StatePtr}; - -use crate::python::epoch::PyEpoch; -use crate::python::state::PyState; -use crate::python::arrays::{ PyMatrix }; - - -// Wrapper for Python-defined state transition - DYNAMIC only -#[pyclass(name="TransitionModel")] -pub struct PyTransitionModel { - pub model: Py, - pub state: PyState, -} - -impl Clone for PyTransitionModel { - fn clone(&self) -> Self { - Python::attach(| py| { - let _ref = self.model.clone_ref(py); - Self { - state: self.state.clone(), - model: _ref, - } - }) - } -} - -#[pymethods] -impl PyTransitionModel { - - #[new] - pub fn new(model: Py, state: PyState) -> Self { - PyTransitionModel { model: model, state } - } - - #[getter] - fn get_model(&self) -> &Py { - &self.model - } - - #[pyo3(name="state")] - fn state_transition(&mut self, epoch: &PyEpoch) -> PyState { - let state = self.state(&epoch.inner).clone(); - PyState { inner: state } - } - - #[pyo3(name="jacobian")] - fn jacobian_matrix<'py>(&mut self, py: Python<'py>, state: &'py PyState) -> PyMatrix<'py> { - let state_data = state.inner.read().unwrap(); - let jac = self.jacobian(&state_data); - jac.as_slice() - .to_pyarray(py) - .reshape([jac.nrows(), jac.ncols()]) - .unwrap() - } - -} - - -impl TransitionModel for PyTransitionModel { - fn state(&mut self, epoch: &Epoch) -> &StatePtr { - - let py_epoch = PyEpoch{ inner: *epoch }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - - // Call Python method with state and epoch - let result = py_obj - .call_method("state", (py_epoch,), None) - .expect("Failed to call state()"); - - let result_state: PyState = result - .extract() - .expect("state() must return "); - - let _state = result_state.inner.read().unwrap(); - - { // lock and update state - let mut state = self.state.inner.write().unwrap(); - state.value = _state.value.clone(); - state.covariance = _state.covariance.clone(); - state.epoch = *epoch; - } - - }); - &self.state.inner - } - - fn jacobian(&self, state: &State) -> na::OMatrix { - let py_state = PyState{ inner: state.ptr() }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - let result = py_obj - .call_method( - "jacobian", - (py_state,), - None) - .expect("Failed to call jacobian()"); - - let result_vec: PyMatrix<'_> = result - .extract() - .expect("jacobian() must return ndarray"); - - // numpy -> ndarray (zero copy view) - let array = unsafe { result_vec.as_array() }; - let (nrows, ncols) = array.dim(); - - // ndarray -> nalgebra (single copy, row-major) - na::DMatrix::::from_iterator( - nrows, - ncols, - array.iter().copied() - ) - }) - } -} \ No newline at end of file diff --git a/src/python/update.rs b/src/python/update.rs deleted file mode 100644 index e04e6d0..0000000 --- a/src/python/update.rs +++ /dev/null @@ -1,158 +0,0 @@ -#![allow(unsafe_op_in_unsafe_fn)] - -use pyo3::prelude::*; -use pyo3::types::PyAny; -use numpy::{ ToPyArray, PyArrayMethods }; - -use crate::common::na; -use crate::measurement::Observation; -use crate::models::{TransitionModel, UpdateModel}; -use crate::state::{ State, StatePtr }; - -use crate::python::epoch::PyEpoch; -use crate::python::state::PyState; -use crate::python::transition::PyTransitionModel; -use crate::python::measurement::{PyMeasurementModel, PyObservation}; -use crate::python::arrays::{ PyMatrix }; - - -#[pyclass(name="UpdateModel")] -pub struct PyUpdateModel { - pub model: Py, - pub state: PyState, - pub transition: PyTransitionModel, - pub measurement: PyMeasurementModel, -} - -impl Clone for PyUpdateModel { - fn clone(&self) -> Self { - Python::attach(| py| { - let _ref = self.model.clone_ref(py); - Self { - model: _ref, - state: self.state.clone(), - transition: self.transition.clone(), - measurement: self.measurement.clone(), - } - }) - } -} - -#[pymethods] -impl PyUpdateModel { - #[new] - pub fn new( - model: Py, - transition: Py, - measurement: Py) -> Self { - - Python::attach(|py| { - - let _transition: PyTransitionModel = transition.extract(py).unwrap(); - let _measurement: PyMeasurementModel = measurement.extract(py).unwrap(); - - PyUpdateModel { - model: model, - state: _transition.state.clone(), - transition: _transition, - measurement: _measurement, - } - }) - } - - #[getter] - fn get_model(&self) -> &Py { - &self.model - } - - #[pyo3(name="state")] - fn state_transition(&mut self, epoch: &PyEpoch) -> PyState { - let state = self.state(&epoch.inner); - PyState { inner: state.clone() } - } - - #[pyo3(name="apply")] - fn apply_update(&mut self, observation: &PyObservation) -> PyState { - let obs = &observation.inner; - let state = self.apply(obs).clone(); - PyState { inner: state } - } - - #[pyo3(name="jacobian")] - fn jacobian_matrix<'py>(&mut self, py: Python<'py>, state: &'py PyState) -> PyMatrix<'py> { - let state_data = state.inner.read().unwrap(); - let jac = self.jacobian(&state_data); - jac.as_slice() - .to_pyarray(py) - .reshape([jac.nrows(), jac.ncols()]) - .unwrap() - } -} - -impl UpdateModel for PyUpdateModel { - - fn state(&mut self, epoch: &crate::epoch::Epoch) -> &StatePtr { - self.transition.state(epoch) - } - - fn apply(&mut self, observation: &Observation) -> &StatePtr { - - let py_obs = PyObservation{ inner: observation.clone() }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - - // Call Python method with state and epoch - let result = py_obj - .call_method( - "apply", - (py_obs,), - None) - .expect("Failed to call apply()"); - - let result_state: PyState = result - .extract() - .expect("apply() must return "); - - let _state = result_state.inner.read().unwrap(); - - { // lock and update state - let mut state = self.state.inner.write().unwrap(); - state.value = _state.value.clone(); - state.covariance = _state.covariance.clone(); - state.epoch = _state.epoch; - } - - }); - &self.state.inner - } - - fn jacobian(&self, state: &State) -> na::OMatrix { - let py_state = PyState{ inner: state.ptr() }; - - Python::attach(|py| { - let py_obj = self.model.bind(py); - let result = py_obj - .call_method( - "jacobian", - (py_state,), - None) - .expect("Failed to call jacobian()"); - - let result_vec: PyMatrix<'_> = result - .extract() - .expect("jacobian() must return ndarray"); - - // numpy -> ndarray (zero copy view) - let array = unsafe { result_vec.as_array() }; - let (nrows, ncols) = array.dim(); - - // ndarray -> nalgebra (single copy, row-major) - na::DMatrix::::from_iterator( - nrows, - ncols, - array.iter().copied() - ) - }) - } -} \ No newline at end of file diff --git a/src/state.rs b/src/state.rs index cc5e316..62102d1 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,133 +1,127 @@ -//! Small typed wrapper around nalgebra static matrices representing a state or -//! observation. `State` is a thin container that carries a `SMatrix` and -//! a simple `epoch` value used for timestamping. - -use std::sync::Arc; -use std::sync::RwLock; -use std::fmt; - -use crate::common::{Composable, na}; -use crate::epoch::Epoch; - -#[derive(Clone, Debug)] -pub struct State -where - T: na::Dim, - na::DefaultAllocator: na::base::allocator::Allocator - + na::allocator::Allocator, -{ - pub value: na::Matrix>, - pub covariance: Option>>, - pub epoch: Epoch, -} - -pub type StateN = State>; -pub type StateDyn = State; -pub type StatePtr = Arc>>; - -impl State -where - T: na::Dim, - na::DefaultAllocator: na::allocator::Allocator - + na::allocator::Allocator, -{ - pub fn new(value: Vec, covariance: Vec>, epoch: Epoch) -> Self { - let dim = value.len(); - State { - value: na::Matrix::::from_row_slice_generic( - T::from_usize(dim), - na::U1, - &value, - ), - covariance: Some(na::Matrix::::from_row_slice_generic( - T::from_usize(dim), - T::from_usize(dim), - &covariance.concat(), - )), - epoch, - } - } - - pub fn zeros(size: usize, epoch: Epoch) -> Self { - State { - value: na::Matrix::::zeros_generic( - T::from_usize(size), - na::U1, - ), - covariance: Some(na::Matrix::::zeros_generic( - T::from_usize(size), - T::from_usize(size), - )), - epoch, - } - } - - pub fn dim(&self) -> usize { - self.value.nrows() - } - - pub fn covariance(&self) -> na::Matrix> { - // TODO: stop this covariance from cloning - self.covariance.clone().unwrap_or_else(|| { - na::Matrix::>::zeros_generic( - T::from_usize(self.dim()), - T::from_usize(self.dim()), - ) - }) - } - - pub fn ptr(&self) -> StatePtr { - Arc::new(RwLock::new(self.clone())) - } -} - -impl Default for State -where - T: na::Dim + na::DimName, - na::DefaultAllocator: na::allocator::Allocator - + na::allocator::Allocator, -{ - fn default() -> Self { - Self { - value: na::Matrix::::zeros(), - covariance: Some(na::Matrix::::identity()), - epoch: Epoch::default(), - } - } -} - -impl Composable for State -where - T: na::Dim, - na::DefaultAllocator: na::allocator::Allocator - + na::allocator::Allocator, -{ - type Output = State; - - fn add(self, other: Self) -> Self { - State { - value: &self.value + &other.value, - covariance: Some(self.covariance() + &other.covariance()), - epoch: self.epoch, - } - } -} - -impl fmt::Display for State -where - T: na::Dim, - na::DefaultAllocator: na::base::allocator::Allocator - + na::allocator::Allocator, -{ - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Format `value`, `covariance`, and `epoch` as desired - writeln!(f, "State:")?; - writeln!(f, "‾‾‾‾‾‾")?; - write!(f, " Value: {:0.6}", self.value)?; // or custom formatting - let _ = match &self.covariance { - Some(cov) => write!(f, " Covariance: {:0.6}", cov), - None => write!(f, " Covariance: None"), - }; - writeln!(f, " Epoch: {}", self.epoch.value) - } -} \ No newline at end of file +//! Small typed wrapper around nalgebra static matrices representing a state or +//! observation. `State` is a thin container that carries a `SMatrix` and +//! a simple `epoch` value used for timestamping. + +use std::fmt; +use std::sync::Arc; +use std::sync::RwLock; + +use crate::common::{Composable, na}; +use crate::epoch::Epoch; + +#[derive(Clone, Debug)] +pub struct State +where + T: na::Dim, + na::DefaultAllocator: + na::base::allocator::Allocator + na::allocator::Allocator, +{ + pub value: na::OVector, + pub covariance: Option>, + pub epoch: Epoch, +} + +pub type StateN = State>; +pub type StateDyn = State; +pub type StatePtr = Arc>>; + +impl State +where + T: na::Dim, + na::DefaultAllocator: na::allocator::Allocator + na::allocator::Allocator, +{ + pub fn new(value: Vec, covariance: Vec>, epoch: Epoch) -> Self { + let dim = value.len(); + State { + value: na::Matrix::::from_row_slice_generic( + T::from_usize(dim), + na::U1, + &value, + ), + covariance: Some(na::Matrix::::from_row_slice_generic( + T::from_usize(dim), + T::from_usize(dim), + &covariance.concat(), + )), + epoch, + } + } + + pub fn zeros(size: usize, epoch: Epoch) -> Self { + State { + value: na::Matrix::::zeros_generic(T::from_usize(size), na::U1), + covariance: Some(na::Matrix::::zeros_generic( + T::from_usize(size), + T::from_usize(size), + )), + epoch, + } + } + + pub fn dim(&self) -> usize { + self.value.nrows() + } + + pub fn covariance(&self) -> na::Matrix> { + // TODO: stop this covariance from cloning + self.covariance.clone().unwrap_or_else(|| { + na::Matrix::>::zeros_generic( + T::from_usize(self.dim()), + T::from_usize(self.dim()), + ) + }) + } + + pub fn ptr(&self) -> StatePtr { + Arc::new(RwLock::new(self.clone())) + } +} + +impl Default for State +where + T: na::Dim + na::DimName, + na::DefaultAllocator: na::allocator::Allocator + na::allocator::Allocator, +{ + fn default() -> Self { + Self { + value: na::Matrix::::zeros(), + covariance: Some(na::Matrix::::identity()), + epoch: Epoch::default(), + } + } +} + +impl Composable for State +where + T: na::Dim, + na::DefaultAllocator: na::allocator::Allocator + na::allocator::Allocator, +{ + type Output = State; + + fn add(self, other: Self) -> Self { + State { + value: &self.value + &other.value, + covariance: Some(self.covariance() + &other.covariance()), + epoch: self.epoch, + } + } +} + +impl fmt::Display for State +where + T: na::Dim, + na::DefaultAllocator: + na::base::allocator::Allocator + na::allocator::Allocator, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // Format `value`, `covariance`, and `epoch` as desired + writeln!(f, "State:")?; + writeln!(f, "‾‾‾‾‾‾")?; + write!(f, " Value: {:0.6}", self.value)?; // or custom formatting + let _ = match &self.covariance { + Some(cov) => write!(f, " Covariance: {:0.6}", cov), + None => write!(f, " Covariance: None"), + }; + writeln!(f, " Epoch: {}", self.epoch.value) + } +} diff --git a/src/tests/logs.rs b/src/tests/logs.rs new file mode 100644 index 0000000..cad978d --- /dev/null +++ b/src/tests/logs.rs @@ -0,0 +1,78 @@ +//! Keeps logs that a test triggers on purpose out of the test output. +//! +//! Some tests give a filter invalid input to check that it refuses cleanly, +//! which logs a warning or error. `expect_logs()` drops log output from the +//! calling test's thread while the returned guard is alive. Logs from every +//! other test still print, so unexpected warnings stay visible. + +use std::cell::Cell; +use std::sync::{Arc, Once}; + +use spdlog::formatter::Formatter; +use spdlog::sink::Sink; +use spdlog::{ErrorHandler, LevelFilter, Logger, Record}; + +thread_local! { + static QUIET: Cell = const { Cell::new(false) }; +} + +/// Forwards records to the original default logger unless the logging +/// thread is quiet. +struct QuietableSink { + inner: Arc, +} + +impl Sink for QuietableSink { + fn log(&self, record: &Record) -> spdlog::Result<()> { + if !QUIET.with(Cell::get) { + self.inner.log(record); + } + Ok(()) + } + + fn flush(&self) -> spdlog::Result<()> { + self.inner.flush(); + Ok(()) + } + + fn level_filter(&self) -> LevelFilter { + LevelFilter::All + } + + fn set_level_filter(&self, _level_filter: LevelFilter) {} + + fn set_formatter(&self, _formatter: Box) {} + + fn set_error_handler(&self, _handler: Option) {} +} + +fn install() { + static INSTALL: Once = Once::new(); + INSTALL.call_once(|| { + let inner = spdlog::default_logger(); + let logger = Logger::builder() + .level_filter(LevelFilter::All) + .sink(Arc::new(QuietableSink { inner })) + .build() + .expect("test logger"); + spdlog::set_default_logger(Arc::new(logger)); + }); +} + +/// Silences log output from this test while alive. +pub struct ExpectedLogs; + +impl Drop for ExpectedLogs { + fn drop(&mut self) { + QUIET.with(|quiet| quiet.set(false)); + } +} + +/// For tests that trigger warnings or errors on purpose: hold the guard for +/// the rest of the test, e.g. `let _quiet = expect_logs();`. +#[must_use = "logs are only silenced while the guard is alive"] +pub fn expect_logs() -> ExpectedLogs { + install(); + QUIET.with(|quiet| quiet.set(true)); + ExpectedLogs +} diff --git a/src/tests/mod.rs b/src/tests/mod.rs index af9628b..f341029 100644 --- a/src/tests/mod.rs +++ b/src/tests/mod.rs @@ -1,2 +1,9 @@ +mod logs; +pub use logs::expect_logs; + +pub mod test_dyn; +pub mod test_ekf; +pub mod test_kalman; pub mod test_state; -pub mod test_kalman_filter; \ No newline at end of file +pub mod test_ud; +pub mod test_ukf; diff --git a/src/tests/test_dyn.rs b/src/tests/test_dyn.rs new file mode 100644 index 0000000..856dd97 --- /dev/null +++ b/src/tests/test_dyn.rs @@ -0,0 +1,146 @@ +use crate::algorithms::{ + ExtendedKalmanFilterDyn, ExtendedKalmanFilterNM, KalmanFilterDyn, KalmanFilterNM, +}; +use crate::common::na; +use crate::epoch::Epoch; +use crate::filter::Filter; +use crate::measurement::Observation; +use crate::models::*; +use crate::state::State; + +const F: [f64; 4] = [1.0, 0.5, 0.0, 1.0]; +const H: [f64; 2] = [1.0, 0.0]; + +fn obs_at(t: i64) -> (f64, Epoch) { + (0.4 * t as f64, Epoch::new(t)) +} + +/// The same models with runtime sizes (`na::Dyn`, as used by Python) must +/// give exactly the statically sized result. +#[test] +fn test_dyn_matches_static_kalman() { + let f = na::SMatrix::::from_row_slice(&F); + let h = na::SMatrix::::from_row_slice(&H); + let q = na::SMatrix::::identity() * 0.05; + let r = na::SMatrix::::new(0.1); + + let state = || { + State::>::new( + vec![0.0, 1.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) + }; + let dyn_state = || { + State::::new( + vec![0.0, 1.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) + }; + + let mut fixed = KalmanFilterNM::new(LinearUpdateNM::<2, 1>::new( + LinearMeasurementNM::<2, 1>::new(h, r), + LinearTransitionN::<2>::new(state().ptr(), f, q), + )); + let mut dynamic = KalmanFilterDyn::new(LinearUpdate::new( + LinearMeasurement::new( + na::DMatrix::from_row_slice(1, 2, &H), + na::DMatrix::from_element(1, 1, 0.1), + ), + LinearTransition::new( + dyn_state().ptr(), + na::DMatrix::from_row_slice(2, 2, &F), + na::DMatrix::identity(2, 2) * 0.05, + ), + )); + + for t in [1, 2, 5, 6] { + let (z, epoch) = obs_at(t); + let a = fixed.update(&Observation { + value: na::SVector::::new(z), + epoch, + }); + let b = dynamic.update(&Observation { + value: na::DVector::from_element(1, z), + epoch, + }); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + + assert_eq!(a.value.as_slice(), b.value.as_slice()); + assert_eq!(a.covariance().as_slice(), b.covariance().as_slice()); + assert_eq!(a.epoch, b.epoch); + } +} + +#[test] +fn test_dyn_matches_static_extended_kalman() { + let range = |x: &[f64]| (x[0] * x[0] + x[1] * x[1]).sqrt(); + + let mut fixed = ExtendedKalmanFilterNM::new(NonLinearUpdateNM::<2, 1>::new( + NonLinearMeasurementNM::<2, 1>::new( + Box::new(move |s| na::SVector::::new(range(s.value.as_slice()))), + Box::new(move |s| { + na::SMatrix::::new( + s.value[0] / range(s.value.as_slice()), + s.value[1] / range(s.value.as_slice()), + ) + }), + na::SMatrix::::new(0.01), + ), + NonLinearTransitionN::<2>::new( + State::>::new( + vec![2.0, 2.0], + vec![vec![4.0, 0.0], vec![0.0, 4.0]], + Epoch::new(0), + ) + .ptr(), + Box::new(|s| s.value * 1.01), + Box::new(|_| na::SMatrix::::identity() * 1.01), + na::SMatrix::::identity() * 0.01, + ), + )); + let mut dynamic = ExtendedKalmanFilterDyn::new(NonLinearUpdate::new( + NonLinearMeasurement::new( + Box::new(move |s| na::DVector::from_element(1, range(s.value.as_slice()))), + Box::new(move |s| { + na::DMatrix::from_row_slice( + 1, + 2, + &[ + s.value[0] / range(s.value.as_slice()), + s.value[1] / range(s.value.as_slice()), + ], + ) + }), + na::DMatrix::from_element(1, 1, 0.01), + ), + NonLinearTransition::new( + State::::new( + vec![2.0, 2.0], + vec![vec![4.0, 0.0], vec![0.0, 4.0]], + Epoch::new(0), + ) + .ptr(), + Box::new(|s| &s.value * 1.01), + Box::new(|_| na::DMatrix::identity(2, 2) * 1.01), + na::DMatrix::identity(2, 2) * 0.01, + ), + )); + + for t in 1..=4 { + let epoch = Epoch::new(t); + let a = fixed.update(&Observation { + value: na::SVector::::new(5.0), + epoch, + }); + let b = dynamic.update(&Observation { + value: na::DVector::from_element(1, 5.0), + epoch, + }); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + + assert_eq!(a.value.as_slice(), b.value.as_slice()); + assert_eq!(a.covariance().as_slice(), b.covariance().as_slice()); + } +} diff --git a/src/tests/test_ekf.rs b/src/tests/test_ekf.rs new file mode 100644 index 0000000..7dcb5d0 --- /dev/null +++ b/src/tests/test_ekf.rs @@ -0,0 +1,150 @@ +use crate::algorithms::{ExtendedKalmanFilterNM, KalmanFilterNM}; +use crate::common::na; +use crate::epoch::Epoch; +use crate::filter::Filter; +use crate::measurement::Observation; +use crate::models::*; +use crate::state::State; + +fn initial_state() -> State> { + State::>::new( + vec![1.0, 1.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) +} + +/// A linear problem run through the EKF must match the linear Kalman filter. +#[test] +fn test_extended_kalman_matches_linear() { + let f = na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[1.5, 0.5])); + let q = &na::SMatrix::::identity() * 0.05; + let h = na::SMatrix::::identity(); + let r = &na::SMatrix::::identity() * 0.01; + + // Linear filter + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<2, 2>::new( + LinearMeasurementNM::<2, 2>::new(h, r), + LinearTransitionN::<2>::new(initial_state().ptr(), f, q), + )); + + // Extended filter + let mut ekf = ExtendedKalmanFilterNM::new(NonLinearUpdateNM::<2, 2>::new( + NonLinearMeasurementNM::<2, 2>::linear(h, r), + NonLinearTransitionN::<2>::new( + initial_state().ptr(), + Box::new(move |state| f * state.value), + Box::new(move |_| f), + q, + ), + )); + + for t in 1..=3 { + let obs = Observation { + value: na::SMatrix::::from_row_slice(&[1.0 + t as f64, 0.5]), + epoch: Epoch::new(t), + }; + + kf.predict(&obs.epoch); + ekf.predict(&obs.epoch); + let kf_state = kf.update(&obs); + let ekf_state = ekf.update(&obs); + + let kf_state = kf_state.read().unwrap(); + let ekf_state = ekf_state.read().unwrap(); + + let dx = (kf_state.value - ekf_state.value).norm(); + let dp = (kf_state.covariance() - ekf_state.covariance()).norm(); + assert!(dx < 1e-12, "State mismatch at t={}: {}", t, dx); + assert!(dp < 1e-12, "Covariance mismatch at t={}: {}", t, dp); + assert_eq!(ekf_state.epoch, Epoch::new(t)); + } +} + +/// Multi-step prediction applies `f` once per epoch, and is idempotent. +#[test] +fn test_extended_kalman_predict_steps() { + let mut ekf = ExtendedKalmanFilterNM::new(NonLinearUpdateNM::<2, 2>::new( + NonLinearMeasurementNM::<2, 2>::linear( + na::SMatrix::::identity(), + na::SMatrix::::identity(), + ), + NonLinearTransitionN::<2>::new( + initial_state().ptr(), + Box::new(|state| state.value * 2.0), + Box::new(|_| na::SMatrix::::identity() * 2.0), + na::SMatrix::::zeros(), + ), + )); + + let state = ekf.predict(&Epoch::new(3)); + assert_eq!( + state.read().unwrap().value, + na::SVector::::new(8.0, 8.0) + ); + assert_eq!( + state.read().unwrap().covariance(), + na::SMatrix::::identity() * 64.0 + ); + + // predicting to the same epoch again must not propagate + let state = ekf.predict(&Epoch::new(3)); + assert_eq!( + state.read().unwrap().value, + na::SVector::::new(8.0, 8.0) + ); +} + +/// A genuinely nonlinear (range) sensor pulls the estimate towards the truth +/// and shrinks the uncertainty. +#[test] +fn test_extended_kalman_range_sensor() { + let truth_range = 5.0; // true position is (3, 4) + + let measurement = NonLinearMeasurementNM::<2, 1>::new( + Box::new(|state| na::SMatrix::::new(state.value.norm())), + Box::new(|state| { + let rng = state.value.norm(); + na::SMatrix::::new(state.value[0] / rng, state.value[1] / rng) + }), + na::SMatrix::::new(0.01), + ); + + // static target + let transition = NonLinearTransitionN::<2>::new( + State::>::new( + vec![2.0, 2.0], + vec![vec![4.0, 0.0], vec![0.0, 4.0]], + Epoch::new(0), + ) + .ptr(), + Box::new(|state| state.value), + Box::new(|_| na::SMatrix::::identity()), + na::SMatrix::::zeros(), + ); + + let mut ekf = + ExtendedKalmanFilterNM::new(NonLinearUpdateNM::<2, 1>::new(measurement, transition)); + + let mut err = (2.0f64 * 2.0 + 2.0 * 2.0).sqrt() - truth_range; + let mut trace = 8.0; + for t in 1..=5 { + let state = ekf.update(&Observation { + value: na::SMatrix::::new(truth_range), + epoch: Epoch::new(t), + }); + let state = state.read().unwrap(); + + let new_err = state.value.norm() - truth_range; + let new_trace = state.covariance().trace(); + assert!( + new_err.abs() < err.abs(), + "Range error did not shrink at t={}", + t + ); + assert!(new_trace < trace, "Uncertainty did not shrink at t={}", t); + err = new_err; + trace = new_trace; + } + assert!(err.abs() < 1e-2, "Range error too large: {}", err); +} diff --git a/src/tests/test_kalman.rs b/src/tests/test_kalman.rs new file mode 100644 index 0000000..6745809 --- /dev/null +++ b/src/tests/test_kalman.rs @@ -0,0 +1,152 @@ +use crate::algorithms::KalmanFilterNM; +use crate::common::na; +use crate::epoch::Epoch; +use crate::filter::Filter; +use crate::measurement::Observation; +use crate::models::*; +use crate::state::State; + +#[test] +fn test_kalman_custom() { + // initial state + let state = State::>::new( + vec![1.0, 1.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) + .ptr(); + + // Linear Transition model + let f = na::SMatrix::::from_diagonal(&na::SVector::from_row_slice(&[1.5, 2.0])); + + let one_step_state = f * state.read().unwrap().value; + + let q = &na::SMatrix::::identity() * 0.05; + + let transition = LinearTransitionN::<2>::new(state.clone(), f, q); + + // Linear Measurement model + let h = na::SMatrix::::identity(); + let r = &na::SMatrix::::identity() * 0.01; + + let measurement = LinearMeasurementNM::<2, 2>::new( + h, // observation model + r, // observation noise + ); + + // Kalman Update Model + let update = LinearUpdateNM::<2, 2>::new(measurement, transition); + + // Filter wrapper + let mut filter = KalmanFilterNM::new(update); + + let obs = Observation { + value: na::SMatrix::::from_row_slice(&[1.0, 1.0]), + epoch: Epoch::new(0), + }; + + // predict + let step_state = filter.predict(&Epoch::new(1)); + let _value = step_state.read().unwrap().value; + + let delta = (_value - one_step_state).norm(); + assert!( + delta < 1e-12, + "Incorrect State update: Abs Err: {}\n Expected: {} Found: {}", + delta, + one_step_state, + _value + ); + + // measurement update + let _ = filter.update(&obs); +} +/// The pseudo-inverse is chosen by shape, so a wide H gets the right inverse +/// (H H+ = I) and a tall H the left inverse (H+ H = I). +#[test] +fn test_moore_penrose_inverse_by_shape() { + use crate::common::linalg::moore_penrose_inverse; + + // det(H^T H) of this wide H rounds to a tiny non-zero value + let wide = na::SMatrix::::new(0.7, 0.9); + let wide_inv = moore_penrose_inverse(&wide).unwrap(); + assert!( + ((wide * wide_inv)[(0, 0)] - 1.0).abs() < 1e-12, + "H H+ = {}", + wide * wide_inv + ); + + let tall = wide.transpose(); + let tall_inv = moore_penrose_inverse(&tall).unwrap(); + assert!( + ((tall_inv * tall)[(0, 0)] - 1.0).abs() < 1e-12, + "H+ H = {}", + tall_inv * tall + ); + + // rank deficient + assert!(moore_penrose_inverse(&na::SMatrix::::zeros()).is_none()); +} + +/// With a non-invertible F, backward prediction leaves the state unchanged. +#[test] +fn test_backward_predict_with_singular_transition() { + let _quiet = crate::tests::expect_logs(); + let state = State::>::new( + vec![1.0, 2.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ); + let mut filter = KalmanFilterNM::new(LinearUpdateNM::<2, 2>::new( + LinearMeasurementNM::<2, 2>::new(na::SMatrix::identity(), na::SMatrix::identity()), + LinearTransitionN::<2>::new( + state.ptr(), + na::SMatrix::::new(1.0, 0.0, 0.0, 0.0), + na::SMatrix::identity(), + ), + )); + + let forward = filter.predict(&Epoch::new(2)).read().unwrap().clone(); + let backward = filter.predict(&Epoch::new(0)).read().unwrap().clone(); + assert_eq!(backward.value, forward.value); + assert_eq!(backward.covariance(), forward.covariance()); + assert_eq!(backward.epoch, Epoch::new(2)); +} + +/// `inverse()` is `None` when a model has no inverse, at any dimension type. +#[test] +fn test_inverse_measurement_is_optional() { + let obs = Observation { + value: na::SVector::::new(2.0), + epoch: Epoch::new(1), + }; + + // full rank wide H: the right inverse + let wide = LinearMeasurementNM::<2, 1>::new( + na::SMatrix::::new(0.7, 0.9), + na::SMatrix::identity(), + ); + let x = wide + .inverse(&obs) + .expect("full rank H has a pseudo-inverse"); + assert!(((0.7 * x.value[0] + 0.9 * x.value[1]) - 2.0).abs() < 1e-12); + + // rank deficient H: no inverse + let deficient = LinearMeasurementNM::<2, 1>::new( + na::SMatrix::::zeros(), + na::SMatrix::identity(), + ); + assert!(deficient.inverse(&obs).is_none()); + + // nonlinear without with_inverse, runtime sized (used to return a 0-length state) + let nonlinear = NonLinearMeasurement::::new( + Box::new(|s| s.value.clone()), + Box::new(|_| na::DMatrix::identity(2, 2)), + na::DMatrix::identity(2, 2), + ); + let dyn_obs = Observation { + value: na::DVector::from_element(2, 1.0), + epoch: Epoch::new(0), + }; + assert!(nonlinear.inverse(&dyn_obs).is_none()); +} diff --git a/src/tests/test_kalman_filter.rs b/src/tests/test_kalman_filter.rs deleted file mode 100644 index f63d8ca..0000000 --- a/src/tests/test_kalman_filter.rs +++ /dev/null @@ -1,66 +0,0 @@ -use crate::measurement::Observation; -use crate::state::State; -use crate::filter::Filter; -use crate::algorithms::KalmanFilterNM; -use crate::common::na as na; -use crate::epoch::Epoch; - -#[test] -fn test_kalman_update_scalar() { - // Small 1D Kalman filter: N=1, M=1 - - let x = State::>::new( - vec![0.0], - vec![vec![1.0]], - Epoch::new(0)).ptr(); - - let mut filter = KalmanFilterNM::<1, 1>::default_from_state(x); - // initial state x = 0 with noise and measurement matrices as identity - - assert_eq!(filter.state().read().unwrap().value[(0, 0)], 0.0); - - let obs = Observation{ - value: na::SMatrix::::from_row_slice(&[1.0]), - epoch: Epoch::new(0), - }; - - // predict (block pattern to avoid borrow conflict) - { let _ = filter.predict(&Epoch::new(1)); } - - // measurement update - let _ = filter.update(&obs); - - // With P=I, R=I, H=I we expect K = 1/2 and x -> 0.5 - let x_val = filter.state().read().unwrap().value[(0, 0)]; - assert!((x_val - 0.5).abs() < 1e-12, "x after update = {}", x_val); -} - -#[test] -fn test_kalman_update_2d() { - let x = State::>::new( - vec![0.0, 0.0], - vec![ - vec![1.0, 0.0], - vec![0.0,1.0]], - Epoch::new(0)).ptr(); - - // N=2, M=1. Observe only the first state element. - let mut filter = KalmanFilterNM::<2, 1>::default_from_state(x); - - let obs = Observation{ - value: na::SMatrix::::from_row_slice(&[2.0]), - epoch: Epoch::new(0), - }; - - // block pattern to avoid borrow conflict - { let _ = filter.predict(&Epoch { value: 1 }); } - { let _ = filter.update(&obs); } - - // Expect first state to be 1.0 (K = [0.5, 0]^T, y = 2) - let state_ptr = filter.state(); - let state = state_ptr.read().unwrap(); - let x0 = state.value[(0, 0)]; - let x1 = state.value[(1, 0)]; - assert!((x0 - 1.0).abs() < 1e-12, "x0 after update = {}", x0); - assert!((x1 - 0.0).abs() < 1e-12, "x1 after update = {}", x1); -} \ No newline at end of file diff --git a/src/tests/test_state.rs b/src/tests/test_state.rs index 30bf491..c46719e 100644 --- a/src/tests/test_state.rs +++ b/src/tests/test_state.rs @@ -1,7 +1,7 @@ -use crate::epoch::Epoch; -use crate::state::{ State, StateN, StateDyn }; use crate::common::Composable; -use crate::common::na as na; +use crate::common::na; +use crate::epoch::Epoch; +use crate::state::{State, StateDyn, StateN}; #[test] fn test_add() { @@ -29,7 +29,6 @@ fn static_alloc() { }; } - #[test] fn runtime_alloc() { // Dynamic state (runtime size) diff --git a/src/tests/test_ud.rs b/src/tests/test_ud.rs new file mode 100644 index 0000000..e334e76 --- /dev/null +++ b/src/tests/test_ud.rs @@ -0,0 +1,405 @@ +use crate::algorithms::{ + ExtendedKalmanFilterNM, KalmanFilterNM, UDExtendedKalmanFilterNM, UDKalmanFilterDyn, + UDKalmanFilterNM, +}; +use crate::common::linalg::udu_decomposition; +use crate::common::na; +use crate::epoch::Epoch; +use crate::filter::Filter; +use crate::measurement::Observation; +use crate::models::*; +use crate::state::State; + +fn state3() -> State> { + State::>::new( + vec![1.0, -0.5, 2.0], + vec![ + vec![2.0, 0.3, 0.1], + vec![0.3, 1.0, -0.2], + vec![0.1, -0.2, 0.5], + ], + Epoch::new(0), + ) +} + +fn close( + a: &na::Matrix, + b: &na::Matrix, +) -> bool +where + S1: na::Storage, + S2: na::Storage, +{ + let diff = a + .iter() + .zip(b.iter()) + .map(|(x, y)| (x - y) * (x - y)) + .sum::() + .sqrt(); + diff <= 1e-9 * (1.0 + b.norm()) +} + +#[test] +fn test_udu_decomposition() { + let m = state3().covariance(); + let (u, d) = udu_decomposition(&m); + for i in 0..3 { + assert_eq!(u[(i, i)], 1.0); + for j in 0..i { + assert_eq!(u[(i, j)], 0.0, "U must be upper triangular"); + } + } + let rebuilt = u * na::SMatrix::::from_diagonal(&d) * u.transpose(); + assert!(close(&rebuilt, &m), "{} vs {}", rebuilt, m); + + // semi-definite: a zero variance gives a zero pivot, not a NaN + let semi = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0); + let (u, d) = udu_decomposition(&semi); + assert!(u.iter().chain(d.iter()).all(|x| x.is_finite())); + assert!(close( + &(u * na::SMatrix::::from_diagonal(&d) * u.transpose()), + &semi + )); +} + +/// The UD filter is the Kalman filter in factored form: same estimates, for +/// diagonal and non-diagonal R (decorrelated), non-diagonal Q and epoch gaps. +#[test] +fn test_ud_matches_kalman() { + let f = na::SMatrix::::new(1.0, 0.1, 0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.95); + let q = na::SMatrix::::new(0.02, 0.005, 0.0, 0.005, 0.01, 0.002, 0.0, 0.002, 0.03); + let h = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0, 1.0, 1.0); + + for r in [ + na::SMatrix::::new(0.1, 0.0, 0.0, 0.2), // diagonal + na::SMatrix::::new(0.1, 0.04, 0.04, 0.2), // correlated + ] { + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(state3().ptr(), f, q), + )); + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(state3().ptr(), f, q), + )); + + for t in [1, 2, 5, 6, 9] { + let obs = Observation { + value: na::SVector::::new(0.3 * t as f64, -0.1 * t as f64), + epoch: Epoch::new(t), + }; + let a = kf.update(&obs); + let b = ud.update(&obs); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert!( + close(&b.value, &a.value), + "t={}: {} vs {}", + t, + b.value, + a.value + ); + assert!( + close(&b.covariance(), &a.covariance()), + "t={}: covariance", + t + ); + assert_eq!(a.epoch, b.epoch); + } + } +} + +/// The UD EKF matches the EKF for nonlinear models and several measurements. +#[test] +fn test_ud_extended_matches_extended() { + let transition = || { + NonLinearTransitionN::<2>::new( + State::>::new( + vec![2.0, 2.0], + vec![vec![1.0, 0.2], vec![0.2, 1.0]], + Epoch::new(0), + ) + .ptr(), + Box::new(|s| { + na::SVector::::new(s.value[0] + 0.1 * s.value[1].sin(), 0.98 * s.value[1]) + }), + Box::new(|s| na::SMatrix::::new(1.0, 0.1 * s.value[1].cos(), 0.0, 0.98)), + na::SMatrix::::identity() * 0.01, + ) + }; + // range and bearing + let measurement = || { + NonLinearMeasurementNM::<2, 2>::new( + Box::new(|s| na::SVector::::new(s.value.norm(), s.value[1].atan2(s.value[0]))), + Box::new(|s| { + let (x, y) = (s.value[0], s.value[1]); + let r2 = x * x + y * y; + let r = r2.sqrt(); + na::SMatrix::::new(x / r, y / r, -y / r2, x / r2) + }), + na::SMatrix::::new(0.05, 0.0, 0.0, 0.01), + ) + }; + + let mut ekf = + ExtendedKalmanFilterNM::new(NonLinearUpdateNM::<2, 2>::new(measurement(), transition())); + let mut ud = UDExtendedKalmanFilterNM::new(NonLinearUDUpdateNM::<2, 2>::new( + measurement(), + transition(), + )); + + for t in [1, 2, 4, 5] { + let obs = Observation { + value: na::SVector::::new(3.0, 0.8), + epoch: Epoch::new(t), + }; + let a = ekf.update(&obs); + let b = ud.update(&obs); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert!( + close(&b.value, &a.value), + "t={}: {} vs {}", + t, + b.value, + a.value + ); + assert!( + close(&b.covariance(), &a.covariance()), + "t={}: covariance", + t + ); + } +} + +/// Badly scaled problem: the factors keep P symmetric with non-negative D. +#[test] +fn test_ud_stays_positive_semi_definite() { + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<2, 1>::new( + LinearMeasurementNM::<2, 1>::new( + na::SMatrix::::new(1.0, 1.0), + na::SMatrix::::new(1e-8), + ), + LinearTransitionN::<2>::new( + State::>::new( + vec![0.0, 0.0], + vec![vec![1e8, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) + .ptr(), + na::SMatrix::::identity(), + na::SMatrix::::zeros(), + ), + )); + + for t in 1..50 { + let state = ud.update(&Observation { + value: na::SVector::::new(1.0), + epoch: Epoch::new(t), + }); + let p = state.read().unwrap().covariance(); + assert_eq!(p, p.transpose()); + assert!( + ud.update.factors.d.iter().all(|d| *d >= 0.0), + "t={}: D = {}", + t, + ud.update.factors.d + ); + } +} + +/// Setting the state from outside is picked up: the filter continues from it +/// exactly as a Kalman filter started there would. +#[test] +fn test_ud_refactorises_after_outside_change() { + let f = na::SMatrix::::identity() * 1.1; + let q = na::SMatrix::::identity() * 0.01; + let h = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0, 1.0, 0.0); + let r = na::SMatrix::::identity() * 0.1; + + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(state3().ptr(), f, q), + )); + ud.update(&Observation { + value: na::SVector::::new(1.0, 1.0), + epoch: Epoch::new(1), + }); + + // replace the state from outside + let mut fresh = state3(); + fresh.epoch = Epoch::new(1); + *ud.state.write().unwrap() = fresh.clone(); + + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(fresh.ptr(), f, q), + )); + + let obs = Observation { + value: na::SVector::::new(0.5, -0.5), + epoch: Epoch::new(3), + }; + let a = kf.update(&obs); + let b = ud.update(&obs); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert!(close(&b.value, &a.value)); + assert!(close(&b.covariance(), &a.covariance())); +} + +/// Runtime sized (`na::Dyn`, as used by Python) matches statically sized exactly. +#[test] +fn test_ud_dyn_matches_static() { + let s = state3(); + let f = na::SMatrix::::new(1.0, 0.1, 0.0, 0.0, 1.0, 0.1, 0.0, 0.0, 0.95); + let q = na::SMatrix::::identity() * 0.02; + let h = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0, 1.0, 1.0); + let r = na::SMatrix::::new(0.1, 0.04, 0.04, 0.2); + let dynm = |m: &dyn Fn(usize, usize) -> f64, rows, cols| na::DMatrix::from_fn(rows, cols, m); + + let mut fixed = UDKalmanFilterNM::new(LinearUDUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(s.ptr(), f, q), + )); + let mut dynamic = UDKalmanFilterDyn::new(LinearUDUpdate::new( + LinearMeasurement::new(dynm(&|i, j| h[(i, j)], 2, 3), dynm(&|i, j| r[(i, j)], 2, 2)), + LinearTransition::new( + State::::new( + s.value.as_slice().to_vec(), + (0..3) + .map(|i| (0..3).map(|j| s.covariance()[(i, j)]).collect()) + .collect(), + Epoch::new(0), + ) + .ptr(), + dynm(&|i, j| f[(i, j)], 3, 3), + dynm(&|i, j| q[(i, j)], 3, 3), + ), + )); + + for t in [1, 3, 4] { + let z = [0.3 * t as f64, -0.2]; + let a = fixed.update(&Observation { + value: na::SVector::::from_row_slice(&z), + epoch: Epoch::new(t), + }); + let b = dynamic.update(&Observation { + value: na::DVector::from_row_slice(&z), + epoch: Epoch::new(t), + }); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert_eq!(a.value.as_slice(), b.value.as_slice()); + assert_eq!(a.covariance().as_slice(), b.covariance().as_slice()); + } +} + +/// Same-epoch predict does nothing; backward predict and a singular +/// innovation leave the state unchanged. +#[test] +fn test_ud_edge_cases() { + let _quiet = crate::tests::expect_logs(); + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new( + na::SMatrix::::zeros(), + na::SMatrix::::zeros(), + ), + LinearTransitionN::<3>::new( + state3().ptr(), + na::SMatrix::::identity() * 2.0, + na::SMatrix::::zeros(), + ), + )); + let before = state3(); + + ud.predict(&Epoch::new(0)); + ud.predict(&Epoch::new(-1)); + let after = ud.state(); + assert_eq!(after.read().unwrap().value, before.value); + assert_eq!(after.read().unwrap().epoch, Epoch::new(0)); + + // H = 0 and R = 0: the innovation variance is zero + let state = ud.update(&Observation { + value: na::SVector::::new(1.0, 1.0), + epoch: Epoch::new(0), + }); + assert_eq!(state.read().unwrap().value, before.value); +} + +/// A perfect measurement (R = 0) of a state the first factor column does not +/// see is handled like the Kalman filter, not rejected. +#[test] +fn test_ud_perfect_measurement() { + let h = na::SMatrix::::new(0.0, 1.0); + let r = na::SMatrix::::zeros(); + let state = || { + State::>::new( + vec![0.0, 0.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + Epoch::new(0), + ) + }; + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<2, 1>::new( + LinearMeasurementNM::<2, 1>::new(h, r), + LinearTransitionN::<2>::new(state().ptr(), na::SMatrix::identity(), na::SMatrix::zeros()), + )); + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<2, 1>::new( + LinearMeasurementNM::<2, 1>::new(h, r), + LinearTransitionN::<2>::new(state().ptr(), na::SMatrix::identity(), na::SMatrix::zeros()), + )); + + let obs = Observation { + value: na::SVector::::new(2.0), + epoch: Epoch::new(0), + }; + let a = kf.update(&obs); + let b = ud.update(&obs); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert!(close(&b.value, &a.value), "{} vs {}", b.value, a.value); + assert!( + close(&b.covariance(), &a.covariance()), + "{} vs {}", + b.covariance(), + a.covariance() + ); + assert!( + ud.update + .factors + .d + .iter() + .all(|d| d.is_finite() && *d >= 0.0) + ); +} + +/// Q is factorised once and reused, but a Q changed between steps +/// (non-stationary process noise) is picked up. +#[test] +fn test_ud_follows_changed_process_noise() { + let f = na::SMatrix::::identity(); + let h = na::SMatrix::::new(1.0, 0.0, 0.0, 0.0, 1.0, 0.0); + let r = na::SMatrix::::identity() * 0.1; + let q1 = na::SMatrix::::identity() * 0.01; + let q2 = na::SMatrix::::new(0.5, 0.1, 0.0, 0.1, 0.3, 0.0, 0.0, 0.0, 0.2); + + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(state3().ptr(), f, q1), + )); + let mut ud = UDKalmanFilterNM::new(LinearUDUpdateNM::<3, 2>::new( + LinearMeasurementNM::<3, 2>::new(h, r), + LinearTransitionN::<3>::new(state3().ptr(), f, q1), + )); + + for t in 1..=6 { + if t == 4 { + kf.update.transition.q = q2; + ud.update.transition.q = q2; + } + let obs = Observation { + value: na::SVector::::new(0.1 * t as f64, 0.2), + epoch: Epoch::new(t), + }; + let a = kf.update(&obs); + let b = ud.update(&obs); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert!(close(&b.value, &a.value), "t={}", t); + assert!(close(&b.covariance(), &a.covariance()), "t={}", t); + } +} diff --git a/src/tests/test_ukf.rs b/src/tests/test_ukf.rs new file mode 100644 index 0000000..1e406d8 --- /dev/null +++ b/src/tests/test_ukf.rs @@ -0,0 +1,258 @@ +use crate::algorithms::{KalmanFilterNM, UnscentedKalmanFilterDyn, UnscentedKalmanFilterNM}; +use crate::common::na; +use crate::epoch::Epoch; +use crate::filter::Filter; +use crate::measurement::Observation; +use crate::models::*; +use crate::state::State; + +fn state2(value: [f64; 2], var: f64) -> State> { + State::>::new( + value.to_vec(), + vec![vec![var, 0.0], vec![0.0, var]], + Epoch::new(0), + ) +} + +/// The unscented transform is exact for linear models, so the UKF must +/// match the linear Kalman filter, whatever the sigma point parameters. +#[test] +fn test_unscented_matches_linear_kalman() { + let f = na::SMatrix::::new(1.0, 0.5, 0.0, 1.0); + let h = na::SMatrix::::new(1.0, 0.0); + let q = na::SMatrix::::identity() * 0.05; + let r = na::SMatrix::::new(0.1); + + for params in [ + UnscentedParams::default(), + UnscentedParams { + alpha: 0.1, + beta: 2.0, + kappa: 1.0, + }, + ] { + let mut kf = KalmanFilterNM::new(LinearUpdateNM::<2, 1>::new( + LinearMeasurementNM::<2, 1>::new(h, r), + LinearTransitionN::<2>::new(state2([0.0, 1.0], 1.0).ptr(), f, q), + )); + let mut ukf = UnscentedKalmanFilterNM::new(UnscentedUpdateNM::<2, 1>::new( + UnscentedMeasurementNM::<2, 1>::new(Box::new(move |s| h * s.value), r), + UnscentedTransitionN::<2>::new( + state2([0.0, 1.0], 1.0).ptr(), + Box::new(move |s| f * s.value), + q, + ) + .with_params(params), + )); + + for t in [1, 2, 4, 5] { + let obs = Observation { + value: na::SVector::::new(0.4 * t as f64), + epoch: Epoch::new(t), + }; + let a = kf.update(&obs); + let b = ukf.update(&obs); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + + assert!( + (a.value - b.value).norm() < 1e-9, + "{:?}: state {} vs {}", + params, + a.value, + b.value + ); + assert!( + (a.covariance() - b.covariance()).norm() < 1e-9, + "{:?}: covariance", + params + ); + assert_eq!(a.epoch, b.epoch); + } + } +} + +/// For f(x) = x^2 with x ~ N(m, s^2), the default sigma points recover the +/// exact mean m^2 + s^2 and variance 4 m^2 s^2 + 2 s^4. (The EKF gives m^2.) +#[test] +fn test_unscented_transform_moments() { + let (m, var) = (2.0, 0.5); + let state = State::>::new(vec![m], vec![vec![var]], Epoch::new(0)).ptr(); + let mut transition = UnscentedTransitionN::<1>::new( + state, + Box::new(|s| s.value.map(|x| x * x)), + na::SMatrix::::zeros(), + ); + + let out = transition.state(&Epoch::new(1)).read().unwrap().clone(); + assert!( + (out.value[0] - (m * m + var)).abs() < 1e-12, + "mean {}", + out.value[0] + ); + let expected_var = 4.0 * m * m * var + 2.0 * var * var; + assert!( + (out.covariance()[(0, 0)] - expected_var).abs() < 1e-12, + "variance {}", + out.covariance()[(0, 0)] + ); +} + +/// Range-only sensor on a static target. Matches an independent numpy UKF +/// (same sigma points). The range constrains only the radius, so with large +/// tangential uncertainty the posterior is ring-shaped and its mean lies +/// inside the ring: E|x|^2 = |mean|^2 + tr(P) approaches 25, not |mean|. +#[test] +fn test_unscented_range_sensor() { + // radius of the mean after each update, from the numpy reference + let reference = [3.8639, 4.3361, 4.4525, 4.4974, 4.5210]; + + let mut ukf = UnscentedKalmanFilterNM::new(UnscentedUpdateNM::<2, 1>::new( + UnscentedMeasurementNM::<2, 1>::new( + Box::new(|s| na::SVector::::new(s.value.norm())), + na::SMatrix::::new(0.01), + ), + UnscentedTransitionN::<2>::new( + state2([2.0, 2.0], 4.0).ptr(), + Box::new(|s| s.value), + na::SMatrix::::zeros(), + ), + )); + + let mut second_moment = 8.0 + 8.0; + for (t, expected) in (1..=5).zip(reference) { + let state = ukf.update(&Observation { + value: na::SVector::::new(5.0), + epoch: Epoch::new(t), + }); + let state = state.read().unwrap(); + + assert!( + (state.value.norm() - expected).abs() < 1e-4, + "t={}: radius {} vs reference {}", + t, + state.value.norm(), + expected + ); + + let moment = state.value.norm_squared() + state.covariance().trace(); + assert!( + moment > second_moment && moment < 25.0, + "t={}: E|x|^2 = {}", + t, + moment + ); + second_moment = moment; + } +} + +/// Runtime sized (`na::Dyn`, as used by Python) matches statically sized exactly. +#[test] +fn test_unscented_dyn_matches_static() { + let range = |x: &[f64]| (x[0] * x[0] + x[1] * x[1]).sqrt(); + + let mut fixed = UnscentedKalmanFilterNM::new(UnscentedUpdateNM::<2, 1>::new( + UnscentedMeasurementNM::<2, 1>::new( + Box::new(move |s| na::SVector::::new(range(s.value.as_slice()))), + na::SMatrix::::new(0.01), + ), + UnscentedTransitionN::<2>::new( + state2([2.0, 2.0], 4.0).ptr(), + Box::new(|s| s.value * 1.01), + na::SMatrix::::identity() * 0.01, + ), + )); + let mut dynamic = UnscentedKalmanFilterDyn::new(UnscentedUpdate::new( + UnscentedMeasurement::new( + Box::new(move |s| na::DVector::from_element(1, range(s.value.as_slice()))), + na::DMatrix::from_element(1, 1, 0.01), + ), + UnscentedTransition::new( + State::::new( + vec![2.0, 2.0], + vec![vec![4.0, 0.0], vec![0.0, 4.0]], + Epoch::new(0), + ) + .ptr(), + Box::new(|s| &s.value * 1.01), + na::DMatrix::identity(2, 2) * 0.01, + ), + )); + + for t in 1..=4 { + let epoch = Epoch::new(t); + let a = fixed.update(&Observation { + value: na::SVector::::new(5.0), + epoch, + }); + let b = dynamic.update(&Observation { + value: na::DVector::from_element(1, 5.0), + epoch, + }); + let (a, b) = (a.read().unwrap(), b.read().unwrap()); + assert_eq!(a.value.as_slice(), b.value.as_slice()); + assert_eq!(a.covariance().as_slice(), b.covariance().as_slice()); + } +} + +/// Same-epoch predict does nothing; backward predict and a covariance that is +/// not positive definite leave the state unchanged. +#[test] +fn test_unscented_edge_cases() { + // let _quiet = crate::tests::expect_logs(); + let mut transition = UnscentedTransitionN::<2>::new( + state2([1.0, 1.0], 1.0).ptr(), + Box::new(|s| s.value * 2.0), + na::SMatrix::::zeros(), + ); + + let before = transition.state.read().unwrap().clone(); + transition.state(&Epoch::new(0)); + transition.state(&Epoch::new(-1)); + let after = transition.state.read().unwrap().clone(); + assert_eq!(before.value, after.value); + assert_eq!(before.epoch, after.epoch); + + let mut bad = UnscentedTransitionN::<2>::new( + state2([1.0, 1.0], -1.0).ptr(), + Box::new(|s| s.value * 2.0), + na::SMatrix::::zeros(), + ); + bad.state(&Epoch::new(1)); + let after = bad.state.read().unwrap().clone(); + assert_eq!(after.value, na::SVector::::new(1.0, 1.0)); + assert_eq!(after.epoch, Epoch::new(0)); +} + +/// The measurement reports its noise, and uses the inverse supplied with +/// `with_inverse`. +#[test] +fn test_unscented_measurement_noise_and_inverse() { + let r = na::SMatrix::::new(0.25); + let measurement = UnscentedMeasurementNM::<2, 1>::new( + Box::new(|s| na::SVector::::new(s.value[0])), + r, + ) + .with_inverse(Box::new(|obs| { + State::>::new( + vec![obs.value[0], 0.0], + vec![vec![1.0, 0.0], vec![0.0, 1.0]], + obs.epoch, + ) + })); + + assert_eq!(measurement.noise(), r); + let obs = Observation { + value: na::SVector::::new(3.0), + epoch: Epoch::new(2), + }; + let x = measurement.inverse(&obs).expect("an inverse was supplied"); + assert_eq!(x.value, na::SVector::::new(3.0, 0.0)); + assert_eq!(x.epoch, Epoch::new(2)); + + // without with_inverse there is no inverse + let plain = UnscentedMeasurementNM::<2, 1>::new( + Box::new(|s| na::SVector::::new(s.value[0])), + r, + ); + assert!(plain.inverse(&obs).is_none()); +} diff --git a/tests/test_kalman_filter.rs b/tests/test_kalman_filter.rs deleted file mode 100644 index f0324b5..0000000 --- a/tests/test_kalman_filter.rs +++ /dev/null @@ -1,30 +0,0 @@ -use determined::algorithms::KalmanFilterNM; -use determined::filter::Filter; -use determined::common::na as na; -use determined::measurement::Observation; -use determined::state::{ State, StatePtr }; -use determined::epoch::Epoch; - -#[test] -fn test_kalman_filter() { - - let x: StatePtr> = State::>::new( - vec![1.0], - vec![vec![1.0]], - Epoch{ value: 0 }).ptr(); - - let mut kf: KalmanFilterNM::<1, 1> = - KalmanFilterNM::default_from_state(x); - - let epoch = Epoch { value: 0 }; - let z = Observation::> { - value: na::SMatrix::::from_row_slice(&[1.2]), - epoch: epoch, - }; - - let _ = kf.predict(&epoch); - let _ = kf.update(&z); - - let s = kf.state(); - println!("state = {:?}", s); -} \ No newline at end of file diff --git a/tests_py/harness/models.py b/tests_py/harness/models.py index e24ca16..85f3995 100644 --- a/tests_py/harness/models.py +++ b/tests_py/harness/models.py @@ -10,12 +10,15 @@ def __init__(self, h, r): def projection(self, state: dt.State): z = self.h @ state.value - return dt.Observation(z, state.epoch) + return dt.Observation(z, epoch=state.epoch) def inverse(self, observation: dt.Observation): x = self.h_inv @ observation.value cov = np.eye(x.size) - return dt.State(x, cov, observation.epoch) + return dt.State(x, covariance=cov, epoch=observation.epoch) + + def noise(self): + return self.r def jacobian(self, state: dt.State): return self.jac @@ -34,7 +37,7 @@ def state(self, epoch: dt.Epoch): # state / covariance update value = self.f @ self._state.value cov = self.f @ self._state.covariance @ self.f.T + self.q - self._state = dt.State(value, cov, epoch) + self._state = dt.State(value, covariance=cov, epoch=epoch) return self._state @@ -61,7 +64,7 @@ def apply(self, observation: dt.Observation): # compute gains h = self.m.model.h h_t = h.transpose() - r = self.m.model.r + r = self.m.noise() s = h @ _cov @ h_t + r s_inv = np.linalg.pinv(s) @@ -72,8 +75,4 @@ def apply(self, observation: dt.Observation): value = _state.value + k @ y cov = (np.eye(value.size) - k @ h) @ _cov - return dt.State(value, cov, epoch) - - def jacobian(self, state: dt.State): - jac = self.t.jacobian(state) - return jac \ No newline at end of file + return dt.State(value, covariance=cov, epoch=epoch) \ No newline at end of file diff --git a/tests_py/test_algorithm.py b/tests_py/test_custom_filter.py similarity index 63% rename from tests_py/test_algorithm.py rename to tests_py/test_custom_filter.py index 5bb04e5..acacf4a 100644 --- a/tests_py/test_algorithm.py +++ b/tests_py/test_custom_filter.py @@ -16,7 +16,7 @@ def setUpClass(cls): cls.value = np.array([1.0, 1.0]) cls.cov = np.eye(cls.value.size) cls.epoch = dt.Epoch(0) - cls.obs = dt.Observation(np.array([0.0, 1.0]), cls.epoch) + cls.obs = dt.Observation(np.array([0.0, 1.0]), epoch=cls.epoch) cls.f = np.diag(np.array([2.0, 1.5])) cls.q = np.eye(cls.f.shape[0]) @@ -31,19 +31,19 @@ def setUpClass(cls): # Rebuild models at each test to avoid persistent state def setUp(self): - self.state = dt.State(self.value, self.cov, self.epoch) + self.state = dt.State(self.value, covariance=self.cov, epoch=self.epoch) self.t = LinearTransition(self.f, self.q, self.state) - self.transition = dt.TransitionModel(self.t, self.state) + self.transition = dt.TransitionModel(self.t, state=self.state) self.m = LinearMeasurement(self.h, self.r) self.measurement = dt.MeasurementModel(self.m) self.u = LinearUpdate(self.transition, self.measurement) - self.update = dt.UpdateModel(self.u, self.transition, self.measurement) + self.update = dt.UpdateModel(self.u, transition=self.transition, measurement=self.measurement) # this is basically just a type-erased wrapper - self.algorithm = dt.KalmanFilter(self.update) + self.algorithm = dt.CustomFilter(self.update) return super().setUp() def test_predict(self): @@ -61,3 +61,18 @@ def test_predict(self): def test_apply(self): new_state = self.algorithm.update(self.obs) assert isinstance(new_state, dt.State), "state should be dt.State" + + def test_reset_without_reset_method_warns(self): + with self.assertLogs("determined", level="WARNING") as logs: + self.algorithm.reset() + self.assertIn("No update.reset()", logs.output[0]) + + def test_reset_error_propagates(self): + class Failing(LinearUpdate): + def reset(self): + raise RuntimeError("reset failed") + + update = dt.UpdateModel(Failing(self.transition, self.measurement), + transition=self.transition, measurement=self.measurement) + with self.assertRaises(RuntimeError): + dt.CustomFilter(update).reset() diff --git a/tests_py/test_ekf.py b/tests_py/test_ekf.py new file mode 100644 index 0000000..4de0cf3 --- /dev/null +++ b/tests_py/test_ekf.py @@ -0,0 +1,310 @@ +import unittest +import numpy as np + +import determined as dt + + +def make_ekf(state, f_mat, q, h_mat, r): + return dt.ExtendedKalmanFilter( + state, + f=lambda x, t: f_mat @ x, + jac_f=lambda x, t: f_mat, + q=q, + h=lambda x, t: h_mat @ x, + jac_h=lambda x, t: h_mat, + r=r, + ) + + +def make_state(value, cov=None, epoch=0): + value = np.asarray(value, dtype=float) + cov = np.eye(value.size) if cov is None else cov + return dt.State(value, covariance=cov, epoch=dt.Epoch(epoch)) + + +class TestExtendedKalmanFilter(unittest.TestCase): + + def test_predict(self): + state = make_state([1.0, 1.0]) + f = np.diag([2.0, 3.0]) + ekf = make_ekf(state, f, np.zeros((2, 2)), np.eye(2), np.eye(2)) + + out = ekf.predict(dt.Epoch(2)) + np.testing.assert_allclose(out.value, [4.0, 9.0]) + np.testing.assert_allclose(out.covariance, np.diag([16.0, 81.0])) + self.assertEqual(out.epoch.value, 2) + + # zero-copy: shares storage with the state it was built from + np.testing.assert_array_equal(state.value, out.value) + + # predicting to the same epoch again must not propagate + out = ekf.predict(dt.Epoch(2)) + np.testing.assert_allclose(out.value, [4.0, 9.0]) + + def test_returned_states_are_live_views(self): + """Documented behaviour: results change on the next call.""" + state = make_state([1.0, 1.0]) + ekf = make_ekf(state, np.diag([2.0, 2.0]), np.zeros((2, 2)), np.eye(2), np.eye(2)) + + first = ekf.predict(dt.Epoch(1)) + kept = first.value.copy() + ekf.predict(dt.Epoch(2)) + + np.testing.assert_allclose(kept, [2.0, 2.0]) + np.testing.assert_allclose(first.value, [4.0, 4.0]) + + def test_matches_reference_kalman_filter(self): + """A linear problem must match a closed-form KF written in numpy.""" + f = np.array([[1.0, 0.5], [0.0, 1.0]]) # deliberately not symmetric + q = np.eye(2) * 0.05 + h = np.array([[1.0, 0.0]]) # non-square (1 x 2) + r = np.array([[0.1]]) + value, cov = np.array([0.0, 1.0]), np.eye(2) + + ekf = make_ekf(make_state(value, cov), f, q, h, r) + + x, p = value.copy(), cov.copy() + for t in range(1, 5): + z = np.array([0.4 * t]) + ekf.predict(dt.Epoch(t)) + out = ekf.update(dt.Observation(z, epoch=dt.Epoch(t))) + + x, p = f @ x, f @ p @ f.T + q + s = h @ p @ h.T + r + k = p @ h.T @ np.linalg.inv(s) + x, p = x + k @ (z - h @ x), (np.eye(2) - k @ h) @ p + + np.testing.assert_allclose(out.value, x, atol=1e-12) + np.testing.assert_allclose(out.covariance, p, atol=1e-12) + + def test_covariance_stays_symmetric_positive(self): + """Joseph form: badly scaled problem must keep P symmetric and PSD.""" + state = make_state([0.0, 0.0], np.diag([1e8, 1.0])) + h = np.array([[1.0, 1.0]]) + ekf = make_ekf(state, np.eye(2), np.zeros((2, 2)), h, np.array([[1e-8]])) + + for t in range(1, 50): + out = ekf.update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(t))) + cov = out.covariance + np.testing.assert_allclose(cov, cov.T, rtol=0, atol=1e-6) + self.assertGreaterEqual(np.linalg.eigvalsh((cov + cov.T) / 2).min(), + -1e-9 * np.abs(cov).max()) + + def test_range_sensor(self): + """Nonlinear range sensor pulls the estimate to the truth.""" + truth_range = 5.0 + state = make_state([2.0, 2.0], np.eye(2) * 4.0) + ekf = dt.ExtendedKalmanFilter( + state, + f=lambda x, t: x, + jac_f=lambda x, t: np.eye(2), + q=np.zeros((2, 2)), + h=lambda x, t: np.array([np.linalg.norm(x)]), + jac_h=lambda x, t: (x / np.linalg.norm(x))[None, :], + r=np.array([[0.01]]), + ) + + err = abs(np.linalg.norm(state.value) - truth_range) + trace = np.trace(state.covariance) + for t in range(1, 6): + out = ekf.update(dt.Observation(np.array([truth_range]), epoch=dt.Epoch(t))) + new_err = abs(np.linalg.norm(out.value) - truth_range) + new_trace = np.trace(out.covariance) + self.assertLess(new_err, err) + self.assertLess(new_trace, trace) + err, trace = new_err, new_trace + self.assertLess(err, 1e-2) + + def test_models_receive_epoch(self): + """f / jac_f get the epoch the step starts at, h / jac_h the measurement epoch.""" + f_epochs, h_epochs = [], [] + + def f(x, t): + f_epochs.append(t) + return x + t # time-varying drift + + def h(x, t): + h_epochs.append(t) + return x * t + + state = make_state([0.0]) + ekf = dt.ExtendedKalmanFilter( + state, + f=f, jac_f=lambda x, t: np.eye(1), q=np.zeros((1, 1)), + h=h, jac_h=lambda x, t: np.eye(1) * t, r=np.eye(1), + ) + + out = ekf.predict(dt.Epoch(3)) # steps from t=0, 1, 2 + self.assertEqual(f_epochs, [0, 1, 2]) + np.testing.assert_allclose(out.value, [0 + 1 + 2]) + + ekf.update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(3))) + self.assertEqual(h_epochs, [3]) + + def test_callables_may_return_array_likes(self): + """Lists / (N, 1) arrays are coerced, as AD libraries return their own types.""" + state = make_state([1.0, 2.0]) + ekf = dt.ExtendedKalmanFilter( + state, + f=lambda x, t: [float(x[0]), float(x[1])], + jac_f=lambda x, t: [[1, 0], [0, 1]], + q=np.zeros((2, 2)), + h=lambda x, t: x.reshape(2, 1), + jac_h=lambda x, t: np.eye(2), + r=np.eye(2), + ) + out = ekf.update(dt.Observation(np.array([1.0, 2.0]), epoch=dt.Epoch(1))) + np.testing.assert_allclose(out.value, [1.0, 2.0]) + + def test_reset_and_state_setter(self): + state = make_state([1.0, 1.0], np.eye(2) * 3.0) + ekf = make_ekf(state, np.eye(2), np.eye(2), np.eye(2), np.eye(2)) + + ekf.reset() + np.testing.assert_array_equal(ekf.state.value, [0.0, 0.0]) + np.testing.assert_array_equal(ekf.state.covariance, np.eye(2)) + + ekf.state = make_state([5.0, 6.0], epoch=3) + np.testing.assert_array_equal(ekf.state.value, [5.0, 6.0]) + self.assertEqual(ekf.state.epoch.value, 3) + + def test_arguments_are_keyword_only(self): + eye = np.eye(2) + with self.assertRaises(TypeError): + dt.ExtendedKalmanFilter( + make_state([1.0, 1.0]), lambda x, t: x, lambda x, t: eye, eye, + lambda x, t: x, lambda x, t: eye, eye) + + def test_bad_inputs_raise(self): + state = make_state([1.0, 1.0]) + eye = np.eye(2) + + def build(**overrides): + args = dict(f=lambda x, t: x, jac_f=lambda x, t: eye, q=eye, + h=lambda x, t: x, jac_h=lambda x, t: eye, r=eye) + args.update(overrides) + return dt.ExtendedKalmanFilter(state, **args) + + with self.assertRaises(TypeError): + build(f=1) + with self.assertRaises(ValueError): + build(q=np.eye(3)) + + # Jacobian of the wrong shape surfaces as ValueError, not a panic + with self.assertRaises(ValueError): + build(jac_f=lambda x, t: np.eye(3)).predict(dt.Epoch(1)) + + # wrong observation size + with self.assertRaises(ValueError): + build().update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(1))) + + # errors raised inside the callable propagate + def boom(x, t): + raise RuntimeError("boom") + with self.assertRaises(RuntimeError): + build(f=boom).predict(dt.Epoch(1)) + + +class TestCallbackSafety(unittest.TestCase): + + def test_error_in_callback_leaves_state_unchanged(self): + calls = [] + + def f(x, t): + calls.append(t) + if t == 1: + raise RuntimeError("boom at t=1") + return x + 1.0 + + state = make_state([0.0]) + ekf = dt.ExtendedKalmanFilter( + state, f=f, jac_f=lambda x, t: np.eye(1), q=np.eye(1), + h=lambda x, t: x, jac_h=lambda x, t: np.eye(1), r=np.eye(1)) + + with self.assertRaises(RuntimeError): + ekf.predict(dt.Epoch(3)) # fails on the second step + + # nothing was half-applied, and no further steps ran after the error + np.testing.assert_array_equal(state.value, [0.0]) + np.testing.assert_array_equal(state.covariance, np.eye(1)) + self.assertEqual(state.epoch.value, 0) + self.assertEqual(calls, [0, 1]) + + def test_callback_may_read_the_state(self): + """Model functions run without the state lock held, so this must not deadlock.""" + state = make_state([1.0, 2.0]) + seen = [] + + def f(x, t): + seen.append(state.value.copy()) + return x + + ekf = dt.ExtendedKalmanFilter( + state, f=f, jac_f=lambda x, t: np.eye(2), q=np.eye(2), + h=lambda x, t: x, jac_h=lambda x, t: np.eye(2), r=np.eye(2)) + ekf.predict(dt.Epoch(1)) + np.testing.assert_array_equal(seen[0], [1.0, 2.0]) + + +class TestKeywordOnlyConstructors(unittest.TestCase): + + def test_positional_arguments_after_the_first_are_rejected(self): + value, cov, epoch = np.zeros(2), np.eye(2), dt.Epoch(0) + with self.assertRaises(TypeError): + dt.State(value, cov, epoch) + with self.assertRaises(TypeError): + dt.Observation(value, epoch) + dt.State(value, covariance=cov, epoch=epoch) + dt.Observation(value, epoch=epoch) + + +class TestMatrixLayout(unittest.TestCase): + """Round trips of non-square / non-symmetric matrices through the bindings. + + These pass regardless of the internal layout; the non-square `H` in + `test_matches_reference_kalman_filter` is what checks the layout is correct. + """ + + def test_nonsquare_jacobian_roundtrip(self): + jac = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + + class Measurement: + def projection(self, s): return dt.Observation(np.zeros(2), epoch=s.epoch) + def inverse(self, o): return dt.State(np.zeros(3), covariance=np.eye(3), epoch=o.epoch) + def jacobian(self, s): return jac + + model = dt.MeasurementModel(Measurement()) + state = make_state(np.zeros(3)) + np.testing.assert_array_equal(model.jacobian(state), jac) + + def test_covariance_roundtrip(self): + cov = np.array([[1.0, 2.0], [3.0, 4.0]]) + state = make_state(np.zeros(2), cov) + np.testing.assert_array_equal(state.covariance, cov) + + +class TestLogging(unittest.TestCase): + """Rust-side log messages arrive through Python's `logging`.""" + + def test_singular_update_is_logged(self): + ekf = dt.ExtendedKalmanFilter( + make_state([1.0, 1.0]), + f=lambda x, t: x, jac_f=lambda x, t: np.eye(2), q=np.zeros((2, 2)), + h=lambda x, t: np.zeros(1), jac_h=lambda x, t: np.zeros((1, 2)), + r=np.zeros((1, 1))) + + with self.assertLogs("determined", level="ERROR") as logs: + ekf.update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(0))) + self.assertIn("singular", logs.output[0]) + + def test_backward_predict_is_logged(self): + state = make_state([1.0, 1.0], epoch=5) + ekf = make_ekf(state, np.eye(2), np.eye(2), np.eye(2), np.eye(2)) + + with self.assertLogs("determined", level="WARNING") as logs: + ekf.predict(dt.Epoch(2)) + self.assertIn("Backward", logs.output[0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests_py/test_kalman.py b/tests_py/test_kalman.py new file mode 100644 index 0000000..e6c3e3f --- /dev/null +++ b/tests_py/test_kalman.py @@ -0,0 +1,144 @@ +import unittest +import numpy as np + +import determined as dt + + +def make_state(value, cov=None, epoch=0): + value = np.asarray(value, dtype=float) + cov = np.eye(value.size) if cov is None else cov + return dt.State(value, covariance=cov, epoch=dt.Epoch(epoch)) + + +F = np.array([[1.0, 0.5], [0.0, 1.0]]) # deliberately not symmetric +Q = np.eye(2) * 0.05 +H = np.array([[1.0, 0.0]]) # non-square (1 x 2) +R = np.array([[0.1]]) + + +def make_kf(state=None): + return dt.KalmanFilter(state or make_state([0.0, 1.0]), f=F, q=Q, h=H, r=R) + + +class TestKalmanFilter(unittest.TestCase): + + def test_matches_reference(self): + """Matches a closed-form KF in numpy, including multi-epoch gaps.""" + kf = make_kf() + x, p = np.array([0.0, 1.0]), np.eye(2) + + epoch = 0 + for gap in [1, 1, 3, 2]: + epoch += gap + z = np.array([0.4 * epoch]) + kf.predict(dt.Epoch(epoch)) + out = kf.update(dt.Observation(z, epoch=dt.Epoch(epoch))) + + for _ in range(gap): + x, p = F @ x, F @ p @ F.T + Q + s = H @ p @ H.T + R + k = p @ H.T @ np.linalg.inv(s) + x, p = x + k @ (z - H @ x), (np.eye(2) - k @ H) @ p + + np.testing.assert_allclose(out.value, x, atol=1e-10) + np.testing.assert_allclose(out.covariance, p, atol=1e-10) + self.assertEqual(out.epoch.value, epoch) + + def test_update_propagates_to_observation_epoch(self): + """update() alone (no predict) brings the state to the observation epoch.""" + with_predict, without = make_kf(), make_kf() + obs = dt.Observation(np.array([1.0]), epoch=dt.Epoch(3)) + + with_predict.predict(dt.Epoch(3)) + a = with_predict.update(obs) + b = without.update(obs) + + np.testing.assert_allclose(a.value, b.value) + np.testing.assert_allclose(a.covariance, b.covariance) + + def test_matches_extended_kalman_filter(self): + kf = make_kf() + ekf = dt.ExtendedKalmanFilter( + make_state([0.0, 1.0]), + f=lambda x, t: F @ x, jac_f=lambda x, t: F, q=Q, + h=lambda x, t: H @ x, jac_h=lambda x, t: H, r=R) + + for t in range(1, 5): + obs = dt.Observation(np.array([0.3 * t]), epoch=dt.Epoch(t)) + a, b = kf.update(obs), ekf.update(obs) + np.testing.assert_allclose(a.value, b.value, atol=1e-12) + np.testing.assert_allclose(a.covariance, b.covariance, atol=1e-12) + + def test_predict_is_idempotent_and_state_is_shared(self): + state = make_state([0.0, 1.0]) + kf = make_kf(state) + + first = kf.predict(dt.Epoch(2)).value.copy() + again = kf.predict(dt.Epoch(2)) + np.testing.assert_allclose(again.value, first) + np.testing.assert_allclose(first, F @ F @ [0.0, 1.0]) + + # zero-copy: the state passed in is the filter's state + np.testing.assert_array_equal(state.value, first) + self.assertEqual(state.epoch.value, 2) + + def test_reset_and_state_setter(self): + kf = make_kf() + kf.predict(dt.Epoch(1)) + + kf.reset() + np.testing.assert_array_equal(kf.state.value, [0.0, 0.0]) + np.testing.assert_array_equal(kf.state.covariance, np.eye(2)) + + kf.state = make_state([5.0, 6.0], epoch=3) + np.testing.assert_array_equal(kf.state.value, [5.0, 6.0]) + self.assertEqual(kf.state.epoch.value, 3) + + with self.assertRaises(ValueError): + kf.state = make_state([1.0, 2.0, 3.0]) + + def test_arguments_are_keyword_only(self): + with self.assertRaises(TypeError): + dt.KalmanFilter(make_state([0.0, 1.0]), F, Q, H, R) + + def test_bad_shapes_raise(self): + state = make_state([0.0, 1.0]) + good = dict(f=F, q=Q, h=H, r=R) + + for key, bad in [("f", np.eye(3)), ("q", np.eye(3)), ("h", np.ones((1, 3))), + ("r", np.eye(2))]: + with self.subTest(key=key), self.assertRaises(ValueError): + dt.KalmanFilter(state, **{**good, key: bad}) + + with self.assertRaises(ValueError): + make_kf().update(dt.Observation(np.array([1.0, 2.0]), epoch=dt.Epoch(1))) + + def test_backward_predict_inverts_forward(self): + """Uses the Rust LinearTransition, which propagates backwards with F^-1.""" + state = make_state([0.0, 1.0], np.diag([2.0, 3.0])) + kf = make_kf(state) + value, cov = state.value.copy(), state.covariance.copy() + + kf.predict(dt.Epoch(3)) + out = kf.predict(dt.Epoch(0)) + np.testing.assert_allclose(out.value, value, atol=1e-12) + np.testing.assert_allclose(out.covariance, cov, atol=1e-12) + self.assertEqual(out.epoch.value, 0) + + def test_singular_update_is_logged(self): + singular = dt.KalmanFilter( + make_state([1.0, 1.0]), f=np.eye(2), q=np.zeros((2, 2)), + h=np.zeros((1, 2)), r=np.zeros((1, 1))) + with self.assertLogs("determined", level="ERROR") as logs: + singular.update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(0))) + self.assertIn("singular", logs.output[0]) + + +class TestCustomFilter(unittest.TestCase): + + def test_is_a_distinct_class_from_kalman_filter(self): + self.assertIsNot(dt.CustomFilter, dt.KalmanFilter) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests_py/test_measurement.py b/tests_py/test_measurement.py index e37b7e2..d44a509 100644 --- a/tests_py/test_measurement.py +++ b/tests_py/test_measurement.py @@ -13,7 +13,7 @@ def setUpClass(cls): cls.value = np.array([1.0, 0.0]) cls.epoch = dt.Epoch(0) - cls.state = dt.State(cls.value, np.eye(cls.value.size), cls.epoch) + cls.state = dt.State(cls.value, covariance=np.eye(cls.value.size), epoch=cls.epoch) h = np.array([ [1.0, 0.0], @@ -39,3 +39,16 @@ def test_jacobian(self): jac = self.measurement.jacobian(self.state) assert jac is not None, "Jacobian is None!" + + def test_noise(self): + np.testing.assert_array_equal(self.measurement.noise(), self.model.r) + + def test_inverse_may_be_none(self): + class NoInverse: + def projection(self, state): return dt.Observation(state.value, epoch=state.epoch) + def inverse(self, observation): return None + def noise(self): return np.eye(2) + + measurement = dt.MeasurementModel(NoInverse()) + obs = measurement.projection(self.state) + self.assertIsNone(measurement.inverse(obs)) diff --git a/tests_py/test_state.py b/tests_py/test_state.py index bb3c563..275b9a9 100644 --- a/tests_py/test_state.py +++ b/tests_py/test_state.py @@ -14,6 +14,14 @@ def test_constructor(self): ]) epoch = dt.Epoch(0) - state = dt.State(value, cov, epoch) + state = dt.State(value, covariance=cov, epoch=epoch) np.testing.assert_array_equal(value, state.value, "state value arrays not equal") + + def test_covariance_shape_is_checked(self): + with self.assertRaises(ValueError): + dt.State(np.zeros(2), covariance=np.eye(3), epoch=dt.Epoch(0)) + with self.assertRaises(ValueError): + dt.State(np.zeros(2), covariance=np.ones((2, 3)), epoch=dt.Epoch(0)) + # no covariance is still allowed + dt.State(np.zeros(2), covariance=None, epoch=dt.Epoch(0)) diff --git a/tests_py/test_transition.py b/tests_py/test_transition.py index a8e1eaa..f2f3504 100644 --- a/tests_py/test_transition.py +++ b/tests_py/test_transition.py @@ -14,7 +14,7 @@ def setUpClass(cls): cls.cov = np.eye(cls.value.size) cls.epoch = dt.Epoch(0) - cls.state = dt.State(cls.value, cls.cov, cls.epoch) + cls.state = dt.State(cls.value, covariance=cls.cov, epoch=cls.epoch) f = np.diag(np.array([2.0, 1.5])) q = np.eye(f.shape[0]) @@ -25,7 +25,7 @@ def setUpClass(cls): return super().setUpClass() def setUp(self): - self.transition = dt.TransitionModel(self.model, self.state) + self.transition = dt.TransitionModel(self.model, state=self.state) return super().setUp() def test_transition(self): @@ -36,3 +36,39 @@ def test_transition(self): def test_jacobian(self): jac = self.transition.jacobian(self.state) assert jac is not None, "Jacobian is None!" + + +class TestSharedState(unittest.TestCase): + """Python models may return the State object they were built with.""" + + def test_transition_returning_its_own_state(self): + class Static: + def __init__(self, state): self._state = state + def state(self, epoch): return self._state + def jacobian(self, state): return np.eye(2) + + state = dt.State(np.array([1.0, 2.0]), covariance=np.eye(2), epoch=dt.Epoch(0)) + transition = dt.TransitionModel(Static(state), state=state) + out = transition.state(dt.Epoch(1)) + np.testing.assert_array_equal(out.value, [1.0, 2.0]) + self.assertEqual(out.epoch.value, 1) + + def test_update_returning_its_own_state(self): + class Static: + def __init__(self, state): self._state = state + def state(self, epoch): return self._state + def jacobian(self, state): return np.eye(2) + + class Measurement: + def projection(self, s): return dt.Observation(s.value, epoch=s.epoch) + def noise(self): return np.eye(2) + + class Update: + def __init__(self, state): self._state = state + def apply(self, observation): return self._state + + state = dt.State(np.array([1.0, 2.0]), covariance=np.eye(2), epoch=dt.Epoch(0)) + transition = dt.TransitionModel(Static(state), state=state) + update = dt.UpdateModel(Update(state), transition=transition, measurement=dt.MeasurementModel(Measurement())) + out = update.apply(dt.Observation(np.array([0.0, 0.0]), epoch=dt.Epoch(0))) + np.testing.assert_array_equal(out.value, [1.0, 2.0]) diff --git a/tests_py/test_ud.py b/tests_py/test_ud.py new file mode 100644 index 0000000..6882e33 --- /dev/null +++ b/tests_py/test_ud.py @@ -0,0 +1,154 @@ +import unittest +import numpy as np + +import determined as dt + + +def make_state(value, cov=None, epoch=0): + value = np.asarray(value, dtype=float) + cov = np.eye(value.size) if cov is None else cov + return dt.State(value, covariance=cov, epoch=dt.Epoch(epoch)) + + +P0 = np.array([[2.0, 0.3, 0.1], [0.3, 1.0, -0.2], [0.1, -0.2, 0.5]]) +F = np.array([[1.0, 0.1, 0.0], [0.0, 1.0, 0.1], [0.0, 0.0, 0.95]]) +Q = np.array([[0.02, 0.005, 0.0], [0.005, 0.01, 0.002], [0.0, 0.002, 0.03]]) +H = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 1.0]]) +R_CORRELATED = np.array([[0.1, 0.04], [0.04, 0.2]]) + + +class TestUDKalmanFilter(unittest.TestCase): + + def test_matches_kalman_filter(self): + """Same estimates as KalmanFilter, for diagonal and correlated R.""" + for r in [np.diag([0.1, 0.2]), R_CORRELATED]: + kf = dt.KalmanFilter(make_state([1.0, -0.5, 2.0], P0), f=F, q=Q, h=H, r=r) + ud = dt.UDKalmanFilter(make_state([1.0, -0.5, 2.0], P0), f=F, q=Q, h=H, r=r) + + for t in [1, 2, 5, 6, 9]: + obs = dt.Observation(np.array([0.3 * t, -0.1 * t]), epoch=dt.Epoch(t)) + a, b = kf.update(obs), ud.update(obs) + np.testing.assert_allclose(b.value, a.value, rtol=1e-9, atol=1e-12) + np.testing.assert_allclose(b.covariance, a.covariance, rtol=1e-9, atol=1e-12) + self.assertEqual(a.epoch.value, b.epoch.value) + + def test_covariance_is_symmetric(self): + state = make_state([0.0, 0.0], np.diag([1e8, 1.0])) + ud = dt.UDKalmanFilter(state, f=np.eye(2), q=np.zeros((2, 2)), + h=np.array([[1.0, 1.0]]), r=np.array([[1e-8]])) + for t in range(1, 50): + cov = ud.update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(t))).covariance + np.testing.assert_array_equal(cov, cov.T) + self.assertGreaterEqual(np.linalg.eigvalsh(cov).min(), -1e-9 * np.abs(cov).max()) + + def test_state_setter_refactorises(self): + """After the state is replaced, the filter continues from it like a fresh KalmanFilter.""" + ud = dt.UDKalmanFilter(make_state([1.0, -0.5, 2.0], P0), f=F, q=Q, h=H, r=R_CORRELATED) + ud.update(dt.Observation(np.array([1.0, 1.0]), epoch=dt.Epoch(1))) + + ud.state = make_state([0.5, 0.5, 0.5], np.diag([3.0, 2.0, 1.0]), epoch=1) + kf = dt.KalmanFilter(make_state([0.5, 0.5, 0.5], np.diag([3.0, 2.0, 1.0]), epoch=1), + f=F, q=Q, h=H, r=R_CORRELATED) + + obs = dt.Observation(np.array([0.2, -0.3]), epoch=dt.Epoch(3)) + a, b = kf.update(obs), ud.update(obs) + np.testing.assert_allclose(b.value, a.value, rtol=1e-9, atol=1e-12) + np.testing.assert_allclose(b.covariance, a.covariance, rtol=1e-9, atol=1e-12) + + def test_backward_predict_is_not_supported(self): + state = make_state([1.0, -0.5, 2.0], P0, epoch=5) + ud = dt.UDKalmanFilter(state, f=F, q=Q, h=H, r=R_CORRELATED) + with self.assertLogs("determined", level="WARNING") as logs: + out = ud.predict(dt.Epoch(2)) + self.assertIn("Backward", logs.output[0]) + self.assertEqual(out.epoch.value, 5) + + def test_reset(self): + ud = dt.UDKalmanFilter(make_state([1.0, -0.5, 2.0], P0), f=F, q=Q, h=H, r=R_CORRELATED) + ud.reset() + np.testing.assert_array_equal(ud.state.value, np.zeros(3)) + np.testing.assert_array_equal(ud.state.covariance, np.eye(3)) + # continues from the reset state + kf = dt.KalmanFilter(make_state(np.zeros(3)), f=F, q=Q, h=H, r=R_CORRELATED) + obs = dt.Observation(np.array([0.2, -0.3]), epoch=dt.Epoch(1)) + np.testing.assert_allclose(ud.update(obs).value, kf.update(obs).value, rtol=1e-9, atol=1e-12) + + def test_bad_inputs_raise(self): + state = make_state([1.0, -0.5, 2.0], P0) + good = dict(f=F, q=Q, h=H, r=R_CORRELATED) + with self.assertRaises(TypeError): + dt.UDKalmanFilter(state, F, Q, H, R_CORRELATED) + for key, bad in [("f", np.eye(2)), ("q", np.eye(2)), ("h", np.ones((2, 2))), ("r", np.eye(3))]: + with self.subTest(key=key), self.assertRaises(ValueError): + dt.UDKalmanFilter(state, **{**good, key: bad}) + with self.assertRaises(ValueError): + dt.UDKalmanFilter(state, **good).update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(1))) + + +def nonlinear_models(): + f = lambda x, t: np.array([x[0] + 0.1 * np.sin(x[1]), 0.98 * x[1]]) + jac_f = lambda x, t: np.array([[1.0, 0.1 * np.cos(x[1])], [0.0, 0.98]]) + h = lambda x, t: np.array([np.hypot(x[0], x[1]), np.arctan2(x[1], x[0])]) + + def jac_h(x, t): + r2 = x @ x + r = np.sqrt(r2) + return np.array([[x[0] / r, x[1] / r], [-x[1] / r2, x[0] / r2]]) + + return dict(f=f, jac_f=jac_f, q=np.eye(2) * 0.01, h=h, jac_h=jac_h, + r=np.array([[0.05, 0.01], [0.01, 0.02]])) + + +class TestUDExtendedKalmanFilter(unittest.TestCase): + + def test_matches_extended_kalman_filter(self): + p = np.array([[1.0, 0.2], [0.2, 1.0]]) + ekf = dt.ExtendedKalmanFilter(make_state([2.0, 2.0], p), **nonlinear_models()) + ud = dt.UDExtendedKalmanFilter(make_state([2.0, 2.0], p), **nonlinear_models()) + + for t in [1, 2, 4, 5]: + obs = dt.Observation(np.array([3.0, 0.8]), epoch=dt.Epoch(t)) + a, b = ekf.update(obs), ud.update(obs) + np.testing.assert_allclose(b.value, a.value, rtol=1e-9, atol=1e-12) + np.testing.assert_allclose(b.covariance, a.covariance, rtol=1e-9, atol=1e-12) + + def test_error_in_callback_leaves_state_unchanged_and_filter_usable(self): + models = nonlinear_models() + calls = {"n": 0} + h_ok = models["h"] + + def h(x, t): + calls["n"] += 1 + if calls["n"] == 2: + raise RuntimeError("boom") + return h_ok(x, t) + + p = np.array([[1.0, 0.2], [0.2, 1.0]]) + state = make_state([2.0, 2.0], p) + ud = dt.UDExtendedKalmanFilter(state, **{**models, "h": h}) + ekf = dt.ExtendedKalmanFilter(make_state([2.0, 2.0], p), **models) + + obs1 = dt.Observation(np.array([3.0, 0.8]), epoch=dt.Epoch(1)) + ud.update(obs1) + ekf.update(obs1) + before = (state.value.copy(), state.covariance.copy(), state.epoch.value) + + with self.assertRaises(RuntimeError): + ud.update(dt.Observation(np.array([3.1, 0.7]), epoch=dt.Epoch(2))) + np.testing.assert_array_equal(state.value, before[0]) + np.testing.assert_array_equal(state.covariance, before[1]) + self.assertEqual(state.epoch.value, before[2]) + + # the factors were restored with the state: it continues like the EKF + obs3 = dt.Observation(np.array([3.2, 0.75]), epoch=dt.Epoch(3)) + np.testing.assert_allclose(ud.update(obs3).value, ekf.update(obs3).value, rtol=1e-9, atol=1e-12) + + def test_arguments_are_keyword_only(self): + m = nonlinear_models() + with self.assertRaises(TypeError): + dt.UDExtendedKalmanFilter(make_state([2.0, 2.0]), m["f"], m["jac_f"], m["q"], + m["h"], m["jac_h"], m["r"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests_py/test_ukf.py b/tests_py/test_ukf.py new file mode 100644 index 0000000..9d2dc8a --- /dev/null +++ b/tests_py/test_ukf.py @@ -0,0 +1,171 @@ +import unittest +import numpy as np + +import determined as dt + + +def make_state(value, cov=None, epoch=0): + value = np.asarray(value, dtype=float) + cov = np.eye(value.size) if cov is None else cov + return dt.State(value, covariance=cov, epoch=dt.Epoch(epoch)) + + +def reference_ukf_step(x, p, f, q, alpha, beta, kappa): + """Independent numpy unscented transform through f, plus q.""" + n = x.size + lam = alpha ** 2 * (n + kappa) - n + l = np.linalg.cholesky((n + lam) * p) + sigma = np.vstack([x, x + l.T, x - l.T]) + wm = np.full(2 * n + 1, 1 / (2 * (n + lam))) + wc = wm.copy() + wm[0] = lam / (n + lam) + wc[0] = wm[0] + 1 - alpha ** 2 + beta + return sigma, wm, wc + + +def reference_predict(x, p, f, q, t, params): + sigma, wm, wc = reference_ukf_step(x, p, f, q, *params) + y = np.array([f(s, t) for s in sigma]) + mean = wm @ y + d = y - mean + return mean, (wc * d.T) @ d + q + + +def reference_update(x, p, z, h, r, t, params): + sigma, wm, wc = reference_ukf_step(x, p, None, None, *params) + zs = np.array([h(s, t) for s in sigma]) + zm = wm @ zs + dz, dx = zs - zm, sigma - x + s = (wc * dz.T) @ dz + r + pxz = (wc * dx.T) @ dz + k = pxz @ np.linalg.inv(s) + return x + k @ (z - zm), p - k @ s @ k.T + + +class TestUnscentedKalmanFilter(unittest.TestCase): + + def test_matches_numpy_reference(self): + """Nonlinear models, process noise, epoch gaps, non-default parameters.""" + params = (0.5, 2.0, 1.0) + f = lambda x, t: np.array([x[0] + 0.1 * np.sin(x[1]), 0.95 * x[1] + 0.01 * t]) + h = lambda x, t: np.array([np.hypot(x[0], x[1]), np.arctan2(x[1], x[0])]) + q, r = np.eye(2) * 0.01, np.diag([0.1, 0.01]) + + x, p = np.array([1.0, 0.5]), np.eye(2) * 0.3 + ukf = dt.UnscentedKalmanFilter( + make_state(x, p), f=f, q=q, h=h, r=r, + alpha=params[0], beta=params[1], kappa=params[2]) + + epoch = 0 + for gap in [1, 2, 1, 3]: + for _ in range(gap): + x, p = reference_predict(x, p, f, q, epoch, params) + epoch += 1 + z = np.array([1.2, 0.3 + 0.01 * epoch]) + x, p = reference_update(x, p, z, h, r, epoch, params) + + out = ukf.update(dt.Observation(z, epoch=dt.Epoch(epoch))) + np.testing.assert_allclose(out.value, x, atol=1e-10) + np.testing.assert_allclose(out.covariance, (p + p.T) / 2, atol=1e-10) + self.assertEqual(out.epoch.value, epoch) + + def test_matches_kalman_filter_on_linear_models(self): + f_mat = np.array([[1.0, 0.5], [0.0, 1.0]]) + h_mat = np.array([[1.0, 0.0]]) + q, r = np.eye(2) * 0.05, np.array([[0.1]]) + + kf = dt.KalmanFilter(make_state([0.0, 1.0]), f=f_mat, q=q, h=h_mat, r=r) + ukf = dt.UnscentedKalmanFilter( + make_state([0.0, 1.0]), f=lambda x, t: f_mat @ x, q=q, h=lambda x, t: h_mat @ x, r=r) + + for t in [1, 2, 4, 5]: + obs = dt.Observation(np.array([0.4 * t]), epoch=dt.Epoch(t)) + a, b = kf.update(obs), ukf.update(obs) + np.testing.assert_allclose(a.value, b.value, atol=1e-10) + np.testing.assert_allclose(a.covariance, b.covariance, atol=1e-10) + + def test_transform_recovers_exact_moments_of_square(self): + """x ~ N(m, s2), f(x) = x^2: mean m^2 + s2, variance 4 m^2 s2 + 2 s2^2.""" + m, s2 = 2.0, 0.5 + ukf = dt.UnscentedKalmanFilter( + make_state([m], np.array([[s2]])), + f=lambda x, t: x ** 2, q=np.zeros((1, 1)), h=lambda x, t: x, r=np.eye(1)) + + out = ukf.predict(dt.Epoch(1)) + np.testing.assert_allclose(out.value, [m * m + s2]) + np.testing.assert_allclose(out.covariance, [[4 * m * m * s2 + 2 * s2 * s2]]) + + def test_models_receive_epoch_at_every_sigma_point(self): + calls = [] + + def f(x, t): + calls.append(t) + return x + + ukf = dt.UnscentedKalmanFilter( + make_state([0.0, 0.0]), f=f, q=np.eye(2), h=lambda x, t: x, r=np.eye(2)) + ukf.predict(dt.Epoch(2)) + self.assertEqual(calls, [0] * 5 + [1] * 5) # 2n + 1 = 5 sigma points per step + + def test_error_in_callback_leaves_state_unchanged(self): + def h(x, t): + raise RuntimeError("boom") + + state = make_state([1.0, 2.0]) + ukf = dt.UnscentedKalmanFilter( + state, f=lambda x, t: x + 1.0, q=np.eye(2), h=h, r=np.eye(2)) + + with self.assertRaises(RuntimeError): + ukf.update(dt.Observation(np.array([0.0, 0.0]), epoch=dt.Epoch(1))) + + # the predict step inside update was rolled back too + np.testing.assert_array_equal(state.value, [1.0, 2.0]) + np.testing.assert_array_equal(state.covariance, np.eye(2)) + self.assertEqual(state.epoch.value, 0) + + def test_not_positive_definite_covariance_is_logged(self): + state = make_state([1.0, 1.0], -np.eye(2)) + ukf = dt.UnscentedKalmanFilter( + state, f=lambda x, t: 2 * x, q=np.zeros((2, 2)), h=lambda x, t: x, r=np.eye(2)) + + with self.assertLogs("determined", level="ERROR") as logs: + out = ukf.predict(dt.Epoch(1)) + self.assertIn("positive definite", logs.output[0]) + np.testing.assert_array_equal(out.value, [1.0, 1.0]) + self.assertEqual(out.epoch.value, 0) + + def test_reset_and_state_setter(self): + ukf = dt.UnscentedKalmanFilter( + make_state([1.0, 1.0], np.eye(2) * 3), f=lambda x, t: x, q=np.eye(2), + h=lambda x, t: x, r=np.eye(2)) + + ukf.reset() + np.testing.assert_array_equal(ukf.state.value, [0.0, 0.0]) + np.testing.assert_array_equal(ukf.state.covariance, np.eye(2)) + + ukf.state = make_state([5.0, 6.0], epoch=3) + np.testing.assert_array_equal(ukf.state.value, [5.0, 6.0]) + self.assertEqual(ukf.state.epoch.value, 3) + + def test_bad_inputs_raise(self): + state, eye = make_state([1.0, 1.0]), np.eye(2) + + def build(**overrides): + args = dict(f=lambda x, t: x, q=eye, h=lambda x, t: x, r=eye) + args.update(overrides) + return dt.UnscentedKalmanFilter(state, **args) + + with self.assertRaises(TypeError): + dt.UnscentedKalmanFilter(state, lambda x, t: x, eye, lambda x, t: x, eye) + with self.assertRaises(TypeError): + build(h=None) + with self.assertRaises(ValueError): + build(q=np.eye(3)) + with self.assertRaises(ValueError): + build(f=lambda x, t: np.zeros(3)).predict(dt.Epoch(1)) + with self.assertRaises(ValueError): + build().update(dt.Observation(np.array([1.0]), epoch=dt.Epoch(1))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests_py/test_update.py b/tests_py/test_update.py index 8b34994..66c04fd 100644 --- a/tests_py/test_update.py +++ b/tests_py/test_update.py @@ -17,7 +17,7 @@ def setUpClass(cls): cls.value = np.array([1.0, 1.0]) cls.cov = np.eye(cls.value.size) cls.epoch = dt.Epoch(0) - cls.obs = dt.Observation(np.array([0.0, 1.0]), cls.epoch) + cls.obs = dt.Observation(np.array([0.0, 1.0]), epoch=cls.epoch) cls.f = np.diag(np.array([2.0, 1.5])) cls.q = np.eye(cls.f.shape[0]) @@ -32,16 +32,16 @@ def setUpClass(cls): # Rebuild models at each test to avoid persistent state def setUp(self): - self.state = dt.State(self.value, self.cov, self.epoch) + self.state = dt.State(self.value, covariance=self.cov, epoch=self.epoch) self.t = LinearTransition(self.f, self.q, self.state) - self.transition = dt.TransitionModel(self.t, self.state) + self.transition = dt.TransitionModel(self.t, state=self.state) self.m = LinearMeasurement(self.h, self.r) self.measurement = dt.MeasurementModel(self.m) self.model = LinearUpdate(self.transition, self.measurement) - self.update = dt.UpdateModel(self.model, self.transition, self.measurement) + self.update = dt.UpdateModel(self.model, transition=self.transition, measurement=self.measurement) return super().setUp() def test_predict(self): @@ -52,7 +52,3 @@ def test_predict(self): def test_apply(self): new_state = self.update.apply(self.obs) assert isinstance(new_state, dt.State), "state should be dt.State" - - def test_jacobian(self): - jac = self.update.jacobian(self.state) - assert jac is not None, "Jacobian is None!"