Particle filter — 512 particles, systematic resample

//! Bootstrap particle filter on a 1-D sinusoidal target with Gaussian
//! observation noise.
//!
//! Model:
//!   true:  x_t = A sin(ω t)
//!   obs:   y_t = x_t + ε,  ε ~ N(0, σ²)
//!   prop:  x_t = x_{t-1} + step + N(0, q²)   (random walk + drift)
//!   wt:    w ∝ exp(-(x - y)² / 2σ²)
//!   ess:   resample (systematic) when ESS < N/2.
//!
//! Deterministic given seed. V-class L0-closed.

use serde::{Deserialize, Serialize};

use crate::rng::SplitMix64;

#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct PFilterParams {
    pub n_particles: u32,
    pub amplitude: f32,
    pub omega: f32,
    pub obs_sigma: f32,
    pub proc_sigma: f32,
}

impl Default for PFilterParams {
    fn default() -> Self {
        Self {
            n_particles: 512,
            amplitude: 1.5,
            omega: 0.15,
            obs_sigma: 0.4,
            proc_sigma: 0.15,
        }
    }
}

pub struct PFilterState {
    pub rng: SplitMix64,
    pub particles: Vec<f32>,
    pub weights: Vec<f32>,
    pub t: f32,
    pub posterior_mean: f32,
    pub ess: f32,
}

impl PFilterState {
    pub fn new(p: &PFilterParams, seed: u64) -> Self {
        let n = p.n_particles as usize;
        let mut rng = SplitMix64::new(seed);
        let particles: Vec<f32> = (0..n)
            .map(|_| p.amplitude * rng.next_normal() * 0.2)
            .collect();
        let weights = vec![1.0 / n as f32; n];
        Self { rng, particles, weights, t: 0.0, posterior_mean: 0.0, ess: n as f32 }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PFilterStepReport {
    pub bit_ops: u64,
    pub ess: f32,
    pub resampled: bool,
    /// True target state at this step.
    pub truth: f32,
    /// Observation drawn for this step.
    pub obs: f32,
}

/// One filter step. Mutates `state`; returns telemetry.
pub fn pfilter_step(state: &mut PFilterState, p: &PFilterParams, dt: f32) -> PFilterStepReport {
    state.t += dt;
    let truth = p.amplitude * (p.omega * state.t).sin();
    let obs = truth + p.obs_sigma * state.rng.next_normal();
    let drift = p.amplitude * p.omega * (p.omega * state.t).cos() * dt;

    // ── Propagate ──
    for x in state.particles.iter_mut() {
        *x += drift + p.proc_sigma * state.rng.next_normal();
    }

    // ── Weight ──
    let inv2s2 = 1.0 / (2.0 * p.obs_sigma * p.obs_sigma);
    let mut w_sum = 0.0f32;
    for (x, w) in state.particles.iter().zip(state.weights.iter_mut()) {
        let d = x - obs;
        *w = (-d * d * inv2s2).exp();
        w_sum += *w;
    }
    if w_sum > 0.0 {
        for w in state.weights.iter_mut() {
            *w /= w_sum;
        }
    } else {
        let n = state.weights.len() as f32;
        for w in state.weights.iter_mut() {
            *w = 1.0 / n;
        }
    }

    // ── ESS ──
    let ess: f32 = 1.0 / state.weights.iter().map(|w| w * w).sum::<f32>();
    state.ess = ess;

    // ── Posterior mean ──
    state.posterior_mean = state
        .particles
        .iter()
        .zip(state.weights.iter())
        .map(|(x, w)| x * w)
        .sum();

    // ── Resample (systematic) when ESS < N/2 ──
    let n = state.particles.len();
    let mut resampled = false;
    if ess < n as f32 * 0.5 {
        let mut new_particles = Vec::with_capacity(n);
        let u0 = state.rng.next_f32() / n as f32;
        let mut c = state.weights[0];
        let mut i = 0usize;
        for j in 0..n {
            let u = u0 + j as f32 / n as f32;
            while u > c && i + 1 < n {
                i += 1;
                c += state.weights[i];
            }
            new_particles.push(state.particles[i]);
        }
        state.particles = new_particles;
        let inv_n = 1.0 / n as f32;
        for w in state.weights.iter_mut() {
            *w = inv_n;
        }
        resampled = true;
    }

    // Bit-ops: O(N) for propagate, weight, ESS, mean, plus an optional
    // O(N) resample. Conservatively 8 f32 ops/particle, possibly doubled.
    let mult = if resampled { 16 } else { 8 };
    let bit_ops = (n as u64) * mult * 32;

    PFilterStepReport { bit_ops, ess, resampled, truth, obs }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn posterior_mean_tracks_truth() {
        let p = PFilterParams::default();
        let mut s = PFilterState::new(&p, 7);
        let mut err = 0.0f32;
        let n_steps = 500;
        for _ in 0..n_steps {
            let r = pfilter_step(&mut s, &p, 0.1);
            err += (s.posterior_mean - r.truth).abs();
        }
        let mae = err / n_steps as f32;
        // Should track within roughly one obs-sigma.
        assert!(mae < p.obs_sigma * 1.5, "MAE = {mae}, expected < {}", p.obs_sigma * 1.5);
    }

    #[test]
    fn resample_fires_when_ess_drops() {
        let p = PFilterParams::default();
        let mut s = PFilterState::new(&p, 1);
        let mut any_resample = false;
        for _ in 0..200 {
            let r = pfilter_step(&mut s, &p, 0.1);
            if r.resampled { any_resample = true; }
        }
        assert!(any_resample, "expected at least one resample in 200 steps");
    }

    #[test]
    fn weights_are_normalised() {
        let p = PFilterParams::default();
        let mut s = PFilterState::new(&p, 3);
        for _ in 0..100 {
            pfilter_step(&mut s, &p, 0.1);
            let sum: f32 = s.weights.iter().sum();
            assert!((sum - 1.0).abs() < 1e-3, "weights sum {sum}");
        }
    }
}

This is the exact Rust file compiled into the WASM module the page just loaded. Every line that runs on your device is here. The receipt is a function of this code, not a bespoke benchmark.