src: implement ego tracking models, and port shaders from renderbug into here

src/tasks/mpu.rs

@@ -0,0 +1,180 @@
+use core::cell::RefCell;
+
+use embassy_embedded_hal::shared_bus::asynch::i2c::I2cDevice;
+use embassy_sync::{blocking_mutex::raw::{CriticalSectionRawMutex, NoopRawMutex}, channel::{DynamicSender, Sender}};
+use embassy_time::{Delay, Timer};
+use esp_hal::{i2c::master::I2c, Async};
+use log::{info, warn};
+use mpu6050_dmp::{address::Address, calibration::CalibrationParameters, error_async::Error, sensor_async::Mpu6050};
+use nalgebra::Vector3;
+use core::f32::consts::PI;
+use nalgebra::ComplexField;
+
+use crate::{backoff::Backoff, ego::{alignment::SensorAlignment, tilt::TiltEstimator}, events::{Measurement, Notification}};
+
+const G: f32 = 9.80665;
+const SENS_2G: f32 = 16384.0;
+const DEG2RAD: f32 = PI / 180.0;
+
+/// Convert raw accel (i16) to m/s^2
+fn accel_raw_to_ms2(raw: i16) -> f32 {
+    (raw as f32 / SENS_2G) * G
+}
+
+/// Convert raw gyro (i16) to rad/s. For MPU6050 typical gyro scale at ±2000 deg/s -> 16.4 LSB/deg/s
+/// Adjust if you configured different full-scale range.
+fn gyro_raw_to_rads(raw: i16) -> f32 {
+    // lsb_per_dps e.g. 16.4 for ±2000 dps
+    (raw as f32 / 16.4) * (DEG2RAD)
+}
+
+/// Rotate body accel into ENU and remove gravity; returns (east, north, up) linear accel in m/s^2
+pub fn accel_body_to_enu_lin(body_acc: (f32,f32,f32), roll: f32, pitch: f32, yaw: f32) -> (f32,f32,f32) {
+    let (ax, ay, az) = body_acc;
+
+    let (sr, cr) = roll.sin_cos();
+    let (sp, cp) = pitch.sin_cos();
+    let (sy, cy) = yaw.sin_cos();
+
+    // Build R = Rz(yaw) * Ry(pitch) * Rx(roll)
+    // multiply R * a_body
+    let r11 =  cy*cp;
+    let r12 =  cy*sp*sr - sy*cr;
+    let r13 =  cy*sp*cr + sy*sr;
+
+    let r21 =  sy*cp;
+    let r22 =  sy*sp*sr + cy*cr;
+    let r23 =  sy*sp*cr - cy*sr;
+
+    let r31 = -sp;
+    let r32 =  cp*sr;
+    let r33 =  cp*cr;
+
+    // a_nav = R * a_body
+    let ax_nav = r11*ax + r12*ay + r13*az; // east
+    let ay_nav = r21*ax + r22*ay + r23*az; // north
+    let az_nav = r31*ax + r32*ay + r33*az; // up
+
+    // remove gravity (nav frame gravity points down in NED; in our ENU 'up' axis positive up)
+    // If accel measured gravity as +9.8 on sensor when static and z down convention used above,
+    // then az_nav will include +g in 'up' direction; subtract G.
+    let ax_lin = ax_nav;
+    let ay_lin = ay_nav;
+    let az_lin = az_nav - G; // linear acceleration in up axis
+
+    (ax_lin, ay_lin, az_lin)
+}
+
+#[embassy_executor::task]
+pub async fn mpu_task(events: Sender<'static, NoopRawMutex, Measurement, 4>, sink: DynamicSender<'static, Notification>, bus: I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>) {
+    let backoff = Backoff::from_millis(5);
+    let busref = RefCell::new(Some(bus));
+
+    backoff.forever().attempt::<_, (), ()>(async || {
+        let mut sensor = backoff.forever().attempt(async || {
+            info!("Initializing MPU");
+            match Mpu6050::new(busref.replace(None).unwrap(), Address::default()).await.map_err(|e| { e.i2c }) {
+                Err(i2c) => {
+                    busref.replace(Some(i2c));
+                    Err(())
+                },
+                Ok(mut sensor) => {
+                    match backoff.attempt(async || { mpu_init(&mut sensor).await }).await {
+                        Err(_) => {
+                            busref.replace(Some(sensor.release()));
+                            Err(())
+                        },
+                        Ok(_) => Ok(sensor)
+                    }
+                }
+            }
+        }).await?;
+
+        let sensor_ref = &mut sensor;
+
+        let alignment = {
+            info!("Locating earth...");
+            let mut accel_samples = [Vector3::default(); 100];
+            let mut gyro_samples = [Vector3::default(); 100];
+            for (idx, (accel_sample, gyro_sample)) in accel_samples.iter_mut().zip(gyro_samples.iter_mut()).enumerate() {
+                let (accel, gyro) = backoff.attempt(async || { sensor_ref.motion6().await }).await.unwrap();
+                *accel_sample = Vector3::new(accel_raw_to_ms2(accel.x()), accel_raw_to_ms2(accel.y()), accel_raw_to_ms2(accel.z()));
+                *gyro_sample = Vector3::new(gyro_raw_to_rads(gyro.x()), gyro_raw_to_rads(gyro.y()), gyro_raw_to_rads(gyro.z()));
+                info!("{idx}/200 {accel_sample:?} {gyro_sample:?}");
+                Timer::after_millis(5).await;
+            }
+
+            SensorAlignment::new(accel_samples.as_slice(), gyro_samples.as_slice())
+        };
+
+        info!("MPU is ready! earth={alignment:?}");
+        sink.send(Notification::CalibrationComplete).await;
+        let mut tilt = TiltEstimator::new(0.3);
+        //TODO: Need to read in a constant buffer of accelerometer measurements, which we can then use to determine where "forward" points in the body frame when converting from the sensor frame.
+        // From there, we can rotate the body frame into the world frame using gps headings to generate a compass north
+        fn lowpass(prev: f32, current: f32, alpha: f32) -> f32 {
+            prev + alpha * (current - prev) // alpha in (0,1), small alpha = heavy smoothing
+        }
+        let mut prev_accel = Vector3::default();
+        loop {
+            match backoff.attempt(async || { sensor_ref.motion6().await }).await {
+                Ok((accel_data, gyro_data)) => {
+
+                    // Apply the initial alignment correction to put it into the body frame where X is forward
+                    let adjusted = alignment.apply(&Vector3::new(
+                        accel_raw_to_ms2(accel_data.x()),
+                        accel_raw_to_ms2(accel_data.y()),
+                        accel_raw_to_ms2(accel_data.z())
+                    ));
+                    let adjusted_gyro = alignment.apply_gyro(&Vector3::new(
+                        gyro_raw_to_rads(gyro_data.x()),
+                        gyro_raw_to_rads(gyro_data.y()),
+                        gyro_raw_to_rads(gyro_data.z())
+                    ));
+
+                    let filtered = Vector3::new(
+                        lowpass(prev_accel.x, adjusted.x, 0.3),
+                        lowpass(prev_accel.y, adjusted.y, 0.3),
+                        lowpass(prev_accel.z, adjusted.z, 0.3),
+                    );
+                    prev_accel = adjusted;
+
+                    // Apply current rotation and accel values to rotate our values back into the ENU frame
+                    tilt.update(0.01, (filtered.x, filtered.y, filtered.z), (adjusted_gyro.x, adjusted_gyro.y, adjusted_gyro.z));
+                    let (roll, pitch, yaw) = tilt.angles();
+                    let (ax_e, ay_n, _az_up) = accel_body_to_enu_lin((filtered.x,filtered.y,filtered.z), roll, pitch, yaw);
+
+                    let measured = Vector3::new(ax_e, ay_n, adjusted_gyro.z);
+                    events.send(
+                        Measurement::IMU(
+                            measured
+                        )
+                    ).await;
+                    Timer::after_millis(10).await;
+                },
+                Err(e) => {
+                    warn!("Failed to read MPU motion data! {e:?}");
+                    busref.replace(Some(sensor.release()));
+                    return Err(());
+                }
+            };
+        }
+    }).await.ok();
+}
+
+async fn mpu_init(sensor: &mut Mpu6050<I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>>) -> Result<(), Error<I2cDevice<'static, CriticalSectionRawMutex, I2c<'static, Async>>>> {
+    let mut delay = Delay;
+    let backoff = Backoff::from_millis(10);
+    info!("Initializing MPU");
+    backoff.attempt(async || { sensor.initialize_dmp(&mut delay).await }).await?;
+    backoff.attempt(async || { sensor.set_digital_lowpass_filter(mpu6050_dmp::config::DigitalLowPassFilter::Filter3).await }).await?;
+    info!("Calibrating MPU");
+    backoff.attempt(async || {
+        // TODO: Store calibration data with set_accel_calibration, set_gyro_calibration
+        sensor.calibrate(&mut delay, &CalibrationParameters::new(
+            mpu6050_dmp::accel::AccelFullScale::G2,
+            mpu6050_dmp::gyro::GyroFullScale::Deg2000,
+            mpu6050_dmp::calibration::ReferenceGravity::Zero
+        )).await.map(|_| { Ok(()) })
+    }).await?
+}