Add down estimator plot

Change-Id: Id3fe6fef762cde1e77d09f779c451b3c2e3a471e
diff --git a/frc971/analysis/BUILD b/frc971/analysis/BUILD
index b72634a..2223955 100644
--- a/frc971/analysis/BUILD
+++ b/frc971/analysis/BUILD
@@ -79,6 +79,7 @@
         "//aos:configuration_ts_fbs",
         "//aos/network/www:demo_plot",
         "//aos/network/www:proxy",
+        "//frc971/control_loops/drivetrain:down_estimator_plotter",
         "//frc971/control_loops/drivetrain:drivetrain_plotter",
         "//frc971/control_loops/drivetrain:robot_state_plotter",
         "//frc971/wpilib:imu_plotter",
diff --git a/frc971/analysis/plot_index.ts b/frc971/analysis/plot_index.ts
index bfc8d53..1eba066 100644
--- a/frc971/analysis/plot_index.ts
+++ b/frc971/analysis/plot_index.ts
@@ -24,6 +24,7 @@
 import * as proxy from 'org_frc971/aos/network/www/proxy';
 import {plotImu} from 'org_frc971/frc971/wpilib/imu_plotter';
 import {plotDrivetrain} from 'org_frc971/frc971/control_loops/drivetrain/drivetrain_plotter';
+import {plotDownEstimator} from 'org_frc971/frc971/control_loops/drivetrain/down_estimator_plotter';
 import {plotRobotState} from
     'org_frc971/frc971/control_loops/drivetrain/robot_state_plotter'
 import {plotDemo} from 'org_frc971/aos/network/www/demo_plot';
@@ -83,6 +84,7 @@
   ['Demo', new PlotState(plotDiv, plotDemo)],
   ['IMU', new PlotState(plotDiv, plotImu)],
   ['Drivetrain', new PlotState(plotDiv, plotDrivetrain)],
+  ['Down Estimator', new PlotState(plotDiv, plotDownEstimator)],
   ['Robot State', new PlotState(plotDiv, plotRobotState)],
   ['C++ Plotter', new PlotState(plotDiv, plotData)],
 ]);
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index c420333..1da9db3 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -726,6 +726,18 @@
 )
 
 ts_library(
+    name = "down_estimator_plotter",
+    srcs = ["down_estimator_plotter.ts"],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        "//aos/network/www:aos_plotter",
+        "//aos/network/www:colors",
+        "//aos/network/www:proxy",
+        "//frc971/wpilib:imu_plot_utils",
+    ],
+)
+
+ts_library(
     name = "drivetrain_plotter",
     srcs = ["drivetrain_plotter.ts"],
     target_compatible_with = ["@platforms//os:linux"],
diff --git a/frc971/control_loops/drivetrain/down_estimator_plotter.ts b/frc971/control_loops/drivetrain/down_estimator_plotter.ts
new file mode 100644
index 0000000..c6e414c
--- /dev/null
+++ b/frc971/control_loops/drivetrain/down_estimator_plotter.ts
@@ -0,0 +1,160 @@
+// Provides a basic plot for debugging IMU-related issues on a robot.
+import {AosPlotter} from 'org_frc971/aos/network/www/aos_plotter';
+import {ImuMessageHandler} from 'org_frc971/frc971/wpilib/imu_plot_utils';
+import * as proxy from 'org_frc971/aos/network/www/proxy';
+import {BLUE, BROWN, CYAN, GREEN, PINK, RED, WHITE} from 'org_frc971/aos/network/www/colors';
+
+import Connection = proxy.Connection;
+
+export function plotDownEstimator(conn: Connection, element: Element): void {
+  const width = 900;
+  const height = 400;
+  const aosPlotter = new AosPlotter(conn);
+
+  const status = aosPlotter.addMessageSource(
+      '/drivetrain', 'frc971.control_loops.drivetrain.Status');
+
+  const imu = aosPlotter.addRawMessageSource(
+      '/drivetrain', 'frc971.IMUValuesBatch',
+      new ImuMessageHandler(conn.getSchema('frc971.IMUValuesBatch')));
+
+  const accelPlot = aosPlotter.addPlot(element, [0, 0], [width, height]);
+  accelPlot.plot.getAxisLabels().setTitle(
+      'Estimated Accelerations (x = forward, y = lateral, z = vertical)');
+  accelPlot.plot.getAxisLabels().setYLabel('Acceleration (m/s/s)');
+  accelPlot.plot.getAxisLabels().setXLabel('Monotonic Reading Time (sec)');
+
+  const accelX = accelPlot.addMessageLine(status, ['down_estimator', 'accel_x']);
+  accelX.setColor(RED);
+  const accelY = accelPlot.addMessageLine(status, ['down_estimator', 'accel_y']);
+  accelY.setColor(GREEN);
+  const accelZ = accelPlot.addMessageLine(status, ['down_estimator', 'accel_z']);
+  accelZ.setColor(BLUE);
+
+  const velPlot = aosPlotter.addPlot(element, [0, height], [width, height]);
+  velPlot.plot.getAxisLabels().setTitle('Raw IMU Integrated Velocity');
+  velPlot.plot.getAxisLabels().setYLabel('Velocity (m/s)');
+  velPlot.plot.getAxisLabels().setXLabel('Monotonic Reading Time (sec)');
+
+  const velX = velPlot.addMessageLine(status, ['down_estimator', 'velocity_x']);
+  velX.setColor(RED);
+  const velY = velPlot.addMessageLine(status, ['down_estimator', 'velocity_y']);
+  velY.setColor(GREEN);
+  const velZ = velPlot.addMessageLine(status, ['down_estimator', 'velocity_z']);
+  velZ.setColor(BLUE);
+
+  const gravityPlot = aosPlotter.addPlot(element, [0, height * 2], [width, height]);
+  gravityPlot.plot.getAxisLabels().setTitle('Accelerometer Magnitudes');
+  gravityPlot.plot.getAxisLabels().setXLabel('Monotonic Sent Time (sec)');
+  gravityPlot.plot.setDefaultYRange([0.95, 1.05]);
+  const gravityLine =
+      gravityPlot.addMessageLine(status, ['down_estimator', 'gravity_magnitude']);
+  gravityLine.setColor(RED);
+  gravityLine.setDrawLine(false);
+  const accelMagnitudeLine =
+      gravityPlot.addMessageLine(imu, ['acceleration_magnitude_filtered']);
+  accelMagnitudeLine.setColor(BLUE);
+  accelMagnitudeLine.setLabel('Filtered Accel Magnitude (0.1 sec)');
+  accelMagnitudeLine.setDrawLine(false);
+
+  const orientationPlot =
+      aosPlotter.addPlot(element, [0, height * 3], [width, height]);
+  orientationPlot.plot.getAxisLabels().setTitle('Orientation');
+  orientationPlot.plot.getAxisLabels().setYLabel('Angle (rad)');
+
+  const roll = orientationPlot.addMessageLine(
+      status, ['down_estimator', 'lateral_pitch']);
+  roll.setColor(RED);
+  roll.setLabel('roll');
+  const pitch = orientationPlot.addMessageLine(
+      status, ['down_estimator', 'longitudinal_pitch']);
+  pitch.setColor(GREEN);
+  pitch.setLabel('pitch');
+  const yaw = orientationPlot.addMessageLine(
+      status, ['down_estimator', 'yaw']);
+  yaw.setColor(BLUE);
+  yaw.setLabel('yaw');
+
+  const imuAccelPlot = aosPlotter.addPlot(element, [0, height * 4], [width, height]);
+  imuAccelPlot.plot.getAxisLabels().setTitle('Filtered Accelerometer Readings');
+  imuAccelPlot.plot.getAxisLabels().setYLabel('Acceleration (g)');
+  imuAccelPlot.plot.getAxisLabels().setXLabel('Monotonic Reading Time (sec)');
+
+  const imuAccelX = imuAccelPlot.addMessageLine(imu, ['accelerometer_x']);
+  imuAccelX.setColor(RED);
+  imuAccelX.setDrawLine(false);
+  const imuAccelY = imuAccelPlot.addMessageLine(imu, ['accelerometer_y']);
+  imuAccelY.setColor(GREEN);
+  imuAccelY.setDrawLine(false);
+  const imuAccelZ = imuAccelPlot.addMessageLine(imu, ['accelerometer_z']);
+  imuAccelZ.setColor(BLUE);
+  imuAccelZ.setDrawLine(false);
+
+  const imuAccelXFiltered = imuAccelPlot.addMessageLine(imu, ['accelerometer_x_filtered']);
+  imuAccelXFiltered.setColor(RED);
+  imuAccelXFiltered.setPointSize(0);
+  const imuAccelYFiltered = imuAccelPlot.addMessageLine(imu, ['accelerometer_y_filtered']);
+  imuAccelYFiltered.setColor(GREEN);
+  imuAccelYFiltered.setPointSize(0);
+  const imuAccelZFiltered = imuAccelPlot.addMessageLine(imu, ['accelerometer_z_filtered']);
+  imuAccelZFiltered.setColor(BLUE);
+  imuAccelZFiltered.setPointSize(0);
+
+  const expectedAccelX = imuAccelPlot.addMessageLine(
+      status, ['down_estimator', 'expected_accel_x']);
+  expectedAccelX.setColor(RED);
+  expectedAccelX.setPointSize(0);
+  const expectedAccelY = imuAccelPlot.addMessageLine(
+      status, ['down_estimator', 'expected_accel_y']);
+  expectedAccelY.setColor(GREEN);
+  expectedAccelY.setPointSize(0);
+  const expectedAccelZ = imuAccelPlot.addMessageLine(
+      status, ['down_estimator', 'expected_accel_z']);
+  expectedAccelZ.setColor(BLUE);
+  expectedAccelZ.setPointSize(0);
+
+  const gyroPlot = aosPlotter.addPlot(element, [0, height * 5], [width, height]);
+  gyroPlot.plot.getAxisLabels().setTitle('Gyro Readings');
+  gyroPlot.plot.getAxisLabels().setYLabel('Angular Velocity (rad / sec)');
+  gyroPlot.plot.getAxisLabels().setXLabel('Monotonic Reading Time (sec)');
+
+  const gyroZeroX =
+      gyroPlot.addMessageLine(status, ['zeroing', 'gyro_x_average']);
+  gyroZeroX.setColor(RED);
+  gyroZeroX.setPointSize(0);
+  gyroZeroX.setLabel('Gyro X Zero');
+  const gyroZeroY =
+      gyroPlot.addMessageLine(status, ['zeroing', 'gyro_y_average']);
+  gyroZeroY.setColor(GREEN);
+  gyroZeroY.setPointSize(0);
+  gyroZeroY.setLabel('Gyro Y Zero');
+  const gyroZeroZ =
+      gyroPlot.addMessageLine(status, ['zeroing', 'gyro_z_average']);
+  gyroZeroZ.setColor(BLUE);
+  gyroZeroZ.setPointSize(0);
+  gyroZeroZ.setLabel('Gyro Z Zero');
+
+  const gyroX = gyroPlot.addMessageLine(imu, ['gyro_x']);
+  gyroX.setColor(RED);
+  const gyroY = gyroPlot.addMessageLine(imu, ['gyro_y']);
+  gyroY.setColor(GREEN);
+  const gyroZ = gyroPlot.addMessageLine(imu, ['gyro_z']);
+  gyroZ.setColor(BLUE);
+
+  const statePlot = aosPlotter.addPlot(element, [0, height * 6], [width, height / 2]);
+  statePlot.plot.getAxisLabels().setTitle('IMU State');
+  statePlot.plot.getAxisLabels().setXLabel('Monotonic Sent Time (sec)');
+
+  const zeroedLine =
+      statePlot.addMessageLine(status, ['zeroing', 'zeroed']);
+  zeroedLine.setColor(RED);
+  zeroedLine.setDrawLine(false);
+  const consecutiveStill =
+      statePlot.addMessageLine(status, ['down_estimator', 'consecutive_still']);
+  consecutiveStill.setColor(BLUE);
+  consecutiveStill.setPointSize(0);
+  const faultedLine =
+  statePlot.addMessageLine(status, ['zeroing', 'faulted']);
+  faultedLine.setColor(GREEN);
+  faultedLine.setPointSize(0);
+}
diff --git a/frc971/wpilib/imu_plot_utils.ts b/frc971/wpilib/imu_plot_utils.ts
index 4e6c371..8193c29 100644
--- a/frc971/wpilib/imu_plot_utils.ts
+++ b/frc971/wpilib/imu_plot_utils.ts
@@ -10,10 +10,17 @@
 import IMUValuesBatch = imu.frc971.IMUValuesBatch;
 import IMUValues = imu.frc971.IMUValues;
 
+const FILTER_WINDOW_SIZE = 100;
+
 export class ImuMessageHandler extends MessageHandler {
+  // Calculated magnitude of the measured acceleration from the IMU.
+  private acceleration_magnitudes: number[] = [];
   constructor(private readonly schema: Schema) {
     super(schema);
   }
+  private readScalar(table: Table, fieldName: string): number {
+    return this.parser.readScalar(table, fieldName);
+  }
   addMessage(data: Uint8Array, time: number): void {
     const batch = IMUValuesBatch.getRootAsIMUValuesBatch(
         new ByteBuffer(data) as unknown as flatbuffers.ByteBuffer);
@@ -27,8 +34,44 @@
         console.log(this.parser.toObject(table));
         continue;
       }
-      this.messages.push(new TimestampedMessage(
-          table, message.monotonicTimestampNs().toFloat64() * 1e-9));
+      const time = message.monotonicTimestampNs().toFloat64() * 1e-9;
+      this.messages.push(new TimestampedMessage(table, time));
+      this.acceleration_magnitudes.push(time);
+      this.acceleration_magnitudes.push(Math.hypot(
+          message.accelerometerX(), message.accelerometerY(),
+          message.accelerometerZ()));
+    }
+  }
+
+  // Computes a moving average for a given input, using a basic window centered
+  // on each value.
+  private movingAverageCentered(input: Float32Array): Float32Array {
+    const num_measurements = input.length / 2;
+    const filtered_measurements = new Float32Array(input);
+    for (let ii = 0; ii < num_measurements; ++ii) {
+      let sum = 0;
+      let count = 0;
+      for (let jj = Math.max(0, Math.ceil(ii - FILTER_WINDOW_SIZE / 2));
+           jj < Math.min(num_measurements, ii + FILTER_WINDOW_SIZE / 2); ++jj) {
+        sum += input[jj * 2 + 1];
+        ++count;
+      }
+      filtered_measurements[ii * 2 + 1] = sum / count;
+    }
+    return new Float32Array(filtered_measurements);
+  }
+
+  getField(field: string[]): Float32Array {
+    // Any requested input that ends with "_filtered" will get a moving average
+    // applied to the original field.
+    const filtered_suffix = "_filtered";
+    if (field[0] == "acceleration_magnitude") {
+      return new Float32Array(this.acceleration_magnitudes);
+    } else if (field[0].endsWith(filtered_suffix)) {
+      return this.movingAverageCentered(this.getField(
+          [field[0].slice(0, field[0].length - filtered_suffix.length)]));
+    } else {
+      return super.getField(field);
     }
   }
 }