Added trajectory planner for the arm.

This computes the velocity as a function of path distance.
I'm not quite certain how to verify this other than the fact that I've
plotted it and the python code, and they produce the same answers.

Change-Id: Ibdbdfc78eb0c142c7e6a30047afa099b73964811
diff --git a/y2018/control_loops/superstructure/arm/BUILD b/y2018/control_loops/superstructure/arm/BUILD
index cbb443d..dfdaef4 100644
--- a/y2018/control_loops/superstructure/arm/BUILD
+++ b/y2018/control_loops/superstructure/arm/BUILD
@@ -6,7 +6,10 @@
     hdrs = [
         "trajectory.h",
     ],
-    deps = ["//third_party/eigen"],
+    deps = [
+        ":dynamics",
+        "//third_party/eigen",
+    ],
 )
 
 cc_test(
diff --git a/y2018/control_loops/superstructure/arm/dynamics.cc b/y2018/control_loops/superstructure/arm/dynamics.cc
index 5086e78..87d2178 100644
--- a/y2018/control_loops/superstructure/arm/dynamics.cc
+++ b/y2018/control_loops/superstructure/arm/dynamics.cc
@@ -12,6 +12,8 @@
                    Dynamics::kResistance)
         .finished();
 
+const ::Eigen::Matrix<double, 2, 2> Dynamics::K3_inverse = K3.inverse();
+
 const ::Eigen::Matrix<double, 2, 2> Dynamics::K4 =
     (::Eigen::Matrix<double, 2, 2>()
          << Dynamics::kG1 * Dynamics::kG1 * Dynamics::Kt /
diff --git a/y2018/control_loops/superstructure/arm/dynamics.h b/y2018/control_loops/superstructure/arm/dynamics.h
index 7ad6696..832796a 100644
--- a/y2018/control_loops/superstructure/arm/dynamics.h
+++ b/y2018/control_loops/superstructure/arm/dynamics.h
@@ -52,6 +52,7 @@
 
   // K3, K4 matricies described below.
   static const ::Eigen::Matrix<double, 2, 2> K3;
+  static const ::Eigen::Matrix<double, 2, 2> K3_inverse;
   static const ::Eigen::Matrix<double, 2, 2> K4;
 
   // Generates K1-2 for the arm ODE.
@@ -114,7 +115,7 @@
 
     MatriciesForState(X, &K1, &K2);
 
-    return K3.inverse() * (K1 * alpha_t + K2 * omega_t + K4 * omega_t);
+    return K3_inverse * (K1 * alpha_t + K2 * omega_t + K4 * omega_t);
   }
 
   static const ::Eigen::Matrix<double, 4, 1> UnboundedDiscreteDynamics(
diff --git a/y2018/control_loops/superstructure/arm/trajectory.cc b/y2018/control_loops/superstructure/arm/trajectory.cc
index e5fdfe5..fb1e0a5 100644
--- a/y2018/control_loops/superstructure/arm/trajectory.cc
+++ b/y2018/control_loops/superstructure/arm/trajectory.cc
@@ -4,6 +4,33 @@
 namespace control_loops {
 namespace superstructure {
 namespace arm {
+
+Path::Path(::std::vector<::std::array<double, 6>> list) {
+  distances_.reserve(list.size());
+  ::Eigen::Matrix<double, 2, 1> last_theta;
+  bool first = true;
+
+  // Pull apart the initializer list into points, and compute the distance
+  // lookup for efficient path lookups.
+  for (const auto &e : list) {
+    thetas_.push_back(
+        (::Eigen::Matrix<double, 2, 1>() << e[0], e[1]).finished());
+    omegas_.push_back(
+        (::Eigen::Matrix<double, 2, 1>() << e[2], e[3]).finished());
+    alphas_.push_back(
+        (::Eigen::Matrix<double, 2, 1>() << e[4], e[5]).finished());
+
+    if (first) {
+      distances_.push_back(0.0);
+    } else {
+      distances_.push_back((thetas_.back() - last_theta).norm() +
+                           distances_.back());
+    }
+    first = false;
+    last_theta = thetas_.back();
+  }
+}
+
 Path::Path(::std::initializer_list<::std::array<double, 6>> list) {
   distances_.reserve(list.size());
   ::Eigen::Matrix<double, 2, 1> last_theta;
@@ -29,6 +56,253 @@
     last_theta = thetas_.back();
   }
 }
+
+::std::vector<double> Trajectory::CurvatureOptimizationPass(
+    const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer, double vmax) {
+  ::std::vector<double> max_dvelocity_unfiltered;
+  max_dvelocity_unfiltered.reserve(num_plan_points_);
+  for (size_t i = 0; i < num_plan_points_; ++i) {
+    const double distance = DistanceForIndex(i);
+
+    // TODO(austin): We are looking up the index in the path 3 times here.
+    const ::Eigen::Matrix<double, 2, 1> theta = path_->Theta(distance);
+    const ::Eigen::Matrix<double, 2, 1> omega = path_->Omega(distance);
+    const ::Eigen::Matrix<double, 2, 1> alpha = path_->Alpha(distance);
+    const ::Eigen::Matrix<double, 4, 1> X =
+        (::Eigen::Matrix<double, 4, 1>() << theta(0, 0), 0.0, theta(1, 0), 0.0)
+            .finished();
+    ::Eigen::Matrix<double, 2, 2> K1;
+    ::Eigen::Matrix<double, 2, 2> K2;
+
+    Dynamics::NormilizedMatriciesForState(X, &K1, &K2);
+
+    const ::Eigen::Matrix<double, 2, 2> omega_square =
+        (::Eigen::Matrix<double, 2, 2>() << omega(0, 0), 0.0, 0.0, omega(1, 0))
+            .finished();
+
+    // Here, we can say that
+    //   d^2/dt^2 theta = d^2/dd^2 theta(d) * (d d/dt)^2
+    // Normalize so that the max accel is 1, and take magnitudes. This
+    // gives us the max velocity we can be at each point due to
+    // curvature.
+    const ::Eigen::Matrix<double, 2, 1> vk1 =
+        Dynamics::K3_inverse * (K1 * alpha + K2 * omega_square * omega);
+    const ::Eigen::Matrix<double, 2, 1> vk2 =
+        Dynamics::K3_inverse * Dynamics::K4 * omega;
+
+    double min_good_ddot =
+        ::std::sqrt(1.0 / ::std::max(0.001, (alpha_unitizer * alpha).norm()));
+
+    // Loop through all the various vmin, vmax combinations.
+    for (const double c : {-vmax, vmax}) {
+      // Also loop through saturating theta0 and theta1
+      for (const ::std::pair<double, double> ab :
+           {::std::pair<double, double>{vk1(0, 0), vk2(0, 0)},
+            ::std::pair<double, double>{vk1(1, 0), vk2(1, 0)}}) {
+        const double a = ab.first;
+        const double b = ab.second;
+        const double sqrt_number = b * b - 4.0 * a * c;
+        // Throw out imaginary solutions to the quadratic.
+        if (sqrt_number > 0) {
+          const double sqrt_result = ::std::sqrt(sqrt_number);
+          const double ddot1 = (-b + sqrt_result) / (2.0 * a);
+          const double ddot2 = (-b - sqrt_result) / (2.0 * a);
+          // Loop through both solutions.
+          for (const double ddot : {ddot1, ddot2}) {
+            const ::Eigen::Matrix<double, 2, 1> U =
+                vk1 * ddot * ddot + vk2 * ddot;
+            // Finally, make sure the velocity is positive and valid.
+            if ((U.array().abs() <= vmax + 1e-6).all() && ddot > 0.0) {
+              min_good_ddot = ::std::min(min_good_ddot, ddot);
+            }
+          }
+        }
+      }
+    }
+
+    max_dvelocity_unfiltered.push_back(min_good_ddot);
+  }
+  return max_dvelocity_unfiltered;
+}
+
+double Trajectory::FeasableForwardsAcceleration(
+    double goal_distance, double goal_velocity, double vmax,
+    const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer) {
+  // TODO(austin): We are looking up the index in the path 3 times here.
+  const ::Eigen::Matrix<double, 2, 1> theta = path_->Theta(goal_distance);
+  const ::Eigen::Matrix<double, 2, 1> omega = path_->Omega(goal_distance);
+  const ::Eigen::Matrix<double, 2, 1> alpha = path_->Alpha(goal_distance);
+
+  const ::Eigen::Matrix<double, 4, 1> X =
+      (::Eigen::Matrix<double, 4, 1>() << theta(0, 0), 0.0, theta(1, 0), 0.0)
+          .finished();
+
+  ::Eigen::Matrix<double, 2, 2> K1;
+  ::Eigen::Matrix<double, 2, 2> K2;
+
+  Dynamics::NormilizedMatriciesForState(X, &K1, &K2);
+
+  const ::Eigen::Matrix<double, 2, 2> omega_square =
+      (::Eigen::Matrix<double, 2, 2>() << omega(0, 0), 0.0, 0.0, omega(1, 0))
+          .finished();
+
+  const ::Eigen::Matrix<double, 2, 1> k_constant =
+      Dynamics::K3_inverse * ((K1 * alpha + K2 * omega_square * omega) *
+                                  goal_velocity * goal_velocity +
+                              Dynamics::K4 * omega * goal_velocity);
+  const ::Eigen::Matrix<double, 2, 1> k_scalar =
+      Dynamics::K3_inverse * K1 * omega;
+
+  double goal_acceleration =
+      ::std::sqrt(
+          ::std::max(0.0, 1.0 - ::std::pow((alpha_unitizer * alpha).norm() *
+                                               goal_velocity * goal_velocity,
+                                           2))) /
+      (alpha_unitizer * omega).norm();
+
+  for (double c : {-vmax, vmax}) {
+    for (const ::std::pair<double, double> ab :
+         {::std::pair<double, double>{k_constant(0, 0), k_scalar(0, 0)},
+          ::std::pair<double, double>{k_constant(1, 0), k_scalar(1, 0)}}) {
+      const double a = ab.first;
+      const double b = ab.second;
+      const double voltage_accel = (c - a) / b;
+      const ::Eigen::Matrix<double, 2, 1> U =
+          k_constant + k_scalar * voltage_accel;
+
+      if (voltage_accel > 0.0 && (U.array().abs() <= vmax + 1e-6).all())
+        goal_acceleration = ::std::min(voltage_accel, goal_acceleration);
+    }
+  }
+
+  return goal_acceleration;
+}
+
+double Trajectory::FeasableBackwardsAcceleration(
+    double goal_distance, double goal_velocity, double vmax,
+    const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer) {
+    const ::Eigen::Matrix<double, 2, 1> theta = path_->Theta(goal_distance);
+    const ::Eigen::Matrix<double, 2, 1> omega = path_->Omega(goal_distance);
+    const ::Eigen::Matrix<double, 2, 1> alpha = path_->Alpha(goal_distance);
+    const ::Eigen::Matrix<double, 4, 1> X =
+        (::Eigen::Matrix<double, 4, 1>() << theta(0, 0), 0.0, theta(1, 0), 0.0)
+            .finished();
+    ::Eigen::Matrix<double, 2, 2> K1;
+    ::Eigen::Matrix<double, 2, 2> K2;
+
+    Dynamics::NormilizedMatriciesForState(X, &K1, &K2);
+
+    const ::Eigen::Matrix<double, 2, 2> omega_square =
+        (::Eigen::Matrix<double, 2, 2>() << omega(0, 0), 0.0, 0.0, omega(1, 0))
+            .finished();
+
+    const ::Eigen::Matrix<double, 2, 1> k_constant =
+        Dynamics::K3_inverse *
+        ((K1 * alpha + K2 * omega_square * omega) * goal_velocity * goal_velocity +
+         Dynamics::K4 * omega * goal_velocity);
+    const ::Eigen::Matrix<double, 2, 1> k_scalar =
+        Dynamics::K3_inverse * K1 * omega;
+
+    double goal_acceleration =
+        ::std::sqrt(
+            ::std::max(0.0, 1.0 - ::std::pow(((alpha_unitizer * alpha).norm() *
+                                              goal_velocity * goal_velocity),
+                                             2))) /
+        (alpha_unitizer * omega).norm();
+    for (double c : {-vmax, vmax}) {
+      for (const ::std::pair<double, double> ab :
+           {::std::pair<double, double>{k_constant(0, 0), k_scalar(0, 0)},
+            ::std::pair<double, double>{k_constant(1, 0), k_scalar(1, 0)}}) {
+        const double a = ab.first;
+        const double b = ab.second;
+        // This time, we are doing the other pass.  So, find all the
+        // decelerations (and flip them) to find the prior velocity.
+        const double voltage_accel = (c - a) / b;
+
+        const ::Eigen::Matrix<double, 2, 1> U =
+            k_constant + k_scalar * voltage_accel;
+
+        if (voltage_accel < 0.0 && (U.array().abs() <= vmax + 1e-6).all()) {
+          goal_acceleration = ::std::min(-voltage_accel, goal_acceleration);
+        }
+      }
+    }
+
+    return goal_acceleration;
+}
+
+::std::vector<double> Trajectory::BackwardsOptimizationPass(
+    const ::std::vector<double> &max_dvelocity,
+    const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer, double vmax) {
+  ::std::vector<double> max_dvelocity_back_pass = max_dvelocity;
+
+  // Now, iterate over the list of velocities and constrain the acceleration.
+  for (int i = num_plan_points_ - 2; i >= 0; --i) {
+    const double previous_velocity = max_dvelocity_back_pass[i + 1];
+    const double prev_distance = DistanceForIndex(i + 1);
+
+    // Now, integrate with respect to distance (not time like normal).
+    // Subdivide the plan step size into kNumSteps steps.
+    double integrated_distance = 0.0;
+    double integrated_velocity = previous_velocity;
+    constexpr int kNumSteps = 10;
+    for (int j = 0; j < kNumSteps; ++j) {
+      // Compute the acceleration with respect to time.
+      const double acceleration = FeasableBackwardsAcceleration(
+          prev_distance - integrated_distance, integrated_velocity, vmax,
+          alpha_unitizer);
+
+      // And then integrate it with respect to distance.
+      const double integration_step_size_distance =
+          step_size_ / static_cast<double>(kNumSteps);
+      integrated_distance += integration_step_size_distance;
+      integrated_velocity =
+          ::std::sqrt(2.0 * acceleration * integration_step_size_distance +
+                      ::std::pow(integrated_velocity, 2));
+    }
+    max_dvelocity_back_pass[i] =
+        ::std::min(integrated_velocity, max_dvelocity_back_pass[i]);
+  }
+
+  return max_dvelocity_back_pass;
+}
+
+::std::vector<double> Trajectory::ForwardsOptimizationPass(
+    const ::std::vector<double> &max_dvelocity,
+    const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer, double vmax) {
+  ::std::vector<double> max_dvelocity_forward_pass = max_dvelocity;
+  // Now, iterate over the list of velocities and constrain the acceleration.
+  for (size_t i = 1; i < num_plan_points_; ++i) {
+    const double previous_velocity = max_dvelocity_forward_pass[i - 1];
+    const double previous_distance = DistanceForIndex(i - 1);
+
+    // Now, integrate with respect to distance (not time like normal).
+    // Subdivide the plan step size into kNumSteps steps.
+    double integrated_distance = 0.0;
+    double integrated_velocity = previous_velocity;
+    constexpr int kNumSteps = 10;
+    for (int j = 0; j < kNumSteps; ++j) {
+      // Compute the acceleration with respect to time.
+      const double acceleration = FeasableForwardsAcceleration(
+          previous_distance + integrated_distance, integrated_velocity, vmax,
+          alpha_unitizer);
+
+      // And then integrate it with respect to distance.
+      const double integration_step_size_distance =
+          step_size_ / static_cast<double>(kNumSteps);
+      integrated_distance += integration_step_size_distance;
+      integrated_velocity =
+          ::std::sqrt(2.0 * acceleration * integration_step_size_distance +
+                      ::std::pow(integrated_velocity, 2));
+    }
+
+    max_dvelocity_forward_pass[i] =
+        ::std::min(integrated_velocity, max_dvelocity_forward_pass[i]);
+  }
+
+  return max_dvelocity_forward_pass;
+}
+
 }  // namespace arm
 }  // namespace superstructure
 }  // namespace control_loops
diff --git a/y2018/control_loops/superstructure/arm/trajectory.h b/y2018/control_loops/superstructure/arm/trajectory.h
index a9136b4..2c9d68c 100644
--- a/y2018/control_loops/superstructure/arm/trajectory.h
+++ b/y2018/control_loops/superstructure/arm/trajectory.h
@@ -3,6 +3,8 @@
 
 #include "Eigen/Dense"
 
+#include "y2018/control_loops/superstructure/arm/dynamics.h"
+
 namespace y2018 {
 namespace control_loops {
 namespace superstructure {
@@ -16,6 +18,9 @@
   // Constructs a path from an initializer list of (theta0, theta1, omega0,
   // omega1, alpha0, alpha1)
   Path(::std::initializer_list<::std::array<double, 6>> list);
+  // TODO(austin): I'd like to not have 2 constructors, but it looks like C++
+  // doesn't want that.
+  Path(::std::vector<::std::array<double, 6>> list);
 
   // Returns the length of the path in radians.
   double length() const { return distances_.back(); }
@@ -33,6 +38,8 @@
     return Interpolate(alphas_, distance);
   }
 
+  const ::std::vector<double> &distances() const { return distances_; }
+
  private:
   // Interpolates the function represented by the list of points at the provided
   // distance.  Distances outside the range will get truncated.  This is a
@@ -48,10 +55,10 @@
     }
     // Keep a max and min, and pull them in towards eachother until they bound
     // our point exactly.
-    size_t before = ::std::distance(
+    size_t after = ::std::distance(
         distances_.begin(),
         ::std::lower_bound(distances_.begin(), distances_.end(), distance));
-    size_t after = before + 1;
+    size_t before = after - 1;
     return (distance - distances_[before]) * (points[after] - points[before]) /
                (distances_[after] - distances_[before]) +
            points[before];
@@ -68,6 +75,133 @@
   ::std::vector<double> distances_;
 };
 
+namespace testing {
+class TrajectoryTest_IndicesForDistanceTest_Test;
+}  // namespace testing
+
+// A trajectory is a path and a set of velocities as a function of distance.
+class Trajectory {
+ public:
+  // Constructs a trajectory (but doesn't calculate it) given a path and a step
+  // size.
+  Trajectory(Path *path, double gridsize)
+      : path_(path),
+        num_plan_points_(
+            static_cast<size_t>(::std::ceil(path_->length() / gridsize) + 1)),
+        step_size_(path_->length() /
+                   static_cast<double>(num_plan_points_ - 1)) {}
+
+  // Optimizes the trajectory.  The path will adhere to the constraints that
+  // || angular acceleration * alpha_unitizer || < 1, and the applied voltage <
+  // vmax.
+  void OptimizeTrajectory(const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer,
+                          double vmax) {
+    max_dvelocity_unfiltered_ = CurvatureOptimizationPass(alpha_unitizer, vmax);
+
+    // We need to start and end the trajectory with 0 velocity.
+    max_dvelocity_unfiltered_[0] = 0.0;
+    max_dvelocity_unfiltered_[max_dvelocity_unfiltered_.size() - 1] = 0.0;
+
+    max_dvelocity_ = BackwardsOptimizationPass(max_dvelocity_unfiltered_,
+                                               alpha_unitizer, vmax);
+    max_dvelocity_forward_pass_ = ForwardsOptimizationPass(
+        max_dvelocity_, alpha_unitizer, vmax);
+  }
+
+  // Returns an array of the distances used in the plan.  The starting point
+  // (0.0), and end point (path->length()) are included in the array.
+  ::std::vector<double> DistanceArray() const {
+    ::std::vector<double> result;
+    result.reserve(num_plan_points_);
+    for (size_t i = 0; i < num_plan_points_; ++i) {
+      result.push_back(DistanceForIndex(i));
+    }
+    return result;
+  }
+
+  // Computes the maximum velocity that we can follow the path while adhering to
+  // the constraints that || angular acceleration * alpha_unitizer || < 1, and
+  // the applied voltage < vmax.  Returns the velocities.
+  ::std::vector<double> CurvatureOptimizationPass(
+      const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer, double vmax);
+
+  // Computes the maximum forwards feasable acceleration at a given position
+  // while adhering to the vmax and alpha_unitizer constraints.
+  // This gives us the maximum path distance acceleration (d^2d/dt^2) for any
+  // initial position and velocity for the forwards path.
+  double FeasableForwardsAcceleration(
+      double goal_distance, double goal_velocity, double vmax,
+      const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer);
+
+  // Computes the maximum backwards feasable acceleration at a given position
+  // while adhering to the vmax and alpha_unitizer constraints.
+  // This gives us the maximum path distance acceleration (d^2d/dt^2) for any
+  // initial position and velocity for the backwards path.
+  // Note: positive acceleration means speed up while going in the negative
+  // direction on the path.
+  double FeasableBackwardsAcceleration(
+      double goal_distance, double goal_velocity, double vmax,
+      const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer);
+
+  // Executes the backwards path optimization pass.
+  ::std::vector<double> BackwardsOptimizationPass(
+      const ::std::vector<double> &max_dvelocity,
+      const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer, double vmax);
+
+  // Executes the forwards path optimization pass.
+  ::std::vector<double> ForwardsOptimizationPass(
+      const ::std::vector<double> &max_dvelocity,
+      const ::Eigen::Matrix<double, 2, 2> &alpha_unitizer, double vmax);
+
+  // Returns the number of points used for the plan.
+  size_t num_plan_points() const { return num_plan_points_; }
+
+  // Returns the curvature only velocity plan.
+  const ::std::vector<double> &max_dvelocity_unfiltered() const {
+    return max_dvelocity_unfiltered_;
+  }
+  // Returns the backwards pass + curvature plan.
+  const ::std::vector<double> &max_dvelocity() const { return max_dvelocity_; }
+  // Returns the full plan.  This isn't useful at runtime since we don't want to
+  // be tied to a specific acceleration plan when we hit saturation.
+  const ::std::vector<double> &max_dvelocity_forward_pass() const {
+    return max_dvelocity_forward_pass_;
+  }
+
+ private:
+  friend class testing::TrajectoryTest_IndicesForDistanceTest_Test;
+
+  // Returns the distance along the path for a specific index in the plan.
+  double DistanceForIndex(size_t index) const {
+    return static_cast<double>(index) * step_size_;
+  }
+
+  // Returns the before and after indices for a specific distance in the plan.
+  ::std::pair<size_t, size_t> IndicesForDistance(double distance) const {
+    const double path_length = path_->length();
+    if (distance <= 0.0) {
+      return ::std::pair<size_t, size_t>(0, 1);
+    }
+    const size_t lower_index =
+        ::std::min(static_cast<size_t>(num_plan_points_ - 2),
+                   static_cast<size_t>(::std::floor((num_plan_points_ - 1) *
+                                                    distance / path_length)));
+
+    return ::std::pair<size_t, size_t>(lower_index, lower_index + 1);
+  }
+
+  // The path to follow.
+  const Path *path_ = nullptr;
+  // The number of points in the plan.
+  const size_t num_plan_points_;
+  // A cached version of the step size since we need this a *lot*.
+  const double step_size_;
+
+  ::std::vector<double> max_dvelocity_unfiltered_;
+  ::std::vector<double> max_dvelocity_;
+  ::std::vector<double> max_dvelocity_forward_pass_;
+};
+
 }  // namespace arm
 }  // namespace superstructure
 }  // namespace control_loops
diff --git a/y2018/control_loops/superstructure/arm/trajectory_test.cc b/y2018/control_loops/superstructure/arm/trajectory_test.cc
index 5d5350b..a4e7822 100644
--- a/y2018/control_loops/superstructure/arm/trajectory_test.cc
+++ b/y2018/control_loops/superstructure/arm/trajectory_test.cc
@@ -32,6 +32,19 @@
   }
 }
 
+// Tests a path where picking the right point to interpolate for matters.
+TEST(TrajectoryTest, HarderPath) {
+  Path p({{{0.0, 0.0, 1.0, 0.0, 0.0, 0.0}},
+          {{1.0, 0.0, 1.0, 0.0, 0.0, 0.0}},
+          {{1.0, 1.0, 0.0, 1.0, 0.0, 0.0}}});
+
+  EXPECT_TRUE(p.Theta(0.5).isApprox(
+      (::Eigen::Matrix<double, 2, 1>() << 0.5, 0.0).finished()));
+
+  EXPECT_TRUE(p.Theta(1.5).isApprox(
+      (::Eigen::Matrix<double, 2, 1>() << 1.0, 0.5).finished()));
+}
+
 // Tests that a path with an even number of points works as expected.
 TEST(TrajectoryTest, EvenPointNumbers) {
   Path p({{{0.0, 0.0, 1.0, 0.0, 0.0, 0.0}},
@@ -59,6 +72,47 @@
       (::Eigen::Matrix<double, 2, 1>() << 3.0, 0.0).finished()));
 }
 
+
+// Tests that we can compute the indices of the plan for a given distance correctly.
+TEST(TrajectoryTest, IndicesForDistanceTest) {
+  // Start with a stupid simple plan.
+  Path p({{{0.0, 0.0, 1.0, 0.0, 0.0, 0.0}},
+          {{1.0, 0.0, 1.0, 0.0, 0.0, 0.0}},
+          {{2.0, 0.0, 1.0, 0.0, 0.0, 0.0}},
+          {{3.0, 0.0, 1.0, 0.0, 0.0, 0.0}}});
+  Trajectory t(&p, 0.1);
+
+  // 0 - 3.0 every 0.1 should be 31 points.
+  EXPECT_EQ(t.num_plan_points(), 31);
+
+  // Verify that something centered in a grid cell returns the points on either side.
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(0), static_cast<size_t>(1)),
+            t.IndicesForDistance(0.05));
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(1), static_cast<size_t>(2)),
+            t.IndicesForDistance(0.15));
+
+  // Verify that something on an edge returns the expected result.
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(1), static_cast<size_t>(2)),
+            t.IndicesForDistance(0.1));
+
+  // Verify what small deviations result.
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(1), static_cast<size_t>(2)),
+            t.IndicesForDistance(0.1001));
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(0), static_cast<size_t>(1)),
+            t.IndicesForDistance(0.0999));
+
+  // Verify that blowing past the ends works.
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(29), static_cast<size_t>(30)),
+            t.IndicesForDistance(4.0));
+  EXPECT_EQ(::std::make_pair(static_cast<size_t>(0), static_cast<size_t>(1)),
+            t.IndicesForDistance(-0.1));
+
+  // Verify that the index to distance calculation also works.
+  EXPECT_EQ(0.0, t.DistanceForIndex(0));
+  EXPECT_EQ(0.1, t.DistanceForIndex(1));
+  EXPECT_EQ(3.0, t.DistanceForIndex(30));
+}
+
 }  // namespace testing
 }  // namespace arm
 }  // namespace superstructure