Compute feed forwards voltage for drivetrain splines

And then integrate it all back up to make sure we get to our
destination.  The feed forwards voltage is now pretty precise.

Change-Id: I37b694cffdc0ae5a93946ad299a3b42404a00c96
diff --git a/frc971/control_loops/python/spline.py b/frc971/control_loops/python/spline.py
index f398d19..585d6b8 100644
--- a/frc971/control_loops/python/spline.py
+++ b/frc971/control_loops/python/spline.py
@@ -11,6 +11,7 @@
 import sys
 
 from frc971.control_loops.python import polydrivetrain
+from frc971.control_loops.python import drivetrain
 import y2018.control_loops.python.drivetrain
 
 """This file is my playground for implementing spline following.
@@ -22,6 +23,19 @@
 FLAGS = gflags.FLAGS
 
 
+def RungeKutta(f, y0, t, h, count=1):
+    """4th order RungeKutta integration of dy/dt = f(t, y) starting at X."""
+    y1 = y0
+    dh = h / float(count)
+    for x in xrange(count):
+        k1 = dh * f(t + dh * x, y1)
+        k2 = dh * f(t + dh * x + dh / 2.0, y1 + k1 / 2.0)
+        k3 = dh * f(t + dh * x + dh / 2.0, y1 + k2 / 2.0)
+        k4 = dh * f(t + dh * x + dh, y1 + k3)
+        y1 += (k1 + 2.0 * k2 + 2.0 * k3 + k4) / 6.0
+    return y1
+
+
 def spline(alpha, control_points):
     """Computes a Cubic Bezier curve.
 
@@ -342,6 +356,265 @@
                   numpy.array(dspline_points)[1, :]**2.0)**2.0))
 
 
+def integrate_accel_for_distance(f, v, x, dx):
+    # Use a trick from
+    # https://www.johndcook.com/blog/2012/02/21/care-and-treatment-of-singularities/
+    #
+    # We want to calculate:
+    #   v0 + (integral of dv/dt = f(x, v) from x to x + dx); noting that dv/dt
+    #   is expressed in t, not distance, so we want to do the integral of
+    #   dv/dx = f(x, v) / v.
+    #
+    # Because v can be near zero at the start of the integral (but because f is
+    # nonnegative, v will never go to zero), but the integral should still be
+    # valid, we follow the suggestion and instead calculate
+    # v0 + integral((f(x, v) - f(x0, v0)) / v) + integral(f(x0, v0) / v).
+    #
+    # Using a0 = f(x0, v0), we get the second term as
+    #   integral((f(x, v) - a0) / v)
+    # where when v is zero we will also be at x0/v0 (because v can only start
+    # at zero, not go to zero).
+    #
+    # The second term, integral(a0 / v) requires an approximation.--in
+    # this case, that dv/dt is constant. Thus, we have
+    #  integral(a0 / sqrt(v0^2 + 2*a0*x)) = sqrt(2*a0*dx + v0^2) - sqrt(v0^2)
+    #                                     = sqrt(2 * a0 * dx * v0^2) - v0.
+    #
+    # Because the RungeKutta function returns v0 + the integral, this
+    # gives the statements below.
+
+    a0 = f(x, v)
+
+    def integrablef(t, y):
+        # Since we know that a0 == a(0) and that they are asymptotically the
+        # same at 0, we know that the limit is 0 at 0.  This is true because
+        # when starting from a stop, under sane accelerations, we can assume
+        # that we will start with a constant acceleration.  So, hard-code it.
+        if numpy.abs(y) < 1e-6:
+            return 0.0
+        return (f(t, y) - a0) / y
+
+    return (RungeKutta(integrablef, v, x, dx) - v) + numpy.sqrt(2.0 * a0 * dx + v * v)
+
+
+class Trajectory(object):
+    def __init__(self, path, drivetrain, longitudal_accel, lateral_accel, distance_count):
+        self._path = path
+        self._drivetrain = drivetrain
+        self.distances = numpy.linspace(0.0,
+                                        self._path.length(), distance_count)
+        self._longitudal_accel = longitudal_accel
+        self._lateral_accel = lateral_accel
+
+        self._B_inverse = numpy.linalg.inv(self._drivetrain.B_continuous)
+
+    def create_plan(self, vmax):
+        vmax = 10.0
+        plan = numpy.array(numpy.zeros((len(self.distances), )))
+        plan.fill(vmax)
+        return plan
+
+    def lateral_velocity_curvature(self, distance):
+        return numpy.sqrt(self._lateral_accel /
+                          numpy.linalg.norm(self._path.ddxy(distance)))
+
+    def lateral_accel_pass(self, plan):
+        plan = plan.copy()
+        # TODO(austin): This appears to be doing nothing.
+        for i, distance in enumerate(self.distances):
+            plan[i] = min(plan[i], self.lateral_velocity_curvature(distance))
+        return plan
+
+    def compute_K345(self, current_dtheta, current_ddtheta):
+        # We've now got the equation:
+        #     K2 * d^2x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
+        K1 = numpy.matrix(
+            [[-self._drivetrain.robot_radius_l * current_ddtheta],
+             [self._drivetrain.robot_radius_r * current_ddtheta]])
+        K2 = numpy.matrix(
+            [[1.0 - self._drivetrain.robot_radius_l * current_dtheta],
+             [1.0 + self._drivetrain.robot_radius_r * current_dtheta]])
+
+        # Now, rephrase it as K5 a + K3 v^2 + K4 v = U
+        K3 = self._B_inverse * K1
+        K4 = -self._B_inverse * self._drivetrain.A_continuous * K2
+        K5 = self._B_inverse * K2
+        return K3, K4, K5
+
+    def curvature_voltage_pass(self, plan):
+        plan = plan.copy()
+        for i, distance in enumerate(self.distances):
+            current_ddtheta = self._path.ddtheta(distance)[0]
+            current_dtheta = self._path.dtheta(distance)[0]
+            # We've now got the equation:
+            #     K2 * d^2x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
+            # Now, rephrase it as K3 v^2 + K4 v = U
+            K3, K4, K5 = self.compute_K345(current_dtheta, current_ddtheta)
+            # But, we are going to assume that d^2x/dt^2 = 0
+
+            x = []
+            for a, b in [(K3[0, 0], K4[0, 0]), (K3[1, 0], K4[1, 0])]:
+                for c in [12.0, -12.0]:
+                    middle = b * b - 4.0 * a * c
+                    if middle >= 0.0:
+                        x.append((-b + numpy.sqrt(middle)) / (2.0 * a))
+                        x.append((-b - numpy.sqrt(middle)) / (2.0 * a))
+
+            maxx = 0.0
+            for newx in x:
+                if newx < 0.0:
+                    continue
+                U = (K3 * newx * newx + K4 * newx).T
+                # TODO(austin): We know that one of these *will* be +-12.0.  Only
+                # check the other one.
+                if not (numpy.abs(U) > 12.0 + 1e-6).any():
+                    maxx = max(newx, maxx)
+
+            if maxx == 0.0:
+                print('Could not solve')
+                return None
+            plan[i] = min(plan[i], maxx)
+            pass
+        return plan
+
+    def forward_acceleration(self, x, v):
+        current_ddtheta = self._path.ddtheta(x)[0]
+        current_dtheta = self._path.dtheta(x)[0]
+        # We've now got the equation:
+        #     K2 * d^2x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
+        # Now, rephrase it as K5 a + K3 v^2 + K4 v = U
+        K3, K4, K5 = self.compute_K345(current_dtheta, current_ddtheta)
+
+        C = K3 * v * v + K4 * v
+        # Note: K345 are not quite constant over the step, but we are going
+        # to assume they are for now.
+        accelerations = [(12.0 - C[0, 0]) / K5[0, 0], (12.0 - C[1, 0]) /
+                         K5[1, 0], (-12.0 - C[0, 0]) / K5[0, 0],
+                         (-12.0 - C[1, 0]) / K5[1, 0]]
+        maxa = 0.0
+        for a in accelerations:
+            if a < 0.0:
+                continue
+
+            U = K5 * a + K3 * v * v + K4 * v
+            if not (numpy.abs(U) > 12.0 + 1e-6).any():
+                maxa = max(maxa, a)
+
+        lateral_accel = v * v * numpy.linalg.norm(self._path.ddxy(x))
+        # Constrain the longitudinal acceleration to keep in a pseudo friction
+        # circle.  This will make it so we don't floor it while in a turn and
+        # cause extra wheel slip.
+        long_accel = numpy.sqrt(1.0 - (lateral_accel / self._lateral_accel)**
+                                2.0) * self._longitudal_accel
+        return min(long_accel, maxa)
+
+    def forward_pass(self, plan):
+        plan = plan.copy()
+        for i, distance in enumerate(self.distances):
+            if i == len(self.distances) - 1:
+                break
+
+            plan[i + 1] = min(
+                plan[i + 1],
+                integrate_accel_for_distance(
+                    self.forward_acceleration, plan[i], self.distances[i],
+                    self.distances[i + 1] - self.distances[i]))
+        return plan
+
+    def backward_acceleration(self, x, v):
+        # TODO(austin): Forwards and backwards are quite similar.  Can we
+        # factor this out?
+        current_ddtheta = self._path.ddtheta(x)[0]
+        current_dtheta = self._path.dtheta(x)[0]
+        # We've now got the equation:
+        #     K2 * d^2x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
+        # Now, rephrase it as K5 a + K3 v^2 + K4 v = U
+        K3, K4, K5 = self.compute_K345(current_dtheta, current_ddtheta)
+
+        C = K3 * v * v + K4 * v
+        # Note: K345 are not quite constant over the step, but we are going
+        # to assume they are for now.
+        accelerations = [(12.0 - C[0, 0]) / K5[0, 0], (12.0 - C[1, 0]) /
+                         K5[1, 0], (-12.0 - C[0, 0]) / K5[0, 0],
+                         (-12.0 - C[1, 0]) / K5[1, 0]]
+        mina = 0.0
+        for a in accelerations:
+            if a > 0.0:
+                continue
+
+            U = K5 * a + K3 * v * v + K4 * v
+            if not (numpy.abs(U) > 12.0 + 1e-6).any():
+                mina = min(mina, a)
+
+        lateral_accel = v * v * numpy.linalg.norm(self._path.ddxy(x))
+        # Constrain the longitudinal acceleration to keep in a pseudo friction
+        # circle.  This will make it so we don't floor it while in a turn and
+        # cause extra wheel slip.
+        long_accel = -numpy.sqrt(1.0 - (lateral_accel / self._lateral_accel)**
+                                 2.0) * self._longitudal_accel
+        return max(long_accel, mina)
+
+    def backward_pass(self, plan):
+        plan = plan.copy()
+        for i, distance in reversed(list(enumerate(self.distances))):
+            if i == 0:
+                break
+
+            plan[i - 1] = min(
+                plan[i - 1],
+                integrate_accel_for_distance(
+                    self.backward_acceleration, plan[i], self.distances[i],
+                    self.distances[i - 1] - self.distances[i]))
+        return plan
+
+    # TODO(austin): The plan should probably not be passed in...
+    def ff_voltage(self, plan, distance):
+        if distance < self.distances[1]:
+            after_index = 1
+            before_index = after_index - 1
+            if distance < self.distances[0]:
+                distance = 0.0
+        elif distance > self.distances[-2]:
+            after_index = len(self.distances) - 1
+            before_index = after_index - 1
+            if distance > self.distances[-1]:
+                distance = self.distances[-1]
+        else:
+            after_index = numpy.searchsorted(
+                self.distances, distance, side='right')
+            before_index = after_index - 1
+
+        vforwards = integrate_accel_for_distance(
+            self.forward_acceleration, plan[before_index],
+            self.distances[before_index],
+            distance - self.distances[before_index])
+        vbackward = integrate_accel_for_distance(
+            self.backward_acceleration, plan[after_index],
+            self.distances[after_index],
+            distance - self.distances[after_index])
+
+        vcurvature = self.lateral_velocity_curvature(distance)
+
+        if vcurvature < vforwards and vcurvature < vbackward:
+            accel = 0
+            velocity = vcurvature
+        elif vforwards < vbackward:
+            velocity = vforwards
+            accel = self.forward_acceleration(distance, velocity)
+        else:
+            velocity = vbackward
+            accel = self.backward_acceleration(distance, velocity)
+
+        current_ddtheta = self._path.ddtheta(distance)[0]
+        current_dtheta = self._path.dtheta(distance)[0]
+        # TODO(austin): Factor these out.
+        # We've now got the equation:
+        #     K2 * d^2x/dt^2 + K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
+        # Now, rephrase it as K5 a + K3 v^2 + K4 v = U
+        K3, K4, K5 = self.compute_K345(current_dtheta, current_ddtheta)
+
+        U = K5 * accel + K3 * velocity * velocity + K4 * velocity
+        return U
 
 def main(argv):
     # Build up the control point matrix
@@ -471,8 +744,8 @@
     ivelocity_plot = numpy.matrix(numpy.zeros((2, distance_count)))
     iposition_plot[:, 0] = position.copy()
     ivelocity_plot[:, 0] = velocity.copy()
-    itheta_plot = numpy.zeros((distance_count,))
-    iomega_plot = numpy.zeros((distance_count,))
+    itheta_plot = numpy.zeros((distance_count, ))
+    iomega_plot = numpy.zeros((distance_count, ))
     itheta_plot[0] = theta
     iomega_plot[0] = omega
 
@@ -518,69 +791,144 @@
 
     velocity_drivetrain = polydrivetrain.VelocityDrivetrainModel(
         y2018.control_loops.python.drivetrain.kDrivetrain)
+    position_drivetrain = drivetrain.Drivetrain(
+        y2018.control_loops.python.drivetrain.kDrivetrain)
+
+    longitudal_accel = 3.0
+    lateral_accel = 2.0
+
+    trajectory = Trajectory(
+        path,
+        drivetrain=velocity_drivetrain,
+        longitudal_accel=longitudal_accel,
+        lateral_accel=lateral_accel,
+        distance_count=500)
 
     vmax = numpy.inf
     vmax = 10.0
-    lateral_accel_plan = numpy.array(numpy.zeros((distance_count, )))
-    lateral_accel_plan.fill(vmax)
-    curvature_plan = lateral_accel_plan.copy()
+    lateral_accel_plan = trajectory.lateral_accel_pass(
+        trajectory.create_plan(vmax))
 
-    longitudal_accel = 15.0
-    lateral_accel = 20.0
+    forward_accel_plan = lateral_accel_plan.copy()
+    # Start and end the path stopped.
+    forward_accel_plan[0] = 0.0
+    forward_accel_plan[-1] = 0.0
 
-    for i, distance in enumerate(distances):
-        lateral_accel_plan[i] = min(
-            lateral_accel_plan[i],
-            numpy.sqrt(lateral_accel / numpy.linalg.norm(path.ddxy(distance))))
+    forward_accel_plan = trajectory.forward_pass(forward_accel_plan)
 
-        # We've now got the equation: K1 (dx/dt)^2 = A * K2 * dx/dt + B * U
-        # We need to find the feasible
-        current_ddtheta = path.ddtheta(distance)[0]
-        current_dtheta = path.dtheta(distance)[0]
-        K1 = numpy.matrix(
-            [[-velocity_drivetrain.robot_radius_l * current_ddtheta],
-             [velocity_drivetrain.robot_radius_r * current_ddtheta]])
-        K2 = numpy.matrix(
-            [[1.0 - velocity_drivetrain.robot_radius_l * current_dtheta],
-             [1.0 + velocity_drivetrain.robot_radius_r * current_dtheta]])
-        A = velocity_drivetrain.A_continuous
-        B = velocity_drivetrain.B_continuous
+    backward_accel_plan = trajectory.backward_pass(forward_accel_plan)
 
-        # Now, rephrase it as K5 a + K3 v^2 + K4 v = U
-        # Now, rephrase it as K5 a + K3 v^2 + K4 v = U
-        K3 = numpy.linalg.inv(B) * K1
-        K5 = numpy.linalg.inv(B) * K2
-        K4 = -K5
+    # And now, calculate the left, right voltage as a function of distance.
 
-        x = []
-        for a, b in [(K3[0, 0], K4[0, 0]), (K3[1, 0], K4[1, 0])]:
-            for c in [12.0, -12.0]:
-                middle = b * b - 4.0 * a * c
-                if middle >= 0.0:
-                    x.append((-b + numpy.sqrt(middle)) / (2.0 * a))
-                    x.append((-b - numpy.sqrt(middle)) / (2.0 * a))
-
-        maxx = 0.0
-        for newx in x:
-            if newx < 0.0:
-                continue
-            U = (K3 * newx * newx + K4 * newx).T
-            # TODO(austin): We know that one of these *will* be +-12.0.  Only
-            # check the other one.
-            if not (numpy.abs(U) > 12.0 + 1e-6).any():
-                maxx = max(newx, maxx)
-
-        if maxx == 0.0:
-            print('Could not solve')
-        curvature_plan[i] = min(lateral_accel_plan[i], maxx)
+    # TODO(austin): Factor out the accel and decel functions so we can use them
+    # to calculate voltage as a function of distance.
 
     pylab.figure()
-    pylab.plot(distances, lateral_accel_plan, label='accel pass')
-    pylab.plot(distances, curvature_plan, label='voltage limit pass')
+    pylab.plot(trajectory.distances, lateral_accel_plan, label='accel pass')
+    pylab.plot(trajectory.distances, forward_accel_plan, label='forward pass')
+    pylab.plot(trajectory.distances, backward_accel_plan, label='forward pass')
     pylab.xlabel("distance along spline (m)")
     pylab.ylabel("velocity (m/s)")
     pylab.legend()
 
+    voltages = numpy.matrix(numpy.zeros((2, len(distances))))
+    for i in range(len(distances)):
+        voltages[:, i] = trajectory.ff_voltage(backward_accel_plan,
+                                               distances[i])
+
+    # Now, let's integrate up the feed forwards voltage to see where we end up.
+    spline_state = numpy.matrix(numpy.zeros((4, 1)))
+    # x, y, theta, vl, vr
+    state = numpy.matrix(numpy.zeros((5, 1)))
+    dt = 0.005
+    num_timesteps = 400
+    spline_states = numpy.matrix(numpy.zeros((4, num_timesteps)))
+    states = numpy.matrix(numpy.zeros((5, num_timesteps)))
+    Us = numpy.matrix(numpy.zeros((2, num_timesteps)))
+    t = numpy.array(range(num_timesteps)) * dt
+    for i in range(num_timesteps):
+        spline_distance = 0.5 * (spline_state[0, 0] + spline_state[2, 0])
+        spline_velocity = 0.5 * (spline_state[1, 0] + spline_state[3, 0])
+        def distance_spline_diffeq(t, x):
+            spline_distance = 0.5 * (x[0, 0] + x[2, 0])
+            spline_velocity = 0.5 * (x[1, 0] + x[3, 0])
+            U = trajectory.ff_voltage(
+                backward_accel_plan,
+                spline_distance)
+            dXdt = (position_drivetrain.A_continuous * x +
+                      position_drivetrain.B_continuous * U)
+            return dXdt
+        spline_state = RungeKutta(distance_spline_diffeq, spline_state, i * dt, dt)
+
+        def spline_diffeq(t, x):
+            spline_distance = 0.5 * (spline_state[0, 0] + spline_state[2, 0])
+            spline_velocity = 0.5 * (spline_state[1, 0] + spline_state[3, 0])
+            velocity = x[3:, :]
+            theta = x[2, 0]
+            linear_velocity = (velocity[0, 0] + velocity[1, 0]) / 2.0
+            angular_velocity = (velocity[1, 0] - velocity[0, 0]) / (
+                velocity_drivetrain.robot_radius_l +
+                velocity_drivetrain.robot_radius_r)
+            U = trajectory.ff_voltage(
+                backward_accel_plan,
+                spline_distance + numpy.linalg.norm(x[0:2, 0] - state[0:2, 0]))
+            accel = (velocity_drivetrain.A_continuous * velocity +
+                     velocity_drivetrain.B_continuous * U)
+            return numpy.matrix(
+                [[numpy.cos(theta) * linear_velocity],
+                 [numpy.sin(theta) * linear_velocity], [angular_velocity],
+                 [accel[0, 0]], [accel[1, 0]]])
+
+        U = trajectory.ff_voltage(backward_accel_plan, spline_distance)
+        state = RungeKutta(spline_diffeq, state, i * dt, dt)
+        spline_states[:, i] = spline_state
+        states[:, i] = state
+        Us[:, i] = U
+
+    spline_distances = numpy.array(
+        (numpy.array(spline_states)[0, :] + numpy.array(spline_states)[2, :]) /
+        2.0)
+    pylab.figure()
+    pylab.plot(distances, numpy.array(voltages)[0, :], label='vl')
+    pylab.plot(distances, numpy.array(voltages)[1, :], label='vr')
+    pylab.legend()
+
+    pylab.figure()
+    pylab.plot(
+        numpy.array(spline_points)[0, :],
+        numpy.array(spline_points)[1, :],
+        label='spline')
+    pylab.plot(
+        numpy.array(states)[0, :], numpy.array(states)[1, :], label="robot")
+    pylab.legend()
+
+    pylab.figure()
+    pylab.plot(
+        spline_distances, (numpy.array(states)[0, :] - numpy.array(
+            path.xy(spline_distances))[0, :]) * 100.0,
+        label='robotx_error * 100')
+    pylab.plot(
+        spline_distances, (numpy.array(states)[1, :] - numpy.array(
+            path.xy(spline_distances))[1, :]) * 100.0,
+        label='roboty_error * 100')
+    pylab.plot(
+        spline_distances, (numpy.array(states)[2, :] - numpy.array(
+            path.theta(spline_distances))) * 100.0,
+        label='robottheta_error * 100')
+    pylab.plot(distances, numpy.array(voltages)[0, :], label='voltsl')
+    pylab.plot(distances, numpy.array(voltages)[1, :], label='voltsr')
+    pylab.legend()
+
+    def a(_, x):
+        return 2.0
+        return 2.0 + 0.0001 * x
+
+    v = 0.0
+    for _ in xrange(10):
+        dx = 4.0 / 10.0
+        v = integrate_accel_for_distance(a, v, 0.0, dx)
+    print('v', v)
+
     pylab.show()