edgefirst_tracker/

kalman.rs

// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
// SPDX-License-Identifier: Apache-2.0

use nalgebra::{
    allocator::Allocator, convert, dimension::U4, DVector, DefaultAllocator, Dyn, OMatrix,
    RealField, SVector, U1, U8,
};

/// Constant-velocity Kalman filter for tracking a bounding box in XYAH space.
///
/// The state vector is `[x, y, a, h, ẋ, ẏ, ȧ, ḣ]` (8 dimensions):
///
/// | Component | Meaning |
/// |-----------|---------|
/// | `x` | Centre x coordinate |
/// | `y` | Centre y coordinate |
/// | `a` | Aspect ratio (`width / height`), kept > 0 by a small epsilon floor |
/// | `h` | Height |
/// | `ẋ, ẏ, ȧ, ḣ` | First derivatives of the position components |
///
/// Measurements arriving via [`update`](Self::update) are in 4-dimensional
/// XYAH space (no velocity), projected from the 8-D state by the observation
/// matrix `H`.  The filter uses a scale-adaptive process-noise covariance:
/// noise magnitude scales proportionally with the box height `h`, so
/// fast-moving large objects tolerate wider uncertainty than small static ones.
///
/// This is the same Kalman filter parameterisation used in the original
/// ByteTrack paper (Zhang et al., 2022).
#[derive(Debug, Clone)]
pub struct ConstantVelocityXYAHModel2<R>
where
    R: RealField,
    DefaultAllocator: Allocator<U8, U8>,
    DefaultAllocator: Allocator<U8>,
{
    /// Current state estimate: `[x, y, a, h, ẋ, ẏ, ȧ, ḣ]`.
    pub mean: SVector<R, 8>,

    /// Scale factor for position noise (`1/20` by default). Process noise for
    /// position components is `std_weight_position * h`.
    pub std_weight_position: R,

    /// Scale factor for velocity noise (`1/160` by default). Process noise for
    /// velocity components is `std_weight_velocity * h`.
    pub std_weight_velocity: R,

    /// Measurement-to-innovation scaling applied in [`update`](Self::update).
    /// Set via [`ByteTrackBuilder::track_update`](crate::bytetrack::ByteTrackBuilder::track_update).
    /// Values in `(0, 1]`; lower → smoother (trusts prediction more).
    pub update_factor: R,

    /// State-transition matrix `F` (identity + velocity block). Projects the
    /// state forward by one time step.
    motion_matrix: OMatrix<R, U8, U8>,

    /// Observation matrix `H`. Extracts the `[x, y, a, h]` slice from the
    /// 8-D state for comparison against incoming measurements.
    update_matrix: OMatrix<R, U4, U8>,

    /// Current state covariance matrix `P` (8×8). Grows on each
    /// [`predict`](Self::predict) and shrinks on each [`update`](Self::update).
    pub covariance: OMatrix<R, U8, U8>,
}

/// Distance metric used by [`ConstantVelocityXYAHModel2::gating_distance`].
///
/// `Mahalanobis` accounts for the current state uncertainty and is the
/// theoretically correct gating criterion; `Gaussian` (squared Euclidean) is
/// faster but scale-dependent.  In the current ByteTrack integration only the
/// IoU-based cost matrix is used for assignment; this enum is available for
/// future association strategies that want a filter-based gate.
#[allow(dead_code)]
pub enum GatingDistanceMetric {
    /// Squared Euclidean distance (no covariance weighting).
    Gaussian,
    /// Mahalanobis distance — normalised by the projected state covariance.
    Mahalanobis,
}

impl<R> ConstantVelocityXYAHModel2<R>
where
    R: RealField + Copy,
{
    /// Initialise the filter from the first observation.
    ///
    /// `measurement` is `[x, y, a, h]` in XYAH space.  Velocity components are
    /// initialised to zero.  The initial covariance is set proportionally to
    /// the box height so larger boxes start with wider uncertainty.
    ///
    /// `update_factor` is stored and applied as a measurement gain in
    /// [`update`](Self::update).
    pub fn new(measurement: &[R; 4], update_factor: R) -> Self {
        let ndim = 4;
        let dt: R = convert(1.0);

        let mut motion_matrix = OMatrix::<R, U8, U8>::identity();
        for i in 0..ndim {
            motion_matrix[(i, ndim + i)] = dt * convert(3.0);
        }
        let mut update_matrix = OMatrix::<R, U4, U8>::identity();
        for i in 0..ndim {
            update_matrix[(i, ndim + i)] = dt * convert(1.0);
        }
        let zero: R = convert(0.0);
        let two: R = convert(2.0);
        let ten: R = convert(10.0);
        let height = measurement[3];

        let mean = SVector::<R, 8>::from_row_slice(&[
            measurement[0],
            measurement[1],
            measurement[2],
            measurement[3],
            zero,
            zero,
            zero,
            zero,
        ]);
        let std_weight_position = convert(1.0 / 20.0);
        let std_weight_velocity = convert(1.0 / 160.0);
        let diag = [
            two * std_weight_position * height,
            two * std_weight_position * height,
            convert(0.01),
            two * std_weight_position * height,
            ten * std_weight_velocity * height,
            ten * std_weight_velocity * height,
            convert(0.00001),
            ten * std_weight_velocity * height,
        ];
        let diag = SVector::<R, 8>::from_row_slice(&diag);

        let covariance = OMatrix::<R, U8, U8>::from_diagonal(&diag.component_mul(&diag));
        Self {
            motion_matrix,
            update_matrix,
            mean,
            covariance,
            std_weight_position,
            std_weight_velocity,
            update_factor,
        }
    }

    /// Advance the filter by one time step with no measurement.
    ///
    /// Applies the constant-velocity motion model (`mean' = F * mean`,
    /// `cov' = F * cov * Fᵀ + Q`) using a scale-adaptive process-noise
    /// covariance `Q` proportional to the current box height.  Called once
    /// per frame for every active tracklet before the assignment step.
    pub fn predict(&mut self) {
        let height = self.mean[3];
        let diag = [
            self.std_weight_position * height,
            self.std_weight_position * height,
            convert(0.01),
            self.std_weight_position * height,
            self.std_weight_velocity * height,
            self.std_weight_velocity * height,
            convert(0.00001),
            self.std_weight_velocity * height,
        ];
        let diag = SVector::<R, 8>::from_row_slice(&diag);
        let motion_cov = OMatrix::<R, U8, U8>::from_diagonal(&diag.component_mul(&diag));

        let mean = (self.mean.transpose() * self.motion_matrix.transpose()).transpose();
        let covariance =
            self.motion_matrix * self.covariance * self.motion_matrix.transpose() + motion_cov;
        self.mean = mean;
        self.covariance = covariance;
    }

    /// Project the 8-D state into 4-D measurement space.
    ///
    /// Returns `(projected_mean, projected_covariance)` — the expected
    /// measurement and its uncertainty under the current state estimate.
    /// Used by [`update`](Self::update) and by
    /// [`gating_distance`](Self::gating_distance).
    pub fn project(&self) -> (OMatrix<R, U4, U1>, OMatrix<R, U4, U4>) {
        let height = self.mean[3];
        let diag = [
            self.std_weight_position * height,
            self.std_weight_position * height,
            convert(0.01),
            self.std_weight_position * height,
        ];
        let diag = SVector::<R, 4>::from_row_slice(&diag);
        let innovation_cov = OMatrix::<R, U4, U4>::from_diagonal(&diag.component_mul(&diag));
        let mean = self.update_matrix * self.mean;
        let covariance =
            self.update_matrix * self.covariance * self.update_matrix.transpose() + innovation_cov;
        (mean, covariance)
    }

    /// Correct the state estimate with a new measurement.
    ///
    /// `measurement` is `[x, y, a, h]` in XYAH space, as produced by
    /// [`xyxy_to_xyah`](super::bytetrack) in the ByteTrack update loop.
    ///
    /// The Kalman gain is computed via Cholesky decomposition of the projected
    /// covariance.  If the decomposition fails (degenerate covariance) the
    /// update is silently skipped and the prior state is preserved.
    ///
    /// The innovation is scaled by `self.update_factor` before being applied,
    /// which effectively attenuates the correction — lower values produce
    /// smoother trajectories at the cost of slower response to abrupt motion.
    pub fn update(&mut self, measurement: &[R; 4]) {
        let measurement = SVector::<R, 4>::from_row_slice(&[
            measurement[0],
            measurement[1],
            measurement[2],
            measurement[3],
        ]);

        let (projected_mean, projected_cov) = self.project();
        let cho_factor = match projected_cov.cholesky() {
            None => return,
            Some(v) => v,
        };
        let kalman_gain = cho_factor
            .solve(&(self.covariance * self.update_matrix.transpose()).transpose())
            .transpose();

        let innovation = (measurement - projected_mean).scale(self.update_factor);
        // println!("kalman_gain={}", kalman_gain);
        // println!("innovation={}", innovation);
        let diff = innovation.transpose() * kalman_gain.transpose();
        self.mean += diff.transpose();
        self.covariance -= kalman_gain * projected_cov * kalman_gain.transpose();
        // let new_mean = self.mean + diff.transpose();
        // let new_cov = self.covariance - kalman_gain * projected_cov *
        // kalman_gain.transpose();
    }

    /// Compute per-measurement gating distances from the current state.
    ///
    /// Returns a vector of one scalar per row in `measurements`.  Each scalar
    /// is either the squared Euclidean distance ([`GatingDistanceMetric::Gaussian`])
    /// or the Mahalanobis distance ([`GatingDistanceMetric::Mahalanobis`]) from
    /// the projected state mean to that measurement.
    ///
    /// `only_position` restricts the comparison to the first two components
    /// (`x`, `y`), ignoring aspect and height.
    ///
    /// Not used by the current ByteTrack integration (which gates via IoU);
    /// exposed for future association strategies.
    #[allow(dead_code)]
    pub fn gating_distance(
        &self,
        measurements: &OMatrix<R, Dyn, U4>,
        only_position: bool,
        metric: GatingDistanceMetric,
    ) -> DVector<R> {
        let (m, cov) = self.project();
        let ndims = if only_position { 2 } else { 4 };
        let mean = m.transpose();
        let mean = mean.columns_range(0..ndims);
        let covariance = cov.view_range(0..ndims, 0..ndims);
        let measurements = measurements.columns_range(0..ndims);
        // let _ = only_position;
        // let mean = m.transpose();
        // let covariance = cov;
        // let measurements = measurements;

        let mut mean_broadcast =
            OMatrix::<R, Dyn, U4>::from_element(measurements.shape().0, convert(0.0));
        for mut col in mean_broadcast.row_iter_mut() {
            col.copy_from(&mean);
        }
        let d = measurements - mean_broadcast;
        match metric {
            GatingDistanceMetric::Gaussian => d.component_mul(&d).column_sum(),
            GatingDistanceMetric::Mahalanobis => {
                let cho_factor = match covariance.cholesky() {
                    None => return DVector::<R>::zeros(measurements.shape().0),
                    Some(v) => v,
                };
                let z = cho_factor.solve(&d.transpose());
                z.component_mul(&z).row_sum_tr()
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use nalgebra::{Dyn, OMatrix, U4};

    use super::{ConstantVelocityXYAHModel2, GatingDistanceMetric};
    #[test]
    fn filter() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        t.predict();
        println!("1. t.mean={}", t.mean);
        t.update(&[0.4, 0.5, 1.0, 0.5]);
        t.predict();
        println!("2. t.mean={}", t.mean);
        t.update(&[0.3, 0.5, 1.0, 0.5]);
        t.predict();
        println!("3. t.mean={}", t.mean);
        t.update(&[0.2, 0.5, 1.0, 0.5]);
        t.predict();
        println!("4. t.mean={}", t.mean);
        t.update(&[0.2, 0.5, 1.0, 0.5]);
        t.predict();
        println!("5. t.mean={}", t.mean);
        t.update(&[0.3, 0.5, 1.0, 0.5]);
        t.predict();
        println!("6. t.mean={}", t.mean);
        t.update(&[0.4, 0.5, 1.0, 0.5]);
    }

    #[test]
    fn gating() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        t.predict();
        t.update(&[0.49, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.48, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.47, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.46, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.45, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.44, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.43, 0.5, 1.0, 0.5]);
        t.predict();
        t.update(&[0.42, 0.5, 1.0, 0.5]);
        t.predict();

        // distances range from 0 to 1e6 for maha
        let mut measurements = OMatrix::<f32, Dyn, U4>::from_element(1, 0.0);
        measurements.copy_from_slice(&[0.3, 0.5, 1.0, 0.5]);

        let mut distances = OMatrix::<f32, Dyn, Dyn>::from_element(1, 1, 0.0);
        for mut column in distances.column_iter_mut() {
            let dist = t.gating_distance(&measurements, false, GatingDistanceMetric::Gaussian);
            column.copy_from(&dist);
        }
        let dist = t.gating_distance(&measurements, false, GatingDistanceMetric::Mahalanobis);
        println!("Dist(false, maha): {dist}");

        let dist = t.gating_distance(&measurements, false, GatingDistanceMetric::Gaussian);
        println!("Dist(false, gaussian): {dist}");
    }

    #[test]
    fn test_predict_constant_velocity() {
        // Initialize filter and give it a few updates to establish velocity,
        // then verify predictions drift in the expected direction.
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 0.1, 2.0], 0.25);
        t.predict();
        t.update(&[0.5, 0.5, 0.1, 2.0]);

        // Record position after first predict-update cycle
        let x_before: f32 = t.mean[0];
        let y_before: f32 = t.mean[1];

        // Run several predict-only cycles to let velocity dominate
        for _ in 0..5 {
            t.predict();
        }

        let x_after: f32 = t.mean[0];
        let y_after: f32 = t.mean[1];
        let h_after: f32 = t.mean[3];

        // The position should remain numerically reasonable (no NaN/Inf)
        assert!(x_after.is_finite(), "x should be finite after predictions");
        assert!(y_after.is_finite(), "y should be finite after predictions");
        assert!(
            h_after.is_finite(),
            "height should be finite after predictions"
        );

        // With near-zero velocity the predicted position should not explode
        assert!(
            (x_after - x_before).abs() < 5.0,
            "x drift should be bounded, got delta={}",
            (x_after - x_before).abs()
        );
        assert!(
            (y_after - y_before).abs() < 5.0,
            "y drift should be bounded, got delta={}",
            (y_after - y_before).abs()
        );
    }

    #[test]
    fn test_numerical_stability_1000_cycles() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);

        // Run 1000 predict-only cycles without any update
        for _ in 0..1000 {
            t.predict();
        }

        // Verify no NaN or Inf in the mean vector
        for i in 0..8 {
            let val: f32 = t.mean[i];
            assert!(
                val.is_finite(),
                "mean[{i}] should be finite after 1000 predictions, got {val}",
            );
        }

        // Verify no NaN or Inf in the covariance matrix
        for r in 0..8 {
            for c in 0..8 {
                let val: f32 = t.covariance[(r, c)];
                assert!(
                    val.is_finite(),
                    "covariance[({r},{c})] should be finite after 1000 predictions, got {val}",
                );
            }
        }
    }

    #[test]
    fn test_gating_distance_edge_cases() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        // Run a few predict-update cycles to stabilize
        for _ in 0..3 {
            t.predict();
            t.update(&[0.5, 0.5, 1.0, 0.5]);
        }
        t.predict();

        // Measurement exactly at the predicted state -- distance should be near 0
        let (projected_mean, _) = t.project();
        let mut meas_close = OMatrix::<f32, Dyn, U4>::from_element(1, 0.0);
        meas_close
            .row_mut(0)
            .copy_from_slice(projected_mean.as_slice());

        let dist_close = t.gating_distance(&meas_close, false, GatingDistanceMetric::Mahalanobis);
        assert!(
            dist_close[0].is_finite(),
            "Close-measurement distance should be finite"
        );
        assert!(
            dist_close[0] < 1.0,
            "Distance for exact-match measurement should be near 0, got {}",
            dist_close[0]
        );

        // Measurement far away -- distance should be large
        let mut meas_far = OMatrix::<f32, Dyn, U4>::from_element(1, 0.0);
        meas_far.copy_from_slice(&[10.0, 10.0, 5.0, 10.0]);

        let dist_far = t.gating_distance(&meas_far, false, GatingDistanceMetric::Mahalanobis);
        assert!(
            dist_far[0].is_finite(),
            "Far-measurement distance should be finite"
        );
        assert!(
            dist_far[0] > dist_close[0],
            "Far measurement should have larger distance than close one: {} vs {}",
            dist_far[0],
            dist_close[0]
        );
    }

    #[test]
    fn test_update_moves_mean_toward_measurement() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        t.predict();

        let x_before: f32 = t.mean[0];
        // Update with a measurement shifted to the right
        t.update(&[0.6, 0.5, 1.0, 0.5]);
        let x_after: f32 = t.mean[0];

        assert!(
            x_after > x_before,
            "Mean x should move toward the measurement (0.6), was {x_before}, now {x_after}"
        );
        assert!(
            x_after <= 0.6,
            "Mean x should not overshoot the measurement, got {x_after}"
        );
    }

    #[test]
    fn test_covariance_positive_diagonal() {
        let t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);

        // All diagonal elements of the covariance should be positive
        for i in 0..8 {
            let val: f32 = t.covariance[(i, i)];
            assert!(
                val > 0.0,
                "Covariance diagonal[{i}] should be positive, got {val}"
            );
        }
    }

    #[test]
    fn test_predict_increases_uncertainty() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);

        let cov_before: f32 = t.covariance[(0, 0)];
        t.predict();
        let cov_after: f32 = t.covariance[(0, 0)];

        assert!(
            cov_after > cov_before,
            "Predict should increase position uncertainty: {cov_before} -> {cov_after}"
        );
    }

    #[test]
    fn test_update_decreases_uncertainty() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        t.predict();

        let cov_before: f32 = t.covariance[(0, 0)];
        t.update(&[0.5, 0.5, 1.0, 0.5]);
        let cov_after: f32 = t.covariance[(0, 0)];

        assert!(
            cov_after < cov_before,
            "Update should decrease position uncertainty: {cov_before} -> {cov_after}"
        );
    }

    #[test]
    fn test_gating_distance_gaussian_vs_mahalanobis() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        for _ in 0..3 {
            t.predict();
            t.update(&[0.5, 0.5, 1.0, 0.5]);
        }
        t.predict();

        let mut measurements = OMatrix::<f32, Dyn, U4>::from_element(1, 0.0);
        measurements.copy_from_slice(&[0.6, 0.5, 1.0, 0.5]);

        let dist_gauss = t.gating_distance(&measurements, false, GatingDistanceMetric::Gaussian);
        let dist_maha = t.gating_distance(&measurements, false, GatingDistanceMetric::Mahalanobis);

        assert!(dist_gauss[0].is_finite());
        assert!(dist_maha[0].is_finite());

        // Both should be non-negative for a non-zero offset
        assert!(
            dist_gauss[0] > 0.0,
            "Gaussian distance should be > 0 for offset measurement"
        );
        assert!(
            dist_maha[0] > 0.0,
            "Mahalanobis distance should be > 0 for offset measurement"
        );
    }

    #[test]
    fn test_gating_distance_multiple_measurements() {
        let mut t = ConstantVelocityXYAHModel2::new(&[0.5, 0.5, 1.0, 0.5], 0.25);
        t.predict();
        t.update(&[0.5, 0.5, 1.0, 0.5]);
        t.predict();

        // Two measurements: one close, one far
        let mut measurements = OMatrix::<f32, Dyn, U4>::from_element(2, 0.0);
        measurements
            .row_mut(0)
            .copy_from_slice(&[0.5, 0.5, 1.0, 0.5]); // close
        measurements
            .row_mut(1)
            .copy_from_slice(&[5.0, 5.0, 1.0, 0.5]); // far

        let dists = t.gating_distance(&measurements, false, GatingDistanceMetric::Mahalanobis);
        assert_eq!(dists.len(), 2, "Should return one distance per measurement");
        assert!(dists[0].is_finite());
        assert!(dists[1].is_finite());
        assert!(
            dists[1] > dists[0],
            "Far measurement should have larger distance: {} vs {}",
            dists[1],
            dists[0]
        );
    }

    #[test]
    fn test_initiate_mean_matches_measurement() {
        let measurement = [0.3, 0.7, 1.5, 2.0];
        let t = ConstantVelocityXYAHModel2::new(&measurement, 0.25);

        // Position portion of mean should match the measurement exactly
        let x: f32 = t.mean[0];
        let y: f32 = t.mean[1];
        let a: f32 = t.mean[2];
        let h: f32 = t.mean[3];
        assert!((x - 0.3).abs() < 1e-6, "Mean x should be 0.3, got {x}");
        assert!((y - 0.7).abs() < 1e-6, "Mean y should be 0.7, got {y}");
        assert!((a - 1.5).abs() < 1e-6, "Mean a should be 1.5, got {a}");
        assert!((h - 2.0).abs() < 1e-6, "Mean h should be 2.0, got {h}");

        // Velocity portion should be zero
        for i in 4..8 {
            let v: f32 = t.mean[i];
            assert!((v).abs() < 1e-6, "Velocity mean[{i}] should be 0, got {v}");
        }
    }
}