feat(filter): implement kalman filter

src/config.rs

@@ -11,7 +11,7 @@ use embassy_rp::{
 use packed_struct::{derive::PackedStruct, PackedStruct};
 
 use crate::{
-    packed_float::{PackedFloat, ToPackedFloatArray},
+    helpers::{PackedFloat, ToPackedFloatArray},
     stick::{NotchStatus, NO_OF_CALIBRATION_POINTS, NO_OF_NOTCHES},
     ADDR_OFFSET, FLASH_SIZE,
 };

src/filter.rs

@@ -1,33 +1,68 @@
 use defmt::Format;
-use libm::{fminf, powf};
+use libm::{fmaxf, fminf, powf};
 
-use crate::config::ControllerConfig;
+use crate::{
+    config::{ControllerConfig, StickConfig},
+    helpers::XyValuePair,
+};
+
+macro_rules! run_kalman_on_axis {
+    ($self:ident, $axis:ident, $snapback:expr, $vel_weight1:ident, $vel_weight2:ident, $old_pos_filt:ident, $filter_gains:ident, $old_vel_filt:ident, $old_pos_diff:ident, $accel:ident, $vel_smooth:ident, $stick_distance6:ident) => {
+        if $snapback > 0 {
+            $self.vel_filt.$axis = $vel_weight1 * $self.vel.$axis
+                + (1. - $filter_gains.vel_decay.$axis) * $vel_weight2 * $old_vel_filt.$axis
+                + $filter_gains.vel_pos_factor.$axis * $old_pos_diff.$axis;
+
+            let pos_weight_vel_acc = 1.
+                - fminf(
+                    1.,
+                    $vel_smooth.$axis * $vel_smooth.$axis * $filter_gains.vel_thresh
+                        + $accel.$axis * $accel.$axis * $filter_gains.accel_thresh,
+                );
+            let pos_weight1 = fmaxf(pos_weight_vel_acc, $stick_distance6);
+            let pos_weight2 = 1. - pos_weight1;
+
+            $self.pos_filt.$axis = pos_weight1 * $self.pos.$axis
+                + pos_weight2
+                    * ($old_pos_filt.$axis
+                        + (1. - $filter_gains.vel_damp.$axis) * $self.vel_filt.$axis)
+        } else if $snapback < 0 {
+            let lpf = $old_pos_filt.$axis * $filter_gains.vel_damp.$axis
+                + $self.pos.$axis * (1. - $filter_gains.vel_damp.$axis);
+
+            let pos_weight_vel_acc = 1.
+                - fminf(
+                    1.,
+                    $vel_smooth.$axis * $vel_smooth.$axis * $filter_gains.vel_thresh
+                        + $accel.$axis * $accel.$axis * $filter_gains.accel_thresh,
+                );
+            let pos_weight1 = fmaxf(pos_weight_vel_acc, $stick_distance6);
+            let pos_weight2 = 1. - pos_weight1;
+
+            $self.pos_filt.$axis = pos_weight1 * $self.pos.$axis + pos_weight2 * lpf;
+        } else {
+            $self.pos_filt.$axis = $self.pos.$axis;
+        }
+    };
+}
 
 /// Filter gains for 800Hz, the ones for 1000Hz are provided by `get_norm_gains`
 pub const FILTER_GAINS: FilterGains = FilterGains {
     max_stick: 100.,
-    x_vel_decay: 0.1,
+    vel_decay: XyValuePair { x: 0.1, y: 0.1 },
-    y_vel_decay: 0.1,
+    vel_pos_factor: XyValuePair { x: 0.01, y: 0.01 },
-    x_vel_pos_factor: 0.01,
+    vel_damp: XyValuePair { x: 0.125, y: 0.125 },
-    y_vel_pos_factor: 0.01,
-    x_vel_damp: 0.125,
-    y_vel_damp: 0.125,
     vel_thresh: 1.,
     accel_thresh: 3.,
-    x_smoothing: 0.0,
+    smoothing: XyValuePair { x: 0.0, y: 0.0 },
-    y_smoothing: 0.0,
+    c_smoothing: XyValuePair { x: 0.0, y: 0.0 },
-    c_xsmoothing: 0.0,
-    c_ysmoothing: 0.0,
 };
 
 #[derive(Debug, Clone, Default)]
 pub struct WaveshapingValues {
-    pub old_x_pos: f32,
+    pub old_pos: XyValuePair<f32>,
-    pub old_y_pos: f32,
+    pub old_vel: XyValuePair<f32>,
-    pub old_x_vel: f32,
+    pub old_out: XyValuePair<f32>,
-    pub old_y_vel: f32,
-    pub old_x_out: f32,
-    pub old_y_out: f32,
 }
 
 fn calc_waveshaping_mult(setting: u8) -> f32 {
@@ -54,17 +89,14 @@ pub struct FilterGains {
     /// filtered velocity terms
     /// how fast the filtered velocity falls off in the absence of stick movement.
     /// Probably don't touch this.
-    pub x_vel_decay: f32, //0.1 default for 1.2ms timesteps, larger for bigger timesteps
+    pub vel_decay: XyValuePair<f32>, //0.1 default for 1.2ms timesteps, larger for bigger timesteps
-    pub y_vel_decay: f32,
     /// how much the current position disagreement impacts the filtered velocity.
     /// Probably don't touch this.
-    pub x_vel_pos_factor: f32, //0.01 default for 1.2ms timesteps, larger for bigger timesteps
+    pub vel_pos_factor: XyValuePair<f32>, //0.01 default for 1.2ms timesteps, larger for bigger timesteps
-    pub y_vel_pos_factor: f32,
     /// how much to ignore filtered velocity when computing the new stick position.
     /// DO CHANGE THIS
     /// Higher gives shorter rise times and slower fall times (more pode, less snapback)
-    pub x_vel_damp: f32, //0.125 default for 1.2ms timesteps, smaller for bigger timesteps
+    pub vel_damp: XyValuePair<f32>, //0.125 default for 1.2ms timesteps, smaller for bigger timesteps
-    pub y_vel_damp: f32,
     /// speed and accel thresholds below which we try to follow the stick better
     /// These may need tweaking according to how noisy the signal is
     /// If it's noisier, we may need to add additional filtering
@@ -76,11 +108,9 @@ pub struct FilterGains {
     /// This just applies a low-pass filter.
     /// The purpose is to provide delay for single-axis ledgedashes.
     /// Must be between 0 and 1. Larger = more smoothing and delay.
-    pub x_smoothing: f32,
+    pub smoothing: XyValuePair<f32>,
-    pub y_smoothing: f32,
     /// Same thing but for C-stick
-    pub c_xsmoothing: f32,
+    pub c_smoothing: XyValuePair<f32>,
-    pub c_ysmoothing: f32,
 }
 
 impl FilterGains {
@@ -88,14 +118,14 @@ impl FilterGains {
     pub fn get_normalized_gains(&self, controller_config: &ControllerConfig) -> Self {
         let mut gains = self.clone();
 
-        gains.x_vel_damp = vel_damp_from_snapback(controller_config.astick_config.x_snapback);
+        gains.vel_damp.x = vel_damp_from_snapback(controller_config.astick_config.x_snapback);
-        gains.y_vel_damp = vel_damp_from_snapback(controller_config.astick_config.y_snapback);
+        gains.vel_damp.y = vel_damp_from_snapback(controller_config.astick_config.y_snapback);
 
-        gains.x_smoothing = controller_config.astick_config.x_smoothing as f32 / 10.;
+        gains.smoothing.x = controller_config.astick_config.x_smoothing as f32 / 10.;
-        gains.y_smoothing = controller_config.astick_config.y_smoothing as f32 / 10.;
+        gains.smoothing.y = controller_config.astick_config.y_smoothing as f32 / 10.;
 
-        gains.c_xsmoothing = controller_config.cstick_config.x_smoothing as f32 / 10.;
+        gains.c_smoothing.x = controller_config.cstick_config.x_smoothing as f32 / 10.;
-        gains.c_ysmoothing = controller_config.cstick_config.y_smoothing as f32 / 10.;
+        gains.c_smoothing.y = controller_config.cstick_config.y_smoothing as f32 / 10.;
 
         // The below is assuming the sticks to be polled at 1000Hz
         let time_factor = 1.0 / 1.2;
@@ -106,37 +136,149 @@ impl FilterGains {
 
         FilterGains {
             max_stick: gains.max_stick * gains.max_stick,
-            x_vel_decay: gains.x_vel_decay * time_factor,
+            vel_decay: XyValuePair {
-            y_vel_decay: gains.y_vel_decay * time_factor,
+                x: gains.vel_decay.x * time_factor,
-            x_vel_pos_factor: gains.x_vel_pos_factor * time_factor,
+                y: gains.vel_decay.y * time_factor,
-            y_vel_pos_factor: gains.y_vel_pos_factor * time_factor,
+            },
-            x_vel_damp: gains.x_vel_damp
+            vel_pos_factor: XyValuePair {
-                * match controller_config.astick_config.x_snapback {
+                x: gains.vel_pos_factor.x * time_factor,
-                    a if a >= 0 => time_factor,
+                y: gains.vel_pos_factor.y * time_factor,
-                    _ => 1.0,
+            },
-                },
+            vel_damp: XyValuePair {
-            y_vel_damp: gains.y_vel_damp
+                x: gains.vel_damp.x
-                * match controller_config.astick_config.y_snapback {
+                    * match controller_config.astick_config.x_snapback {
-                    a if a >= 0 => time_factor,
+                        a if a >= 0 => time_factor,
-                    _ => 1.0,
+                        _ => 1.0,
-                },
+                    },
+                y: gains.vel_damp.y
+                    * match controller_config.astick_config.y_snapback {
+                        a if a >= 0 => time_factor,
+                        _ => 1.0,
+                    },
+            },
             vel_thresh,
             accel_thresh,
-            x_smoothing: powf(1.0 - gains.x_smoothing, time_divisor),
+            smoothing: XyValuePair {
-            y_smoothing: powf(1.0 - gains.y_smoothing, time_divisor),
+                x: powf(1.0 - gains.smoothing.x, time_divisor),
-            c_xsmoothing: powf(1.0 - gains.c_xsmoothing, time_divisor),
+                y: powf(1.0 - gains.smoothing.y, time_divisor),
-            c_ysmoothing: powf(1.0 - gains.c_ysmoothing, time_divisor),
+            },
+            c_smoothing: XyValuePair {
+                x: powf(1.0 - gains.c_smoothing.x, time_divisor),
+                y: powf(1.0 - gains.c_smoothing.y, time_divisor),
+            },
         }
     }
 }
 
-pub fn run_kalman(
+#[derive(Clone, Debug, Format, Default)]
-    x_z: f32,
+pub struct KalmanState {
-    y_z: f32,
+    pos: XyValuePair<f32>,
-    controller_config: &ControllerConfig,
+    vel: XyValuePair<f32>,
-    filter_gains: &FilterGains,
+    vel_filt: XyValuePair<f32>,
-) -> (f32, f32) {
+    pos_filt: XyValuePair<f32>,
-    todo!()
+}
+
+impl KalmanState {
+    // runs kalman filter
+    pub fn run_kalman(
+        &mut self,
+        x_z: f32,
+        y_z: f32,
+        stick_config: &StickConfig,
+        filter_gains: &FilterGains,
+    ) -> (f32, f32) {
+        let old_pos = self.pos;
+        let old_vel = self.vel;
+        let old_vel_filt = self.vel_filt;
+        let old_pos_filt = self.pos_filt;
+
+        self.pos.x = x_z;
+        self.pos.y = y_z;
+        self.vel.x = x_z - old_pos.x;
+        self.vel.y = y_z - old_pos.y;
+
+        let vel_smooth = XyValuePair {
+            x: 0.5 * (self.vel.x + old_vel.x),
+            y: 0.5 * (self.vel.y + old_vel.y),
+        };
+        let accel = XyValuePair {
+            x: self.vel.x - old_vel.x,
+            y: self.vel.y - old_vel.y,
+        };
+        let old_pos_diff = XyValuePair {
+            x: old_pos.x - old_pos_filt.x,
+            y: old_pos.y - old_pos_filt.y,
+        };
+
+        let stick_distance2 = fminf(
+            filter_gains.max_stick,
+            self.pos.x * self.pos.x + self.pos.y * self.pos.y,
+        ) / filter_gains.max_stick;
+        let stick_distance6 = stick_distance2 * stick_distance2 * stick_distance2;
+
+        let vel_weight1 = stick_distance2;
+        let vel_weight2 = 1. - vel_weight1;
+
+        //modified velocity to feed into our kalman filter.
+        //We don't actually want an accurate model of the velocity, we want to suppress snapback without adding delay
+        //term 1: weight current velocity according to r^2
+        //term 2: the previous filtered velocity, weighted the opposite and also set to decay
+        //term 3: a corrective factor based on the disagreement between real and filtered position
+
+        //the current position weight used for the filtered position is whatever is larger of
+        //  a) 1 minus the sum of the squares of
+        //    1) the smoothed velocity divided by the velocity threshold
+        //    2) the acceleration divided by the accel threshold
+        //  b) stick r^6
+        //When the stick is moving slowly, we want to weight it highly, in order to achieve
+        //  quick control for inputs such as tilts. We lock out using both velocity and
+        //  acceleration in order to rule out snapback.
+        //When the stick is near the rim, we also want instant response, and we know snapback
+        //  doesn't reach the rim.
+
+        //In calculating the filtered stick position, we have the following components
+        //term 1: current position, weighted according to the above weight
+        //term 2: a predicted position based on the filtered velocity and previous filtered position,
+        //  with the filtered velocity damped, and the overall term weighted inverse of the previous term
+        //term 3: the integral error correction term
+
+        //But if we xSnapback or ySnapback is zero, we skip the calculation
+        run_kalman_on_axis!(
+            self,
+            x,
+            stick_config.x_snapback,
+            vel_weight1,
+            vel_weight2,
+            old_pos_filt,
+            filter_gains,
+            old_vel_filt,
+            old_pos_diff,
+            accel,
+            vel_smooth,
+            stick_distance6
+        );
+
+        run_kalman_on_axis!(
+            self,
+            y,
+            stick_config.y_snapback,
+            vel_weight1,
+            vel_weight2,
+            old_pos_filt,
+            filter_gains,
+            old_vel_filt,
+            old_pos_diff,
+            accel,
+            vel_smooth,
+            stick_distance6
+        );
+
+        self.get_xy()
+    }
+
+    pub fn get_xy(&self) -> (f32, f32) {
+        (self.pos_filt.x, self.pos_filt.y)
+    }
 }
 
 /// This simulates an idealized sort of pode:
@@ -160,11 +302,11 @@ pub fn run_waveshaping(
     let x_factor = calc_waveshaping_mult(x_waveshaping);
     let y_factor = calc_waveshaping_mult(y_waveshaping);
 
-    let x_vel = x_pos - waveshaping_values.old_x_pos;
+    let x_vel = x_pos - waveshaping_values.old_pos.x;
-    let y_vel = y_pos - waveshaping_values.old_y_pos;
+    let y_vel = y_pos - waveshaping_values.old_pos.y;
 
-    let x_vel_smooth = 0.5 * (x_vel + waveshaping_values.old_x_vel);
+    let x_vel_smooth = 0.5 * (x_vel + waveshaping_values.old_vel.x);
-    let y_vel_smooth = 0.5 * (y_vel + waveshaping_values.old_y_vel);
+    let y_vel_smooth = 0.5 * (y_vel + waveshaping_values.old_vel.y);
 
     let old_x_pos_weight = fminf(
         1.,
@@ -177,15 +319,15 @@ pub fn run_waveshaping(
     );
     let new_y_pos_weight = 1. - old_y_pos_weight;
 
-    let x_out = x_pos * new_x_pos_weight + waveshaping_values.old_x_out * old_x_pos_weight;
+    let x_out = x_pos * new_x_pos_weight + waveshaping_values.old_out.x * old_x_pos_weight;
-    let y_out = y_pos * new_y_pos_weight + waveshaping_values.old_y_out * old_y_pos_weight;
+    let y_out = y_pos * new_y_pos_weight + waveshaping_values.old_out.y * old_y_pos_weight;
 
-    waveshaping_values.old_x_pos = x_pos;
+    waveshaping_values.old_pos.x = x_pos;
-    waveshaping_values.old_y_pos = y_pos;
+    waveshaping_values.old_pos.y = y_pos;
-    waveshaping_values.old_x_vel = x_vel_smooth;
+    waveshaping_values.old_vel.x = x_vel_smooth;
-    waveshaping_values.old_y_vel = y_vel_smooth;
+    waveshaping_values.old_vel.y = y_vel_smooth;
-    waveshaping_values.old_x_out = x_out;
+    waveshaping_values.old_out.x = x_out;
-    waveshaping_values.old_y_out = y_out;
+    waveshaping_values.old_out.y = y_out;
 
     (x_out, y_out)
 }

src/helpers.rs

@@ -46,3 +46,9 @@ impl Deref for PackedFloat {
         &self.0
     }
 }
+
+#[derive(Debug, Clone, Format, Default, Copy)]
+pub struct XyValuePair<T> {
+    pub x: T,
+    pub y: T,
+}

src/input.rs

@@ -16,7 +16,7 @@ use libm::{fmaxf, fminf};
 
 use crate::{
     config::ControllerConfig,
-    filter::{run_kalman, run_waveshaping, FilterGains, WaveshapingValues, FILTER_GAINS},
+    filter::{run_waveshaping, FilterGains, KalmanState, WaveshapingValues, FILTER_GAINS},
     gcc_hid::GcReport,
     stick::{linearize, notch_remap, StickParams},
     FLASH_SIZE,
@@ -125,6 +125,7 @@ async fn update_stick_states<
     cstick_waveshaping_values: &mut WaveshapingValues,
     old_stick_pos: &mut StickPositions,
     raw_stick_values: &mut RawStickValues,
+    kalman_state: &mut KalmanState,
 ) -> StickState {
     let mut adc_count = 0u32;
     let mut ax_sum = 0u32;
@@ -203,7 +204,8 @@ async fn update_stick_states<
     raw_stick_values.cx_linearized = pos_cx;
     raw_stick_values.cy_linearized = pos_cy;
 
-    let (x_pos_filt, y_pos_filt) = run_kalman(x_z, y_z, controller_config, &filter_gains);
+    let (x_pos_filt, y_pos_filt) =
+        kalman_state.run_kalman(x_z, y_z, &controller_config.astick_config, &filter_gains);
 
     let (shaped_x, shaped_y) = run_waveshaping(
         x_pos_filt,
@@ -215,9 +217,9 @@ async fn update_stick_states<
     );
 
     let pos_x: f32 =
-        filter_gains.x_smoothing * shaped_x + (1.0 - filter_gains.x_smoothing) * old_stick_pos.x;
+        filter_gains.smoothing.x * shaped_x + (1.0 - filter_gains.smoothing.x) * old_stick_pos.x;
     let pos_y =
-        filter_gains.y_smoothing * shaped_y + (1.0 - filter_gains.y_smoothing) * old_stick_pos.y;
+        filter_gains.smoothing.y * shaped_y + (1.0 - filter_gains.smoothing.y) * old_stick_pos.y;
     old_stick_pos.x = pos_x;
     old_stick_pos.y = pos_y;
 
@@ -235,9 +237,9 @@ async fn update_stick_states<
     old_stick_pos.cx = shaped_cx;
     old_stick_pos.cy = shaped_cy;
 
-    let x_weight_1 = filter_gains.c_xsmoothing;
+    let x_weight_1 = filter_gains.c_smoothing.x;
     let x_weight_2 = 1.0 - x_weight_1;
-    let y_weight_1 = filter_gains.c_ysmoothing;
+    let y_weight_1 = filter_gains.c_smoothing.y;
     let y_weight_2 = 1.0 - y_weight_1;
 
     let pos_cx_filt = x_weight_1 * shaped_cx + x_weight_2 * old_cx_pos;
@@ -401,6 +403,7 @@ pub async fn input_loop(
         let mut old_stick_pos = StickPositions::default();
         let mut cstick_waveshaping_values = WaveshapingValues::default();
         let mut controlstick_waveshaping_values = WaveshapingValues::default();
+        let mut kalman_state = KalmanState::default();
 
         loop {
             current_stick_state = update_stick_states(
@@ -416,6 +419,7 @@ pub async fn input_loop(
                 &mut cstick_waveshaping_values,
                 &mut old_stick_pos,
                 &mut raw_stick_values,
+                &mut kalman_state,
             )
             .await;
 

src/main.rs

@@ -7,8 +7,8 @@
 mod config;
 mod filter;
 mod gcc_hid;
+mod helpers;
 mod input;
-mod packed_float;
 mod stick;
 
 use defmt::{debug, info};

src/stick.rs

@@ -7,8 +7,8 @@ use libm::{atan2f, cosf, fabs, roundf, sinf, sqrtf};
 
 use crate::{
     config::{ControllerConfig, StickConfig, DEFAULT_NOTCH_STATUS},
+    helpers::ToRegularArray,
     input::Stick,
-    packed_float::ToRegularArray,
 };
 
 /// fit order for the linearization