src: implement simulation data sources

src/ego/alignment.rs

@@ -1,56 +1,95 @@
-use log::*;
-use nalgebra::{Matrix3, Rotation3, Unit, Vector3};
+
+use nalgebra::{Rotation3, Unit, Vector3, ComplexField};
+
+use crate::CircularBuffer;
 
 type Sample = Vector3<f32>; // [ax, ay, az] in m/s^2
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Default)]
 pub struct SensorAlignment {
     /// 3x3 transform matrix: raw -> body
     /// includes permutation, sign, and bias subtraction.
     pub transform: Rotation3<f32>,
-    /// bias vector in body frame
-    pub accel_bias: Vector3<f32>,
+    /// bias vector in sensor frame
+    pub bias: Vector3<f32>
+}
 
-    pub gyro_bias: Vector3<f32>,
+#[derive(Debug, Clone, Copy)]
+pub enum Direction {
+    UpZ,
+    DnZ,
+    UpX,
+    DnX,
+    UpY,
+    DnY
+}
+
+impl From<Direction> for Unit<Vector3<f32>> {
+    fn from(value: Direction) -> Self {
+        Unit::new_normalize(match value {
+            Direction::UpZ => Vector3::z(),
+            Direction::DnZ => -Vector3::z(),
+            Direction::UpX => Vector3::x(),
+            Direction::DnX => -Vector3::x(),
+            Direction::UpY => Vector3::y(),
+            Direction::DnY => -Vector3::y(),
+        })
+    }
 }
 
 impl SensorAlignment {
     pub fn apply(&self, raw: &Sample) -> Vector3<f32> {
-        self.transform * raw - self.accel_bias
+        self.transform * (raw - self.bias)
     }
 
-    pub fn apply_gyro(&self, raw: &Sample) -> Vector3<f32> {
-        self.transform * raw - self.gyro_bias
-    }
-
-    pub fn new(accel_samples: &[Sample], gyro_samples: &[Sample]) -> SensorAlignment {
-        let avg: Vector3<f32> = accel_samples.iter().sum::<Vector3<f32>>() / (accel_samples.len() as f32);
-        let down_body = Unit::new_normalize(avg);
-        let down_world = Unit::new_normalize(Vector3::new(0.0, 0.0, -1.0));
-        let rot_down = Rotation3::rotation_between(&down_world, &down_body).unwrap_or_else(Rotation3::identity);
-        let forward_world = Vector3::x();
-
-        let forward_projected = (forward_world - forward_world.dot(&down_world) * down_world.into_inner()).normalize();
-
-        // Transform forward into the body frame
-        let forward_body = rot_down * forward_projected;
-
-        // Construct a proper body frame basis
-        let right_body = down_body.cross(&forward_body).normalize();
-        let forward_body = right_body.cross(&down_body).normalize();
-
-        // Construct rotation matrix from basis vectors
-        let rot = Matrix3::from_columns(&[forward_body, right_body, down_body.into_inner()]);
-        let transform = Rotation3::from_matrix_unchecked(rot);
-
-        let gyro_avg: Vector3<f32> = gyro_samples.iter().sum::<Vector3<f32>>() / (gyro_samples.len() as f32);
-
-        info!("down_body={down_body:?} down_world={down_world:?}");
+    pub fn from_samples(accel_samples: &[Sample], down_direction: Direction) -> SensorAlignment {
+        let bias: Vector3<f32> = accel_samples.iter().sum::<Vector3<f32>>() / (accel_samples.len() as f32);
+        // Figure out the unit vector pointing 'down' at rest
+        let down_body = Unit::new_normalize(bias);
+        // Create our 'true' down vector
+        let down_world: Unit<Vector3<f32>> = down_direction.into();//Unit::new_normalize(Vector3::new(0.0, 0.0, -1.0));
+        // Determine how the axes are rotated between the two
+        let axis = down_body.cross(&down_world);
+        let axis_unit = Unit::new_normalize(axis);
+        let angle = down_body.dot(&down_world).clamp(-1.0, 1.0).acos();
+        // Create a rotation matrix from the angles
+        let transform = Rotation3::from_axis_angle(&axis_unit, angle);
 
         Self {
             transform,
-            accel_bias: transform * avg,
-            gyro_bias: transform * gyro_avg
+            bias,
         }
     }
+}
+
+
+#[derive(Debug)]
+pub struct DirectionEstimator {
+    buf: CircularBuffer<Vector3<f32>>,
+    direction: Direction,
+    pub alignment: SensorAlignment
+}
+
+impl DirectionEstimator {
+    pub fn new(direction: Direction) -> Self {
+        Self {
+            buf: Default::default(),
+            alignment: Default::default(),
+            direction,
+        }
+    }
+}
+
+impl DirectionEstimator {
+    pub fn update(&mut self, accel: Vector3<f32>) {
+        self.buf.insert(accel);
+
+        if self.buf.next_index == 0 {
+            self.alignment = SensorAlignment::from_samples(&self.buf.data, self.direction);
+        }
+    }
+
+    pub fn apply(&self, vec: &Vector3<f32>) -> Vector3<f32> {
+        self.alignment.apply(vec)
+    }
 }

src/ego/kalman.rs

@@ -0,0 +1,87 @@
+#![allow(non_snake_case)] // TODO, because most of this code came from chatgpt and I'm still not sure if it actually works or is completely wrong
+use nalgebra::{SMatrix, SVector};
+
+/// State: [px, py, vx, vy]
+pub type State = SVector<f32, 6>;
+pub type Control = SVector<f32, 3>; // dx, dx, yaw
+pub type GpsMeasurement = SVector<f32, 2>; // gps x, gps y
+
+pub struct Kalman {
+    x: State,                          // state estimate
+    P: SMatrix<f32, 6, 6>,             // covariance
+    F: SMatrix<f32, 6, 6>,             // state transition
+    B: SMatrix<f32, 6, 3>,             // control matrix
+    H: SMatrix<f32, 2, 6>,             // measurement matrix
+    Q: SMatrix<f32, 6, 6>,             // process noise
+    R: SMatrix<f32, 2, 2>,             // measurement noise
+}
+
+impl Kalman {
+    pub fn new(dt: f32) -> Self {
+        let mut kf = Kalman {
+            x: State::zeros(),
+            P: SMatrix::identity(),
+            F: SMatrix::identity(),
+            B: SMatrix::zeros(),
+            H: SMatrix::zeros(),
+            Q: SMatrix::identity() * 0.1,
+            R: SMatrix::identity() * 3.0, // GPS noise
+        };
+
+        // F (state transition)
+        kf.F[(0,2)] = dt;
+        kf.F[(1,3)] = dt;
+        kf.F[(4,5)] = dt;
+
+        // B (control matrix)
+        kf.B[(0,0)] = 0.5 * dt * dt;
+        kf.B[(1,1)] = 0.5 * dt * dt;
+        kf.B[(2,0)] = dt;
+        kf.B[(3,1)] = dt;
+        kf.B[(5,2)] = dt;
+
+        // H (GPS: positions only)
+        kf.H[(0,0)] = 1.0;
+        kf.H[(1,1)] = 1.0;
+
+        kf
+    }
+
+    pub fn predict(&mut self, u: &Control) {
+        self.x = self.F * self.x + self.B * u;
+        self.P = self.F * self.P * self.F.transpose() + self.Q;
+    }
+
+    pub fn update(&mut self, z: &GpsMeasurement) {
+        let y = z - self.H * self.x; // innovation
+        let S = self.H * self.P * self.H.transpose() + self.R;
+        if let Some(s_inv) = S.try_inverse() {
+            let K = self.P * self.H.transpose() * s_inv;
+            self.x += K * y;
+            self.P = (SMatrix::<f32,6,6>::identity() - K * self.H) * self.P;
+        }
+    }
+
+    // Used to indicate stationary behavior and therefore velocities should be zero
+    pub fn update_zupt(&mut self) {
+        let z = SVector::<f32, 2>::zeros(); // velocity = 0
+        let mut H = SMatrix::<f32, 2, 6>::zeros();
+        H[(0,2)] = 1.0; // vx
+        H[(1,3)] = 1.0; // vy
+
+        let R = SMatrix::<f32, 2, 2>::identity() * 0.01; // very confident
+
+        let y = z - H * self.x;
+        let S = H * self.P * H.transpose() + R;
+
+        if let Some(s_inv) = S.try_inverse() {
+            let k = self.P * H.transpose() * s_inv;
+            self.x += k * y;
+            self.P = (SMatrix::<f32,6,6>::identity() - k * H) * self.P;
+        }
+    }
+
+    pub fn state(&self) -> &State {
+        &self.x
+    }
+}

src/ego/mod.rs

@@ -1,3 +1,4 @@
 pub mod alignment;
 pub mod tilt;
-pub mod heading;
+pub mod heading;
+pub mod kalman;