Make drivetrain_test_lib handle IMU transforms

This also updates the 2020 imu transform to better match the 2020
drivetrain. I haven't actually checked the yaw orientation of the IMU
though.

Change-Id: Ib036705269dfc467a60714c456751a6cfca8a67e
diff --git a/frc971/analysis/plot_configs/localizer.pb b/frc971/analysis/plot_configs/localizer.pb
index f81ec89..644b516 100644
--- a/frc971/analysis/plot_configs/localizer.pb
+++ b/frc971/analysis/plot_configs/localizer.pb
@@ -23,6 +23,16 @@
   axes {
     line {
       y_signal {
+        channel: "DrivetrainStatus"
+        field: "trajectory_logging.y"
+      }
+      x_signal {
+        channel: "DrivetrainStatus"
+        field: "trajectory_logging.x"
+      }
+    }
+    line {
+      y_signal {
         channel: "DrivetrainTruthStatus"
         field: "y"
       }
@@ -31,7 +41,6 @@
         field: "x"
       }
     }
-    share_x_axis: false
     line {
       y_signal {
         channel: "DrivetrainStatus"
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index 9fef8f9..30abc87 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -691,6 +691,7 @@
         "improved_down_estimator.h",
     ],
     deps = [
+        ":drivetrain_config",
         ":drivetrain_status_fbs",
         "//aos/events:event_loop",
         "//frc971/control_loops:control_loops_fbs",
@@ -707,6 +708,7 @@
         "improved_down_estimator_test.cc",
     ],
     deps = [
+        ":drivetrain_test_lib",
         "//aos/testing:googletest",
         "//aos/testing:random_seed",
         "//frc971/control_loops/drivetrain:improved_down_estimator",
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 77c52c4..1c0ac4f 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -41,6 +41,7 @@
       gyro_reading_fetcher_(
           event_loop->MakeFetcher<::frc971::sensors::GyroReading>(
               "/drivetrain")),
+      down_estimator_(dt_config),
       localizer_(localizer),
       kf_(dt_config_.make_kf_drivetrain_loop()),
       dt_openloop_(dt_config_, &kf_),
@@ -177,6 +178,9 @@
       case IMUType::IMU_Y:
         last_accel_ = -imu_values_fetcher_->accelerometer_y();
         break;
+      case IMUType::IMU_Z:
+        last_accel_ = imu_values_fetcher_->accelerometer_z();
+        break;
     }
   }
 
diff --git a/frc971/control_loops/drivetrain/drivetrain_config.h b/frc971/control_loops/drivetrain/drivetrain_config.h
index 5a29008..a33ff22 100644
--- a/frc971/control_loops/drivetrain/drivetrain_config.h
+++ b/frc971/control_loops/drivetrain/drivetrain_config.h
@@ -37,6 +37,7 @@
   IMU_X = 0,          // Use the x-axis of the IMU.
   IMU_Y = 1,          // Use the y-axis of the IMU.
   IMU_FLIPPED_X = 2,  // Use the flipped x-axis of the IMU.
+  IMU_Z = 3,          // Use the z-axis of the IMU.
 };
 
 template <typename Scalar = double>
diff --git a/frc971/control_loops/drivetrain/drivetrain_test_lib.cc b/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
index 336025a..6a90122 100644
--- a/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_test_lib.cc
@@ -190,16 +190,23 @@
   auto builder = imu_sender_.MakeBuilder();
   frc971::IMUValues::Builder imu_builder =
       builder.MakeBuilder<frc971::IMUValues>();
-  imu_builder.add_gyro_x(0.0);
-  imu_builder.add_gyro_y(0.0);
-  imu_builder.add_gyro_z(
-      (drivetrain_plant_.X(3, 0) - drivetrain_plant_.X(1, 0)) /
-      (dt_config_.robot_radius * 2.0));
+  const Eigen::Vector3d gyro =
+      dt_config_.imu_transform.inverse() *
+      Eigen::Vector3d(0.0, 0.0,
+                      (drivetrain_plant_.X(3, 0) - drivetrain_plant_.X(1, 0)) /
+                          (dt_config_.robot_radius * 2.0));
+  imu_builder.add_gyro_x(gyro.x());
+  imu_builder.add_gyro_y(gyro.y());
+  imu_builder.add_gyro_z(gyro.z());
   // Acceleration due to gravity, in m/s/s.
   constexpr double kG = 9.807;
-  imu_builder.add_accelerometer_x(last_acceleration_.x() / kG);
-  imu_builder.add_accelerometer_y(last_acceleration_.y() / kG);
-  imu_builder.add_accelerometer_z(1.0);
+  const Eigen::Vector3d accel =
+      dt_config_.imu_transform.inverse() *
+      Eigen::Vector3d(last_acceleration_.x() / kG, last_acceleration_.y() / kG,
+                      1.0);
+  imu_builder.add_accelerometer_x(accel.x());
+  imu_builder.add_accelerometer_y(accel.y());
+  imu_builder.add_accelerometer_z(accel.z());
   imu_builder.add_monotonic_timestamp_ns(
       std::chrono::duration_cast<std::chrono::nanoseconds>(
           event_loop_->monotonic_now().time_since_epoch())
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.h b/frc971/control_loops/drivetrain/improved_down_estimator.h
index 27de867..c1c8e1b 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator.h
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.h
@@ -6,6 +6,7 @@
 
 #include "aos/events/event_loop.h"
 #include "aos/time/time.h"
+#include "frc971/control_loops/drivetrain/drivetrain_config.h"
 #include "frc971/control_loops/drivetrain/drivetrain_status_generated.h"
 #include "frc971/control_loops/runge_kutta.h"
 #include "glog/logging.h"
@@ -276,6 +277,8 @@
 // accelerometer calibration).
 class DrivetrainUkf : public QuaternionUkf {
  public:
+  DrivetrainUkf(const DrivetrainConfig<double> &dt_config)
+      : QuaternionUkf(dt_config.imu_transform) {}
   // UKF for http://kodlab.seas.upenn.edu/uploads/Arun/UKFpaper.pdf
   // Reference in case the link is dead:
   // Kraft, Edgar. "A quaternion-based unscented Kalman filter for orientation
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
index 86ef5e9..2edce5a 100644
--- a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
+++ b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
@@ -4,6 +4,7 @@
 #include <random>
 
 #include "aos/testing/random_seed.h"
+#include "frc971/control_loops/drivetrain/drivetrain_test_lib.h"
 #include "frc971/control_loops/drivetrain/improved_down_estimator.h"
 #include "glog/logging.h"
 #include "gtest/gtest.h"
@@ -91,7 +92,8 @@
 }
 
 TEST(DownEstimatorTest, UkfConstantRotation) {
-  drivetrain::DrivetrainUkf dtukf;
+  drivetrain::DrivetrainUkf dtukf(
+      drivetrain::testing::GetTestDrivetrainConfig());
   const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   EXPECT_EQ(0.0,
             (Eigen::Vector3d(0.0, 0.0, 1.0) - dtukf.H(dtukf.X_hat().coeffs()))
@@ -112,7 +114,8 @@
 
 // Tests that the euler angles in the status message are correct.
 TEST(DownEstimatorTest, UkfEulerStatus) {
-  drivetrain::DrivetrainUkf dtukf;
+  drivetrain::DrivetrainUkf dtukf(
+      drivetrain::testing::GetTestDrivetrainConfig());
   const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   const Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
   const Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
@@ -167,7 +170,8 @@
 // that we are slightly rotated, that we eventually adjust our estimate to be
 // correct.
 TEST(DownEstimatorTest, UkfAccelCorrectsBias) {
-  drivetrain::DrivetrainUkf dtukf;
+  drivetrain::DrivetrainUkf dtukf(
+      drivetrain::testing::GetTestDrivetrainConfig());
   const Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
   Eigen::Matrix<double, 3, 1> measurement;
   // Supply the accelerometer with a slightly off reading to ensure that we
@@ -192,7 +196,8 @@
 // that we are slightly rotated, that we eventually adjust our estimate to be
 // correct.
 TEST(DownEstimatorTest, UkfIgnoreBadAccel) {
-  drivetrain::DrivetrainUkf dtukf;
+  drivetrain::DrivetrainUkf dtukf(
+      drivetrain::testing::GetTestDrivetrainConfig());
   const Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
   Eigen::Matrix<double, 3, 1> measurement;
   // Set up a scenario where, if we naively took the accelerometer readings, we
diff --git a/y2020/control_loops/drivetrain/drivetrain_base.cc b/y2020/control_loops/drivetrain/drivetrain_base.cc
index 62679bc..b919902 100644
--- a/y2020/control_loops/drivetrain/drivetrain_base.cc
+++ b/y2020/control_loops/drivetrain/drivetrain_base.cc
@@ -25,8 +25,8 @@
   static DrivetrainConfig<double> kDrivetrainConfig{
       ::frc971::control_loops::drivetrain::ShifterType::SIMPLE_SHIFTER,
       ::frc971::control_loops::drivetrain::LoopType::CLOSED_LOOP,
-      ::frc971::control_loops::drivetrain::GyroType::IMU_Z_GYRO,
-      ::frc971::control_loops::drivetrain::IMUType::IMU_FLIPPED_X,
+      ::frc971::control_loops::drivetrain::GyroType::IMU_X_GYRO,
+      ::frc971::control_loops::drivetrain::IMUType::IMU_Z,
 
       drivetrain::MakeDrivetrainLoop,
       drivetrain::MakeVelocityDrivetrainLoop,
@@ -53,8 +53,8 @@
       1.2 /* quickturn_wheel_multiplier */,
       1.2 /* wheel_multiplier */,
       true /*pistol_grip_shift_enables_line_follow*/,
-      (Eigen::Matrix<double, 3, 3>() << 0.0, -1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0,
-       1.0)
+      (Eigen::Matrix<double, 3, 3>() << 0.0, 0.0, 1.0, 0.0, -1.0, 0.0, 1.0, 0.0,
+       0.0)
           .finished() /*imu_transform*/,
   };
 
diff --git a/y2020/control_loops/drivetrain/drivetrain_replay_test.cc b/y2020/control_loops/drivetrain/drivetrain_replay_test.cc
index 3e391c8..6033184 100644
--- a/y2020/control_loops/drivetrain/drivetrain_replay_test.cc
+++ b/y2020/control_loops/drivetrain/drivetrain_replay_test.cc
@@ -52,13 +52,19 @@
         reader_.event_loop_factory()->MakeEventLoop("drivetrain", roborio_);
     drivetrain_event_loop_->SkipTimingReport();
 
+    frc971::control_loops::drivetrain::DrivetrainConfig<double> config =
+        GetDrivetrainConfig();
+    // Make the modification required to the imu transform to work with the 2016
+    // logs...
+    config.imu_transform << 0.0, -1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0;
+    config.gyro_type = frc971::control_loops::drivetrain::GyroType::IMU_Z_GYRO;
+
     localizer_ =
         std::make_unique<frc971::control_loops::drivetrain::DeadReckonEkf>(
-            drivetrain_event_loop_.get(), GetDrivetrainConfig());
+            drivetrain_event_loop_.get(), config);
     drivetrain_ =
         std::make_unique<frc971::control_loops::drivetrain::DrivetrainLoop>(
-            GetDrivetrainConfig(), drivetrain_event_loop_.get(),
-            localizer_.get());
+            config, drivetrain_event_loop_.get(), localizer_.get());
 
     test_event_loop_ =
         reader_.event_loop_factory()->MakeEventLoop("drivetrain", roborio_);
@@ -88,7 +94,7 @@
   ASSERT_TRUE(status_fetcher_->has_x());
   ASSERT_TRUE(status_fetcher_->has_y());
   ASSERT_TRUE(status_fetcher_->has_theta());
-  EXPECT_LT(std::abs(status_fetcher_->x()), 0.1);
+  EXPECT_LT(std::abs(status_fetcher_->x()), 0.25);
   // Because the encoders should not be affecting the y or yaw axes, expect a
   // reasonably precise result (although, since this is a real worl dtest, the
   // robot probably did actually move be some non-zero amount).
diff --git a/y2020/control_loops/drivetrain/localizer_test.cc b/y2020/control_loops/drivetrain/localizer_test.cc
index c438af5..c47faa7 100644
--- a/y2020/control_loops/drivetrain/localizer_test.cc
+++ b/y2020/control_loops/drivetrain/localizer_test.cc
@@ -92,7 +92,6 @@
 namespace chrono = std::chrono;
 using frc971::control_loops::drivetrain::testing::DrivetrainSimulation;
 using frc971::control_loops::drivetrain::DrivetrainLoop;
-using frc971::control_loops::drivetrain::testing::GetTestDrivetrainConfig;
 using aos::monotonic_clock;
 
 class LocalizedDrivetrainTest : public aos::testing::ControlLoopTest {
@@ -418,6 +417,22 @@
   EXPECT_TRUE(VerifyEstimatorAccurate(5e-4));
 }
 
+// Tests that we can drive in a straight line and have things estimate
+// correctly.
+TEST_F(LocalizedDrivetrainTest, NoCameraUpdateStraightLine) {
+  SetEnabled(true);
+  set_enable_cameras(false);
+  EXPECT_TRUE(VerifyEstimatorAccurate(1e-7));
+
+  SendGoal(1.0, 1.0);
+
+  RunFor(chrono::seconds(1));
+  VerifyNearGoal();
+  // Due to accelerometer drift, the straight-line driving tends to be less
+  // accurate...
+  EXPECT_TRUE(VerifyEstimatorAccurate(0.1));
+}
+
 // Tests that camera updates with a perfect models results in no errors.
 TEST_F(LocalizedDrivetrainTest, PerfectCameraUpdate) {
   SetEnabled(true);