Bring in Quaternion UKF

This pulls in and cleans up a bit the implementation of the Quaternion
UKF that Austin pulled in. This also makes a couple of minor changes to
the UKF:
1) Explicitly handles large covariances by clipping the covariance.
2) Only accepts accelerometer measurements which are close to exactly
   oneone g in magnitude.

Change-Id: Id1619406ebddf44570e2542dfd5a74620612e9e5
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index f263d14..6f46965 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -642,3 +642,31 @@
         "arm": [],
     }),
 )
+
+cc_library(
+    name = "improved_down_estimator",
+    srcs = [
+        "improved_down_estimator.cc",
+    ],
+    hdrs = [
+        "improved_down_estimator.h",
+    ],
+    deps = [
+        "//frc971/control_loops:runge_kutta",
+        "@com_github_google_glog//:glog",
+        "@org_tuxfamily_eigen//:eigen",
+    ],
+)
+
+cc_test(
+    name = "improved_down_estimator_test",
+    srcs = [
+        "improved_down_estimator_test.cc",
+    ],
+    deps = [
+        "//aos/testing:googletest",
+        "//aos/testing:random_seed",
+        "//frc971/control_loops/drivetrain:improved_down_estimator",
+        "@org_tuxfamily_eigen//:eigen",
+    ],
+)
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.cc b/frc971/control_loops/drivetrain/improved_down_estimator.cc
new file mode 100644
index 0000000..40dc8e3
--- /dev/null
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.cc
@@ -0,0 +1,287 @@
+#include "frc971/control_loops/drivetrain/improved_down_estimator.h"
+
+#include "Eigen/Dense"
+#include "Eigen/Geometry"
+
+namespace frc971 {
+namespace control_loops {
+namespace drivetrain {
+
+Eigen::Matrix<double, 4, 1> ToQuaternionFromRotationVector(
+    const Eigen::Matrix<double, 3, 1> &X, const double max_angle_cap) {
+  const double unclipped_angle = X.norm();
+  const double angle_scalar =
+      (unclipped_angle > max_angle_cap) ? max_angle_cap / unclipped_angle : 1.0;
+  const double angle = unclipped_angle * angle_scalar;
+  const double half_angle = angle * 0.5;
+
+  const double half_angle_squared = half_angle * half_angle;
+
+  // sin(x)/x = 1
+  double sinx_x = 1.0;
+
+  // - x^2/3!
+  double value = half_angle_squared / 6.0;
+  sinx_x -= value;
+
+  // + x^4/5!
+  value = value * half_angle_squared / 20.0;
+  sinx_x += value;
+
+  // - x^6/7!
+  value = value * half_angle_squared / (6.0 * 7.0);
+  sinx_x -= value;
+
+  // + x^8/9!
+  value = value * half_angle_squared / (8.0 * 9.0);
+  sinx_x += value;
+
+  // - x^10/11!
+  value = value * half_angle_squared / (10.0 * 11.0);
+  sinx_x -= value;
+
+  // + x^12/13!
+  value = value * half_angle_squared / (12.0 * 13.0);
+  sinx_x += value;
+
+  // - x^14/15!
+  value = value * half_angle_squared / (14.0 * 15.0);
+  sinx_x -= value;
+
+  // + x^16/17!
+  value = value * half_angle_squared / (16.0 * 17.0);
+  sinx_x += value;
+
+  // To plot the residual in matplotlib, run:
+  // import numpy
+  // import scipy
+  // from matplotlib import pyplot
+  // x = numpy.arange(-numpy.pi, numpy.pi, 0.01)
+  // pyplot.plot(x, 1 - x**2 / scipy.misc.factorial(3) +
+  //                   x**4 / scipy.misc.factorial(5) -
+  //                   x**6 / scipy.misc.factorial(7) +
+  //                   x**8 / scipy.misc.factorial(9) -
+  //                   x ** 10 / scipy.misc.factorial(11) +
+  //                   x ** 12 / scipy.misc.factorial(13) -
+  //                   x ** 14 / scipy.misc.factorial(15) +
+  //                   x ** 16 / scipy.misc.factorial(17) -
+  //                   numpy.sin(x) / x)
+
+  const double scalar = sinx_x * 0.5;
+
+  Eigen::Matrix<double, 4, 1> result;
+  result.block<3, 1>(0, 0) = X * scalar * angle_scalar;
+  result(3, 0) = std::cos(half_angle);
+  return result;
+}
+
+inline Eigen::Matrix<double, 4, 1> MaybeFlipX(
+    const Eigen::Matrix<double, 4, 1> &X) {
+  if (X(3, 0) < 0.0) {
+    return -X;
+  } else {
+    return X;
+  }
+}
+
+Eigen::Matrix<double, 3, 1> ToRotationVectorFromQuaternion(
+    const Eigen::Matrix<double, 4, 1> &X) {
+  // TODO(austin): Verify we still need it.
+  const Eigen::Matrix<double, 4, 1> corrected_X = MaybeFlipX(X);
+  const double half_angle =
+      std::atan2(corrected_X.block<3, 1>(0, 0).norm(), corrected_X(3, 0));
+
+  const double half_angle_squared = half_angle * half_angle;
+
+  // TODO(austin): We are doing a division at the end of this. Do the taylor
+  // series expansion of x/sin(x) instead to avoid this.
+
+  // sin(x)/x = 1
+  double sinx_x = 1.0;
+
+  // - x^2/3!
+  double value = half_angle_squared / 6.0;
+  sinx_x -= value;
+
+  // + x^4/5!
+  value = value * half_angle_squared / 20.0;
+  sinx_x += value;
+
+  // - x^6/7!
+  value = value * half_angle_squared / (6.0 * 7.0);
+  sinx_x -= value;
+
+  // + x^8/9!
+  value = value * half_angle_squared / (8.0 * 9.0);
+  sinx_x += value;
+
+  // - x^10/11!
+  value = value * half_angle_squared / (10.0 * 11.0);
+  sinx_x -= value;
+
+  // + x^12/13!
+  value = value * half_angle_squared / (12.0 * 13.0);
+  sinx_x += value;
+
+  // - x^14/15!
+  value = value * half_angle_squared / (14.0 * 15.0);
+  sinx_x -= value;
+
+  // + x^16/17!
+  value = value * half_angle_squared / (16.0 * 17.0);
+  sinx_x += value;
+
+  const double scalar = 2.0 / sinx_x;
+
+  return corrected_X.block<3, 1>(0, 0) * scalar;
+}
+
+// States are X_hat_bar (position estimate) and P (Covariance)
+
+void QuaternionUkf::Predict(const Eigen::Matrix<double, 3, 1> &U,
+                            const Eigen::Matrix<double, 3, 1> &measurement) {
+  // Compute the sigma points.
+  // Our system is pretty linear. The traditional way of dealing with process
+  // noise is to augment your state vector with the mean of the process noise,
+  // and augment your covariance matrix with the covariance of your process
+  // noise. Sigma points are then computed. These points are then propegated
+  // through the model. This ends up effectively meaning that perturbations
+  // from the unaugmented state with covariance P are propegated through the
+  // model, and points which are at the mean but with perturbations to simulated
+  // process noise are propegated through the system. The covariance is then
+  // calculated from this set of points, and works out to have a covariance of
+  // essentially P + Q.
+  //
+  // Since our noise is just additive, and quaternian rotation preserves
+  // distance, we can add our noise first and it'll be a good representation of
+  // our distance. This will reduce the number of math operations we need to
+  // do. If we break this assumption in the future by adding a nonlinear model
+  // somewhere in this system, we'll have to revisit this assumption.
+
+  // Now, compute the actual sigma points using the columns of S as the
+  // pertubation vectors. The last point is the original mean.
+  const Eigen::Matrix<double, 4, 3 * 2 + 1> X =
+      GenerateSigmaPoints(X_hat_, P_ + Q_);
+
+  // Now, compute Y, the sigma points which have been propegated forwards by the
+  // model.
+  Eigen::Matrix<double, 4, 3 * 2 + 1> Y;
+  for (int i = 0; i < Y.cols(); ++i) {
+    // Y = Transformed sigma points
+    Y.col(i) = A(X.col(i), U);
+  }
+
+  // We now have the sigma points after the model update.
+  // Compute the mean of the transformed sigma point
+  X_hat_ = Eigen::Quaternion<double>(QuaternionMean(Y));
+
+  // And the covariance.
+  Eigen::Matrix<double, 3, 2 * 3 + 1> Wprime;
+  Eigen::Matrix<double, 3, 3> P_prior =
+      ComputeQuaternionCovariance(X_hat_, Y, &Wprime);
+
+  // If the only obvious acceleration is that due to gravity, then accept the
+  // measurement.
+  constexpr double kUseAccelThreshold = 0.1;
+  if (std::abs(measurement.squaredNorm() - 1.0) < kUseAccelThreshold) {
+    P_ = P_prior;
+    return;
+  }
+
+  // TODO(austin): Maybe re-calculate the sigma points here before transforming
+  // them?  Otherwise we can't cleanly decouple the model and measurement updates.
+
+  // Apply the measurement transform to all the sigma points to get a
+  // representation of the distribution of the measurement.
+  Eigen::Matrix<double, kNumMeasurements, 3 * 2 + 1> Z;
+  Z_hat_.setZero();
+  for (int i = 0; i < Z.cols(); ++i) {
+    Z.col(i) = H(Y.col(i));
+
+    // Compute the mean in addition.
+    Z_hat_ += Z.col(i) / Z.cols();
+  }
+
+  // Now compute the measurement covariance.
+  Eigen::Matrix<double, 3, 3> P_zz;
+  P_zz.setZero();
+  Eigen::Matrix<double, 3, 2 * 3 + 1> Zprime;
+  for (int i = 0; i < 7; ++i) {
+    // Compute the error vector for each sigma point.
+    Eigen::Matrix<double, 3, 1> Zprimei = Z.col(i) - Z_hat_;
+
+    // Now, compute the contribution of this sigma point to P_zz.
+    P_zz += 1.0 / 12.0 * Zprimei * Zprimei.transpose();
+    // Save the error for the cross-correlation matrix.
+    Zprime.col(i) = Zprimei;
+  }
+
+  // Compute the measurement error and innovation uncertanty.
+  const Eigen::Matrix<double, kNumMeasurements, kNumMeasurements> P_vv =
+      P_zz + R_;
+
+  // Now compute the cross correlation matrix P_xz.
+  Eigen::Matrix<double, 3, 3> P_xz;
+  P_xz.setZero();
+  for (int i = 0; i < 7; ++i) {
+    // Now, compute the contribution of this sigma point to P_prior.
+    P_xz += 1.0 / 12.0 * Wprime.col(i) * Zprime.col(i).transpose();
+  }
+
+  // Compute the kalman gain.
+  const Eigen::Matrix<double, 3, kNumMeasurements> K =
+      P_xz * P_vv.inverse();
+
+  // Update X_hat and the covariance P
+  X_hat_ = X_hat_ * Eigen::Quaternion<double>(ToQuaternionFromRotationVector(
+                        K * (measurement - Z_hat_)));
+  P_ = P_prior - K * P_vv * K.transpose();
+}
+
+Eigen::Matrix<double, 3, 3> ComputeQuaternionCovariance(
+    const Eigen::Quaternion<double> &mean,
+    const Eigen::Matrix<double, 4, 7> &points,
+    Eigen::Matrix<double, 3, 7> *residual) {
+  Eigen::Matrix<double, 3, 3> P_prior;
+  P_prior.setZero();
+
+  for (int i = 0; i < 7; ++i) {
+    // Compute the error vector for each sigma point.
+    Eigen::Matrix<double, 3, 1> Wprimei = ToRotationVectorFromQuaternion(
+        Eigen::Quaternion<double>(mean).conjugate() *
+        Eigen::Quaternion<double>(points.col(i)));
+    // Now, compute the contribution of this sigma point to P_prior.
+    P_prior += 1.0 / 6.0 * (Wprimei * Wprimei.transpose());
+    // Save the error for the cross-correlation matrix.
+    residual->col(i) = Wprimei;
+  }
+  return P_prior / 2.0;
+}
+
+Eigen::Matrix<double, 4, 3 * 2 + 1> GenerateSigmaPoints(
+    const Eigen::Quaternion<double> &mean,
+    const Eigen::Matrix<double, 3, 3> &covariance) {
+  // Take the matrix square root.
+  Eigen::Matrix<double, 3, 3> S = covariance.llt().matrixL();
+
+  S *= std::sqrt(2.0 * 3.0);
+  // TODO(austin): Make sure the sigma points aren't outside +- PI/2.0.
+  // Otherwise they wrap on themselves and we get a mess.
+
+  // Now, compute the actual sigma points using the columns of S as the
+  // pertubation vectors. The last point is the original mean.
+  Eigen::Matrix<double, 4, 3 * 2 + 1> X;
+  for (int i = 0; i < 3; ++i) {
+    Eigen::Quaternion<double> perturbation(
+        ToQuaternionFromRotationVector(S.col(i), M_PI_2));
+
+    X.col(i * 2) = (mean * perturbation).coeffs();
+    X.col(i * 2 + 1) = (mean * perturbation.conjugate()).coeffs();
+  }
+  X.col(6) = mean.coeffs();
+  return X;
+}
+
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace frc971
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator.h b/frc971/control_loops/drivetrain/improved_down_estimator.h
new file mode 100644
index 0000000..2b47f9e
--- /dev/null
+++ b/frc971/control_loops/drivetrain/improved_down_estimator.h
@@ -0,0 +1,219 @@
+#ifndef FRC971_CONTROL_LOOPS_DRIVETRAIN_IMPROVED_DOWN_ESTIMATOR_H_
+#define FRC971_CONTROL_LOOPS_DRIVETRAIN_IMPROVED_DOWN_ESTIMATOR_H_
+
+#include "Eigen/Dense"
+#include "Eigen/Geometry"
+
+#include "frc971/control_loops/runge_kutta.h"
+#include "glog/logging.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace drivetrain {
+
+// Function to compute the quaternion average of a bunch of quaternions. Each
+// column in the input matrix is a quaternion (optionally scaled by it's
+// weight).
+template <int SM>
+Eigen::Matrix<double, 4, 1> QuaternionMean(
+    Eigen::Matrix<double, 4, SM> input) {
+  // Algorithm to compute the average of a bunch of quaternions:
+  // http://www.acsu.buffalo.edu/~johnc/ave_quat07.pdf
+
+  Eigen::Matrix<double, 4, 4> m = input * input.transpose();
+
+  Eigen::EigenSolver<Eigen::Matrix<double, 4, 4>> solver;
+  solver.compute(m);
+
+  Eigen::EigenSolver<Eigen::Matrix<double, 4, 4>>::EigenvectorsType
+      eigenvectors = solver.eigenvectors();
+  Eigen::EigenSolver<Eigen::Matrix<double, 4, 4>>::EigenvalueType eigenvalues =
+      solver.eigenvalues();
+
+  int max_index = 0;
+  double max_eigenvalue = 0.0;
+  for (int i = 0; i < 4; ++i) {
+    const double eigenvalue = std::abs(eigenvalues(i, 0));
+    if (eigenvalue > max_eigenvalue) {
+      max_eigenvalue = eigenvalue;
+      max_index = i;
+    }
+  }
+
+  // Assume that there shouldn't be any imaginary components to the eigenvector.
+  // I can't prove this is true, but everyone else seems to assume it...
+  // TODO(james): Handle this more rigorously.
+  for (int i = 0; i < 4; ++i) {
+    CHECK_LT(eigenvectors(i, max_index).imag(), 1e-4)
+        << eigenvectors(i, max_index);
+  }
+  return eigenvectors.col(max_index).real().normalized();
+}
+
+// Converts from a quaternion to a rotation vector, where the rotation vector's
+// direction represents the axis to rotate around and its magnitude represents
+// the number of radians to rotate.
+Eigen::Matrix<double, 3, 1> ToRotationVectorFromQuaternion(
+    const Eigen::Matrix<double, 4, 1> &X);
+
+inline Eigen::Matrix<double, 3, 1> ToRotationVectorFromQuaternion(
+    const Eigen::Quaternion<double> &X) {
+  return ToRotationVectorFromQuaternion(X.coeffs());
+};
+
+// Converts from a rotation vector to a quaternion. If you supply max_angle_cap,
+// then the rotation vector's magnitude will be clipped to be no more than
+// max_angle_cap before being converted to a quaternion.
+Eigen::Matrix<double, 4, 1> ToQuaternionFromRotationVector(
+    const Eigen::Matrix<double, 3, 1> &X,
+    const double max_angle_cap = std::numeric_limits<double>::infinity());
+
+// Generates the sigma points to use in the UKF given the current estimate and
+// covariance.
+Eigen::Matrix<double, 4, 3 * 2 + 1> GenerateSigmaPoints(
+    const Eigen::Quaternion<double> &mean,
+    const Eigen::Matrix<double, 3, 3> &covariance);
+
+// Computes the covariance of the noise given the mean and the transformed sigma
+// points. The residual corresponds with the W' variable from the original
+// paper.
+Eigen::Matrix<double, 3, 3> ComputeQuaternionCovariance(
+    const Eigen::Quaternion<double> &mean,
+    const Eigen::Matrix<double, 4, 7> &points,
+    Eigen::Matrix<double, 3, 7> *residual);
+
+// This class provides a quaternion-based Kalman filter for estimating
+// orientation using a 3-axis gyro and 3-axis accelerometer. It does leave open
+// the option of overridding the system process model and the function used to
+// calculate the expected measurement (which is relevant if, e.g., the IMU is
+// not mounted horizontally in the robot).
+class QuaternionUkf {
+ public:
+  // The state is just a quaternion representing the current robot orientaiton.
+  // The zero/identity quaternion (1, 0, 0, 0) implies that the robot is
+  // position flat on the ground with a heading of zero (which, in our normal
+  // field coordinates, means pointed straight away from our driver's station
+  // wall).
+  constexpr static int kNumStates = 4;
+  // Inputs to the system--we use the (x, y, z)  gyro measurements as the inputs
+  // to the system.
+  constexpr static int kNumInputs = 3;
+  // Measurements to use for correcting the estimated system state. These
+  // correspond to (x, y, z) measurements from the accelerometer.
+  constexpr static int kNumMeasurements = 3;
+  QuaternionUkf() {
+    // TODO(james): Tune the process/measurement noises.
+    R_.setIdentity();
+    R_ /= 100.0;
+
+    Q_.setIdentity();
+    Q_ /= 10000.0;
+
+    // Assume that the robot starts flat on the ground pointed straight forward.
+    X_hat_ = Eigen::Quaternion<double>(1.0, 0.0, 0.0, 0.0);
+
+    // TODO(james): Determine an appropriate starting noise estimate. Probably
+    // not too critical.
+    P_.setIdentity();
+    P_ /= 1000.0;
+  }
+
+  // Handles updating the state of the UKF, given the gyro and accelerometer
+  // measurements.
+  void Predict(const Eigen::Matrix<double, kNumInputs, 1> &U,
+               const Eigen::Matrix<double, 3, 1> &measurement);
+
+  // Returns the updated state for X after one time step, given the current
+  // state and gyro measurements.
+  virtual Eigen::Matrix<double, 4, 1> A(
+      const Eigen::Matrix<double, 4, 1> &X,
+      const Eigen::Matrix<double, kNumInputs, 1> &U) const = 0;
+
+  // Returns the current expected accelerometer measurements given the current
+  // state.
+  virtual Eigen::Matrix<double, 3, 1> H(
+      const Eigen::Matrix<double, 4, 1> &X) const = 0;
+
+  // Returns the current estimate of the robot's orientation. Note that this
+  // filter does not have anything other than the gyro with which to estimate
+  // the robot's yaw heading, and so it may need to be corrected for by upstream
+  // filters.
+  const Eigen::Quaternion<double> &X_hat() const { return X_hat_; }
+
+  Eigen::Matrix<double, 3, 1> Z_hat() const { return Z_hat_; };
+
+ private:
+  // Measurement Noise (Uncertainty)
+  Eigen::Matrix<double, kNumInputs, kNumInputs> R_;
+  // Model noise. Note that both this and P are 3 x 3 matrices, despite the
+  // state having 4 dimensions.
+  Eigen::Matrix<double, 3, 3> Q_;
+  // Current estimate covariance.
+  Eigen::Matrix<double, 3, 3> P_;
+
+  // Current state estimate.
+  Eigen::Quaternion<double> X_hat_;
+
+  // Current expected accelerometer measurement.
+  Eigen::Matrix<double, 3, 1> Z_hat_;
+};
+
+class DrivetrainUkf : public QuaternionUkf {
+ public:
+  constexpr static double kDt = 0.00505;
+
+  // 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
+  // tracking." In Proceedings of the Sixth International Conference of
+  // Information Fusion, vol. 1, pp. 47-54. 2003.
+
+  // A good reference for quaternions is available at
+  // http://www.euclideanspace.com/maths/algebra/realNormedAlgebra/quaternions/
+  //
+  // A good reference for angular velocity vectors with quaternions is at
+  // http://www.euclideanspace.com/physics/kinematics/angularvelocity/
+
+  // Creates a rotational velocity vector to be integrated.
+  //
+  // omega is the rotational velocity vector in body coordinates.
+  // q is a matrix with the compononents of the quaternion in it.
+  //
+  // Returns dq / dt
+  static Eigen::Vector4d QuaternionDerivative(Eigen::Vector3d omega,
+                                              const Eigen::Vector4d &q_matrix) {
+    Eigen::Quaternion<double> q(q_matrix);
+
+    Eigen::Quaternion<double> omega_q;
+    omega_q.w() = 0.0;
+    omega_q.vec() = 0.5 * (q * omega);
+
+    Eigen::Quaternion<double> deriv = omega_q * q;
+    return deriv.coeffs();
+  }
+
+  // Moves the robot by the provided rotation vector (U).
+  Eigen::Matrix<double, kNumStates, 1> A(
+      const Eigen::Matrix<double, kNumStates, 1> &X,
+      const Eigen::Matrix<double, kNumInputs, 1> &U) const override {
+    return RungeKutta(
+        std::bind(&QuaternionDerivative, U, std::placeholders::_1), X, kDt);
+  }
+
+  // Returns the expected accelerometer measurement (which is just going to be
+  // 1g downwards).
+  Eigen::Matrix<double, kNumMeasurements, 1> H(
+      const Eigen::Matrix<double, kNumStates, 1> &X) const override {
+    // TODO(austin): Figure out how to compute what the sensors *should* read.
+    Eigen::Quaternion<double> Xquat(X);
+    Eigen::Matrix<double, 3, 1> gprime =
+        Xquat * Eigen::Matrix<double, 3, 1>(0.0, 0.0, -1.0);
+    return gprime;
+  }
+};
+
+}  // namespace drivetrain
+}  // namespace control_loops
+}  // namespace frc971
+
+#endif  // FRC971_CONTROL_LOOPS_DRIVETRAIN_IMPROVED_DOWN_ESTIMATOR_H_
diff --git a/frc971/control_loops/drivetrain/improved_down_estimator_test.cc b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
new file mode 100644
index 0000000..db7eeee
--- /dev/null
+++ b/frc971/control_loops/drivetrain/improved_down_estimator_test.cc
@@ -0,0 +1,280 @@
+#include "frc971/control_loops/runge_kutta.h"
+
+#include <Eigen/Geometry>
+#include <random>
+
+#include "aos/testing/random_seed.h"
+#include "frc971/control_loops/drivetrain/improved_down_estimator.h"
+#include "glog/logging.h"
+#include "gtest/gtest.h"
+
+namespace frc971 {
+namespace control_loops {
+namespace testing {
+
+// Do a known transformation to see if quaternion integration is working
+// correctly.
+TEST(RungeKuttaTest, QuaternionIntegral) {
+  Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
+  Eigen::Vector3d uy = Eigen::Vector3d::UnitY();
+  Eigen::Vector3d uz = Eigen::Vector3d::UnitZ();
+
+  Eigen::Quaternion<double> q(
+      Eigen::AngleAxis<double>(0.5 * M_PI, Eigen::Vector3d::UnitY()));
+
+  Eigen::Quaternion<double> q0(
+      Eigen::AngleAxis<double>(0, Eigen::Vector3d::UnitY()));
+
+  auto qux = q * ux;
+
+  VLOG(1) << "Q is w: " << q.w() << " vec: " << q.vec();
+  VLOG(1) << "ux is " << ux;
+  VLOG(1) << "qux is " << qux;
+
+  // Start by rotating around the X world vector for pi/2
+  Eigen::Quaternion<double> integral1(
+      RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative, ux,
+                           std::placeholders::_1),
+                 q0.coeffs(), 0.5 * M_PI));
+
+  VLOG(1) << "integral1 * uz => " << integral1 * uz;
+
+  // Then rotate around the Y world vector for pi/2
+  Eigen::Quaternion<double> integral2(
+      RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative, uy,
+                           std::placeholders::_1),
+                 integral1.normalized().coeffs(), 0.5 * M_PI));
+
+  VLOG(1) << "integral2 * uz => " << integral2 * uz;
+
+  // Then rotate around the X world vector for -pi/2
+  Eigen::Quaternion<double> integral3(
+      RungeKutta(std::bind(&drivetrain::DrivetrainUkf::QuaternionDerivative,
+                           -ux, std::placeholders::_1),
+                 integral2.normalized().coeffs(), 0.5 * M_PI));
+
+  integral3.normalize();
+
+  VLOG(1) << "Integral is w: " << integral1.w() << " vec: " << integral1.vec()
+          << " norm " << integral1.norm();
+
+  VLOG(1) << "Integral is w: " << integral3.w() << " vec: " << integral3.vec()
+          << " norm " << integral3.norm();
+
+  VLOG(1) << "ux => " << integral3 * ux;
+  EXPECT_NEAR(0.0, (uy - integral3 * ux).norm(), 5e-2);
+}
+
+TEST(RungeKuttaTest, Ukf) {
+  drivetrain::DrivetrainUkf dtukf;
+  Eigen::Vector3d ux = Eigen::Vector3d::UnitX();
+  Eigen::Matrix<double, 3, 1> measurement;
+  measurement.setZero();
+  dtukf.Predict(ux, measurement);
+}
+
+// Tests that small perturbations around a couple quaternions averaged out
+// return the original quaternion.
+TEST(RungeKuttaTest, QuaternionMean) {
+  Eigen::Matrix<double, 4, 7> vectors;
+  vectors.col(0) << 0, 0, 0, 1;
+  for (int i = 0; i < 3; ++i) {
+    Eigen::Matrix<double, 4, 1> perturbation;
+    perturbation.setZero();
+    perturbation(i, 0) = 0.1;
+
+    vectors.col(i * 2 + 1) = vectors.col(0) + perturbation;
+    vectors.col(i * 2 + 2) = vectors.col(0) - perturbation;
+  }
+
+  for (int i = 0; i < 7; ++i) {
+    vectors.col(i).normalize();
+  }
+
+  Eigen::Matrix<double, 4, 1> mean = drivetrain::QuaternionMean(vectors);
+
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_NEAR(mean(i, 0), vectors(i, 0), 0.001) << ": Failed on index " << i;
+  }
+}
+
+// Tests that computing sigma points, and then computing the mean and covariance
+// returns the original answer.
+TEST(RungeKuttaTest, SigmaPoints) {
+  const Eigen::Quaternion<double> mean(
+      Eigen::AngleAxis<double>(M_PI / 2.0, Eigen::Vector3d::UnitX()));
+
+  Eigen::Matrix<double, 3, 3> covariance;
+  covariance << 0.4, -0.1, 0.2, -0.1, 0.6, 0.0, 0.2, 0.0, 0.5;
+  covariance *= 0.1;
+
+  const Eigen::Matrix<double, 4, 3 * 2 + 1> vectors =
+      drivetrain::GenerateSigmaPoints(mean, covariance);
+
+  const Eigen::Matrix<double, 4, 1> calculated_mean =
+      drivetrain::QuaternionMean(vectors);
+
+  VLOG(1) << "actual mean: " << mean.coeffs();
+  VLOG(1) << "calculated mean: " << calculated_mean;
+
+  Eigen::Matrix<double, 3, 3 * 2 + 1> Wprime;
+  Eigen::Matrix<double, 3, 3> calculated_covariance =
+      drivetrain::ComputeQuaternionCovariance(
+          Eigen::Quaternion<double>(calculated_mean), vectors, &Wprime);
+
+  EXPECT_NEAR(1.0,
+              (mean.conjugate().coeffs() * calculated_mean.transpose()).norm(),
+              1e-4);
+
+  EXPECT_NEAR(0.0, (calculated_covariance - covariance).norm(), 1e-8);
+}
+
+// Tests that computing sigma points with a large covariance that will precisely
+// wrap, that we do clip the perturbations.
+TEST(RungeKuttaTest, ClippedSigmaPoints) {
+  const Eigen::Quaternion<double> mean(
+      Eigen::AngleAxis<double>(M_PI / 2.0, Eigen::Vector3d::UnitX()));
+
+  Eigen::Matrix<double, 3, 3> covariance;
+  covariance << 0.4, -0.1, 0.2, -0.1, 0.6, 0.0, 0.2, 0.0, 0.5;
+  covariance *= 100.0;
+
+  const Eigen::Matrix<double, 4, 3 * 2 + 1> vectors =
+      drivetrain::GenerateSigmaPoints(mean, covariance);
+
+  const Eigen::Matrix<double, 4, 1> calculated_mean =
+      drivetrain::QuaternionMean(vectors);
+
+  Eigen::Matrix<double, 3, 3 * 2 + 1> Wprime;
+  Eigen::Matrix<double, 3, 3> calculated_covariance =
+      drivetrain::ComputeQuaternionCovariance(
+          Eigen::Quaternion<double>(calculated_mean), vectors, &Wprime);
+
+  EXPECT_NEAR(1.0,
+              (mean.conjugate().coeffs() * calculated_mean.transpose()).norm(),
+              1e-4);
+
+  const double calculated_covariance_norm = calculated_covariance.norm();
+  const double covariance_norm = covariance.norm();
+  EXPECT_LT(calculated_covariance_norm, covariance_norm / 2.0)
+      << "Calculated covariance should be much smaller than the original "
+         "covariance.";
+}
+
+// Tests that ToRotationVectorFromQuaternion works for a 0 rotation.
+TEST(RungeKuttaTest, ToRotationVectorFromQuaternionAtZero) {
+  Eigen::Matrix<double, 3, 1> vector =
+      drivetrain::ToRotationVectorFromQuaternion(
+          Eigen::Quaternion<double>(
+              Eigen::AngleAxis<double>(0.0, Eigen::Vector3d::UnitX()))
+              .coeffs());
+
+  EXPECT_NEAR(0.0, (vector - Eigen::Vector3d::Zero()).norm(), 1e-4);
+}
+
+// Tests that ToRotationVectorFromQuaternion works for a real rotation.
+TEST(RungeKuttaTest, ToRotationVectorFromQuaternion) {
+  Eigen::Matrix<double, 3, 1> vector =
+      drivetrain::ToRotationVectorFromQuaternion(
+          Eigen::Quaternion<double>(
+              Eigen::AngleAxis<double>(M_PI * 0.5, Eigen::Vector3d::UnitX()))
+              .coeffs());
+
+  EXPECT_NEAR(0.0, (vector - Eigen::Vector3d::UnitX() * M_PI * 0.5).norm(),
+              1e-4);
+}
+
+// Tests that ToRotationVectorFromQuaternion works for a solution with negative
+// coefficients.
+TEST(RungeKuttaTest, ToRotationVectorFromQuaternionNegative) {
+  Eigen::Matrix<double, 3, 1> vector =
+      drivetrain::ToRotationVectorFromQuaternion(
+          Eigen::Quaternion<double>(
+              -Eigen::Quaternion<double>(
+                   Eigen::AngleAxis<double>(M_PI * 0.5,
+                                              Eigen::Vector3d::UnitX()))
+                   .coeffs())
+              .coeffs());
+
+  EXPECT_NEAR(0.0, (vector - Eigen::Vector3d::UnitX() * M_PI * 0.5).norm(),
+              1e-4);
+}
+
+// Tests that ToQuaternionFromRotationVector works for a 0 rotation.
+TEST(RungeKuttaTest, ToQuaternionFromRotationVectorAtZero) {
+  Eigen::Matrix<double, 4, 1> quaternion =
+      drivetrain::ToQuaternionFromRotationVector(Eigen::Vector3d::Zero());
+
+  EXPECT_NEAR(0.0, (quaternion -
+                    Eigen::Quaternion<double>(
+                        Eigen::AngleAxis<double>(0.0, Eigen::Vector3d::UnitX()))
+                        .coeffs()).norm(),
+              1e-4);
+}
+
+// Tests that ToQuaternionFromRotationVector works for a real rotation.
+TEST(RungeKuttaTest, ToQuaternionFromRotationVector) {
+  Eigen::Matrix<double, 4, 1> quaternion =
+      drivetrain::ToQuaternionFromRotationVector(Eigen::Vector3d::UnitX() *
+                                                 M_PI * 0.5);
+
+  EXPECT_NEAR(0.0, (quaternion -
+                    Eigen::Quaternion<double>(
+                        Eigen::AngleAxis<double>(
+                            M_PI * 0.5, Eigen::Vector3d::UnitX())).coeffs())
+
+                       .norm(),
+              1e-4);
+}
+
+// Tests that ToQuaternionFromRotationVector correctly clips a rotation vector
+// that is too large in magnitude.
+TEST(RungeKuttaTest, ToQuaternionFromLargeRotationVector) {
+  const double kMaxAngle = 2.0;
+  const Eigen::Vector3d rotation_vector =
+      Eigen::Vector3d::UnitX() * kMaxAngle * 2.0;
+  const Eigen::Matrix<double, 3, 1> clipped_vector =
+      drivetrain::ToRotationVectorFromQuaternion(
+          drivetrain::ToQuaternionFromRotationVector(rotation_vector,
+                                                     kMaxAngle));
+
+  EXPECT_NEAR(0.0, (rotation_vector / 2.0 - clipped_vector).norm(), 1e-4);
+}
+
+// Tests that ToQuaternionFromRotationVector and ToRotationVectorFromQuaternion
+// works for random rotations.
+TEST(RungeKuttaTest, RandomQuaternions) {
+  std::mt19937 generator(aos::testing::RandomSeed());
+  std::uniform_real_distribution<double> random_scalar(-1.0, 1.0);
+
+  for (int i = 0; i < 1000; ++i) {
+    Eigen::Matrix<double, 3, 1> axis;
+    axis << random_scalar(generator), random_scalar(generator),
+        random_scalar(generator);
+    EXPECT_GE(axis.norm(), 1e-6);
+    axis.normalize();
+
+    const double angle = random_scalar(generator) * M_PI;
+
+    Eigen::Matrix<double, 4, 1> quaternion =
+        drivetrain::ToQuaternionFromRotationVector(axis * angle);
+
+    Eigen::Quaternion<double> answer(Eigen::AngleAxis<double>(angle, axis));
+
+    EXPECT_NEAR(quaternion(3, 0), std::cos(angle / 2.0), 1e-8);
+    EXPECT_NEAR(answer.w(), std::cos(angle / 2.0), 1e-8);
+
+    EXPECT_NEAR(1.0, (answer.coeffs() * quaternion.transpose()).norm(), 1e-6);
+
+    const Eigen::Matrix<double, 3, 1> recalculated_axis =
+        drivetrain::ToRotationVectorFromQuaternion(quaternion);
+
+    EXPECT_NEAR(std::abs(angle), recalculated_axis.norm(), 1e-8);
+
+    EXPECT_NEAR(0.0, (axis * angle - recalculated_axis).norm(), 1e-8);
+  }
+}
+
+}  // namespace testing
+}  // namespace control_loops
+}  // namespace frc971