frc971/control_loops/python/spline.py

#!/usr/bin/python

from __future__ import print_function

from matplotlib import pylab
import gflags
import glog
import numpy
import scipy
import scipy.integrate
import sys

from frc971.control_loops.python import polydrivetrain
import y2018.control_loops.python.drivetrain

"""This file is my playground for implementing spline following.

All splines here are cubic bezier splines.  See
  https://en.wikipedia.org/wiki/B%C3%A9zier_curve for more details.
"""

FLAGS = gflags.FLAGS


def spline(alpha, control_points):
    """Computes a Cubic Bezier curve.

    Args:
        alpha: scalar or list of spline parameters to calculate the curve at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        n x m matrix of spline points.  n is the dimension of the control
        points, and m is the number of points in 'alpha'.
    """
    if numpy.isscalar(alpha):
        alpha = [alpha]
    alpha_matrix = [[(1.0 - a)**3.0, 3.0 * (1.0 - a)**2.0 * a,
                     3.0 * (1.0 - a) * a**2.0, a**3.0] for a in alpha]

    return control_points * numpy.matrix(alpha_matrix).T


def dspline(alpha, control_points):
    """Computes the derivative of a Cubic Bezier curve wrt alpha.

    Args:
        alpha: scalar or list of spline parameters to calculate the curve at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        n x m matrix of spline point derivatives.  n is the dimension of the
        control points, and m is the number of points in 'alpha'.
    """
    if numpy.isscalar(alpha):
        alpha = [alpha]
    dalpha_matrix = [[
        -3.0 * (1.0 - a)**2.0, 3.0 * (1.0 - a)**2.0 + -2.0 * 3.0 *
        (1.0 - a) * a, -3.0 * a**2.0 + 2.0 * 3.0 * (1.0 - a) * a, 3.0 * a**2.0
    ] for a in alpha]

    return control_points * numpy.matrix(dalpha_matrix).T


def ddspline(alpha, control_points):
    """Computes the second derivative of a Cubic Bezier curve wrt alpha.

    Args:
        alpha: scalar or list of spline parameters to calculate the curve at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        n x m matrix of spline point second derivatives.  n is the dimension of
        the control points, and m is the number of points in 'alpha'.
    """
    if numpy.isscalar(alpha):
        alpha = [alpha]
    ddalpha_matrix = [[
        2.0 * 3.0 * (1.0 - a),
        -2.0 * 3.0 * (1.0 - a) + -2.0 * 3.0 * (1.0 - a) + 2.0 * 3.0 * a,
        -2.0 * 3.0 * a + 2.0 * 3.0 * (1.0 - a) - 2.0 * 3.0 * a,
        2.0 * 3.0 * a
    ] for a in alpha]

    return control_points * numpy.matrix(ddalpha_matrix).T


def dddspline(alpha, control_points):
    """Computes the third derivative of a Cubic Bezier curve wrt alpha.

    Args:
        alpha: scalar or list of spline parameters to calculate the curve at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        n x m matrix of spline point second derivatives.  n is the dimension of
        the control points, and m is the number of points in 'alpha'.
    """
    if numpy.isscalar(alpha):
        alpha = [alpha]
    ddalpha_matrix = [[
        -2.0 * 3.0,
        2.0 * 3.0 + 2.0 * 3.0 + 2.0 * 3.0,
        -2.0 * 3.0 - 2.0 * 3.0 - 2.0 * 3.0,
        2.0 * 3.0
    ] for a in alpha]

    return control_points * numpy.matrix(ddalpha_matrix).T


def spline_theta(alpha, control_points, dspline_points=None):
    """Computes the heading of a robot following a Cubic Bezier curve at alpha.

    Args:
        alpha: scalar or list of spline parameters to calculate the heading at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        m array of spline point headings.  m is the number of points in 'alpha'.
    """
    if dspline_points is None:
        dspline_points = dspline(alpha, control_points)

    return numpy.arctan2(
        numpy.array(dspline_points)[1, :], numpy.array(dspline_points)[0, :])


def dspline_theta(alpha,
                  control_points,
                  dspline_points=None,
                  ddspline_points=None):
    """Computes the derivative of the heading at alpha.

    This is the derivative of spline_theta wrt alpha.

    Args:
        alpha: scalar or list of spline parameters to calculate the derivative
            of the heading at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        m array of spline point heading derivatives.  m is the number of points
        in 'alpha'.
    """
    if dspline_points is None:
        dspline_points = dspline(alpha, control_points)

    if ddspline_points is None:
        ddspline_points = ddspline(alpha, control_points)

    dx = numpy.array(dspline_points)[0, :]
    dy = numpy.array(dspline_points)[1, :]

    ddx = numpy.array(ddspline_points)[0, :]
    ddy = numpy.array(ddspline_points)[1, :]

    return 1.0 / (dx**2.0 + dy**2.0) * (dx * ddy - dy * ddx)


def ddspline_theta(alpha,
                   control_points,
                   dspline_points=None,
                   ddspline_points=None,
                   dddspline_points=None):
    """Computes the second derivative of the heading at alpha.

    This is the second derivative of spline_theta wrt alpha.

    Args:
        alpha: scalar or list of spline parameters to calculate the second
            derivative of the heading at.
        control_points: n x 4 matrix of control points.  n[:, 0] is the
            starting point, and n[:, 3] is the ending point.

    Returns:
        m array of spline point heading second derivatives.  m is the number of
        points in 'alpha'.
    """
    if dspline_points is None:
        dspline_points = dspline(alpha, control_points)

    if ddspline_points is None:
        ddspline_points = ddspline(alpha, control_points)

    if dddspline_points is None:
        dddspline_points = dddspline(alpha, control_points)

    dddspline_points = dddspline(alpha, control_points)

    dx = numpy.array(dspline_points)[0, :]
    dy = numpy.array(dspline_points)[1, :]

    ddx = numpy.array(ddspline_points)[0, :]
    ddy = numpy.array(ddspline_points)[1, :]

    dddx = numpy.array(dddspline_points)[0, :]
    dddy = numpy.array(dddspline_points)[1, :]

    return -1.0 / ((dx**2.0 + dy**2.0)**2.0) * (dx * ddy - dy * ddx) * 2.0 * (
        dy * ddy + dx * ddx) + 1.0 / (dx**2.0 + dy**2.0) * (dx * dddy - dy *
                                                            dddx)


class Path(object):
    """Represents a path to follow."""
    def __init__(self, control_points):
        """Constructs a path given the control points."""
        self._control_points = control_points

        def spline_velocity(alpha):
            return numpy.linalg.norm(dspline(alpha, self._control_points), axis=0)

        self._point_distances = [0.0]
        num_alpha = 100
        # Integrate the xy velocity as a function of alpha for each step in the
        # table to get an alpha -> distance calculation.  Gaussian Quadrature
        # is quite accurate, so we can get away with fewer points here than we
        # might think.
        for alpha in numpy.linspace(0.0, 1.0, num_alpha)[:-1]:
            self._point_distances.append(
                scipy.integrate.fixed_quad(spline_velocity, alpha, alpha + 1.0
                                           / (num_alpha - 1.0))[0] +
                self._point_distances[-1])

    def distance_to_alpha(self, distance):
        """Converts distances along the spline to alphas.

        Args:
            distance: A scalar or array of distances to convert

        Returns:
            An array of distances, (1 big if the input was a scalar)
        """
        if numpy.isscalar(distance):
            return numpy.array([self._distance_to_alpha_scalar(distance)])
        else:
            return numpy.array([self._distance_to_alpha_scalar(d) for d in distance])

    def _distance_to_alpha_scalar(self, distance):
        """Helper to compute alpha for a distance for a single scalar."""
        if distance <= 0.0:
            return 0.0
        elif distance >= self.length():
            return 1.0
        after_index = numpy.searchsorted(
            self._point_distances, distance, side='right')
        before_index = after_index - 1

        # Linearly interpolate alpha from our (sorted) distance table.
        return (distance - self._point_distances[before_index]) / (
            self._point_distances[after_index] -
            self._point_distances[before_index]) * (1.0 / (
                len(self._point_distances) - 1.0)) + float(before_index) / (
                    len(self._point_distances) - 1.0)

    def length(self):
        """Returns the length of the spline (in meters)"""
        return self._point_distances[-1]

    # TODO(austin): need a better name...
    def xy(self, distance):
        """Returns the xy position as a function of distance."""
        return spline(self.distance_to_alpha(distance), self._control_points)

    # TODO(austin): need a better name...
    def dxy(self, distance):
        """Returns the xy velocity as a function of distance."""
        dspline_point = dspline(
            self.distance_to_alpha(distance), self._control_points)
        return dspline_point / numpy.linalg.norm(dspline_point, axis=0)

    # TODO(austin): need a better name...
    def ddxy(self, distance):
        """Returns the xy acceleration as a function of distance."""
        alpha = self.distance_to_alpha(distance)
        dspline_points = dspline(alpha, self._control_points)
        ddspline_points = ddspline(alpha, self._control_points)

        norm = numpy.linalg.norm(
            dspline_points, axis=0)**2.0

        return ddspline_points / norm - numpy.multiply(
            dspline_points, (numpy.array(dspline_points)[0, :] *
                             numpy.array(ddspline_points)[0, :] +
                             numpy.array(dspline_points)[1, :] *
                             numpy.array(ddspline_points)[1, :]) / (norm**2.0))

    def theta(self, distance, dspline_points=None):
        """Returns the heading as a function of distance."""
        return spline_theta(
            self.distance_to_alpha(distance),
            self._control_points,
            dspline_points=dspline_points)

    def dtheta(self, distance, dspline_points=None, ddspline_points=None):
        """Returns the angular velocity as a function of distance."""
        alpha = self.distance_to_alpha(distance)
        if dspline_points is None:
            dspline_points = dspline(alpha, self._control_points)
        if ddspline_points is None:
            ddspline_points = ddspline(alpha, self._control_points)

        dtheta_points = dspline_theta(alpha, self._control_points,
                                      dspline_points, ddspline_points)

        return dtheta_points / numpy.linalg.norm(dspline_points, axis=0)

    def ddtheta(self,
                distance,
                dspline_points=None,
                ddspline_points=None,
                dddspline_points=None):
        """Returns the angular acceleration as a function of distance."""
        alpha = self.distance_to_alpha(distance)
        if dspline_points is None:
            dspline_points = dspline(alpha, self._control_points)
        if ddspline_points is None:
            ddspline_points = ddspline(alpha, self._control_points)
        if dddspline_points is None:
            dddspline_points = dddspline(alpha, self._control_points)

        dtheta_points = dspline_theta(alpha, self._control_points,
                                      dspline_points, ddspline_points)
        ddtheta_points = ddspline_theta(alpha, self._control_points,
                                        dspline_points, ddspline_points,
                                        dddspline_points)

        # TODO(austin): Factor out the d^alpha/dd^2.
        return ddtheta_points / numpy.linalg.norm(
            dspline_points, axis=0)**2.0 - numpy.multiply(
                dtheta_points, (numpy.array(dspline_points)[0, :] *
                                numpy.array(ddspline_points)[0, :] +
                                numpy.array(dspline_points)[1, :] *
                                numpy.array(ddspline_points)[1, :]) /
                ((numpy.array(dspline_points)[0, :]**2.0 +
                  numpy.array(dspline_points)[1, :]**2.0)**2.0))



def main(argv):
    # Build up the control point matrix
    start = numpy.matrix([[0.0, 0.0]]).T
    c1 = numpy.matrix([[0.5, 0.0]]).T
    c2 = numpy.matrix([[0.5, 1.0]]).T
    end = numpy.matrix([[1.0, 1.0]]).T
    control_points = numpy.hstack((start, c1, c2, end))

    # The alphas to plot
    alphas = numpy.linspace(0.0, 1.0, 1000)

    # Compute x, y and the 3 derivatives
    spline_points = spline(alphas, control_points)
    dspline_points = dspline(alphas, control_points)
    ddspline_points = ddspline(alphas, control_points)
    dddspline_points = dddspline(alphas, control_points)

    # Compute theta and the two derivatives
    theta = spline_theta(alphas, control_points, dspline_points=dspline_points)
    dtheta = dspline_theta(
        alphas, control_points, dspline_points=dspline_points)
    ddtheta = ddspline_theta(
        alphas,
        control_points,
        dspline_points=dspline_points,
        dddspline_points=dddspline_points)

    # Plot the control points and the spline.
    pylab.figure()
    pylab.plot(
        numpy.array(control_points)[0, :],
        numpy.array(control_points)[1, :],
        '-o',
        label='control')
    pylab.plot(
        numpy.array(spline_points)[0, :],
        numpy.array(spline_points)[1, :],
        label='spline')
    pylab.legend()

    # For grins, confirm that the double integral of the acceleration (with
    # respect to the spline parameter) matches the position.  This lets us
    # confirm that the derivatives are consistent.
    xint_plot = numpy.matrix(numpy.zeros((2, len(alphas))))
    dxint_plot = xint_plot.copy()
    xint = spline_points[:, 0].copy()
    dxint = dspline_points[:, 0].copy()
    xint_plot[:, 0] = xint
    dxint_plot[:, 0] = dxint
    for i in range(len(alphas) - 1):
        xint += (alphas[i + 1] - alphas[i]) * dxint
        dxint += (alphas[i + 1] - alphas[i]) * ddspline_points[:, i]
        xint_plot[:, i + 1] = xint
        dxint_plot[:, i + 1] = dxint

    # Integrate up the spline velocity and heading to confirm that given a
    # velocity (as a function of the spline parameter) and angle, we will move
    # from the starting point to the ending point.
    thetaint_plot = numpy.zeros((len(alphas),))
    thetaint = theta[0]
    dthetaint_plot = numpy.zeros((len(alphas),))
    dthetaint = dtheta[0]
    thetaint_plot[0] = thetaint
    dthetaint_plot[0] = dthetaint

    txint_plot = numpy.matrix(numpy.zeros((2, len(alphas))))
    txint = spline_points[:, 0].copy()
    txint_plot[:, 0] = txint
    for i in range(len(alphas) - 1):
        dalpha = alphas[i + 1] - alphas[i]
        txint += dalpha * numpy.linalg.norm(
            dspline_points[:, i]) * numpy.matrix(
                [[numpy.cos(theta[i])], [numpy.sin(theta[i])]])
        txint_plot[:, i + 1] = txint
        thetaint += dalpha * dtheta[i]
        dthetaint += dalpha * ddtheta[i]
        thetaint_plot[i + 1] = thetaint
        dthetaint_plot[i + 1] = dthetaint


    # Now plot x, dx/dalpha, ddx/ddalpha, dddx/dddalpha, and integrals thereof
    # to perform consistency checks.
    pylab.figure()
    pylab.plot(alphas, numpy.array(spline_points)[0, :], label='x')
    pylab.plot(alphas, numpy.array(xint_plot)[0, :], label='ix')
    pylab.plot(alphas, numpy.array(dspline_points)[0, :], label='dx')
    pylab.plot(alphas, numpy.array(dxint_plot)[0, :], label='idx')
    pylab.plot(alphas, numpy.array(txint_plot)[0, :], label='tix')
    pylab.plot(alphas, numpy.array(ddspline_points)[0, :], label='ddx')
    pylab.plot(alphas, numpy.array(dddspline_points)[0, :], label='dddx')
    pylab.legend()

    # Now do the same for y.
    pylab.figure()
    pylab.plot(alphas, numpy.array(spline_points)[1, :], label='y')
    pylab.plot(alphas, numpy.array(xint_plot)[1, :], label='iy')
    pylab.plot(alphas, numpy.array(dspline_points)[1, :], label='dy')
    pylab.plot(alphas, numpy.array(dxint_plot)[1, :], label='idy')
    pylab.plot(alphas, numpy.array(txint_plot)[1, :], label='tiy')
    pylab.plot(alphas, numpy.array(ddspline_points)[1, :], label='ddy')
    pylab.plot(alphas, numpy.array(dddspline_points)[1, :], label='dddy')
    pylab.legend()

    # And for theta.
    pylab.figure()
    pylab.plot(alphas, theta, label='theta')
    pylab.plot(alphas, dtheta, label='dtheta')
    pylab.plot(alphas, ddtheta, label='ddtheta')
    pylab.plot(alphas, thetaint_plot, label='thetai')
    pylab.plot(alphas, dthetaint_plot, label='dthetai')
    pylab.plot(
        alphas,
        numpy.linalg.norm(
            numpy.array(dspline_points), axis=0),
        label='velocity')

    # Now, repeat as a function of path length as opposed to alpha
    path = Path(control_points)
    distance_count = 1000
    position = path.xy(0.0)
    velocity = path.dxy(0.0)
    theta = path.theta(0.0)
    omega = path.dtheta(0.0)

    iposition_plot = numpy.matrix(numpy.zeros((2, distance_count)))
    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[0] = theta
    iomega_plot[0] = omega

    distances = numpy.linspace(0.0, path.length(), distance_count)

    for i in xrange(len(distances) - 1):
        position += velocity * (distances[i + 1] - distances[i])
        velocity += path.ddxy(distances[i]) * (distances[i + 1] - distances[i])
        iposition_plot[:, i + 1] = position
        ivelocity_plot[:, i + 1] = velocity

        theta += omega * (distances[i + 1] - distances[i])
        omega += path.ddtheta(distances[i]) * (distances[i + 1] - distances[i])
        itheta_plot[i + 1] = theta
        iomega_plot[i + 1] = omega

    pylab.figure()
    pylab.plot(distances, numpy.array(path.xy(distances))[0, :], label='x')
    pylab.plot(distances, numpy.array(iposition_plot)[0, :], label='ix')
    pylab.plot(distances, numpy.array(path.dxy(distances))[0, :], label='dx')
    pylab.plot(distances, numpy.array(ivelocity_plot)[0, :], label='idx')
    pylab.plot(distances, numpy.array(path.ddxy(distances))[0, :], label='ddx')
    pylab.legend()

    pylab.figure()
    pylab.plot(distances, numpy.array(path.xy(distances))[1, :], label='y')
    pylab.plot(distances, numpy.array(iposition_plot)[1, :], label='iy')
    pylab.plot(distances, numpy.array(path.dxy(distances))[1, :], label='dy')
    pylab.plot(distances, numpy.array(ivelocity_plot)[1, :], label='idy')
    pylab.plot(distances, numpy.array(path.ddxy(distances))[1, :], label='ddy')
    pylab.legend()

    pylab.figure()
    pylab.plot(distances, path.theta(distances), label='theta')
    pylab.plot(distances, itheta_plot, label='itheta')
    pylab.plot(distances, path.dtheta(distances), label='omega')
    pylab.plot(distances, iomega_plot, label='iomega')
    pylab.plot(distances, path.ddtheta(distances), label='alpha')
    pylab.legend()

    # TODO(austin): Start creating a velocity plan now that we have all the
    # derivitives of our spline.

    velocity_drivetrain = polydrivetrain.VelocityDrivetrainModel(
        y2018.control_loops.python.drivetrain.kDrivetrain)

    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()

    longitudal_accel = 15.0
    lateral_accel = 20.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))))

        # 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

        # 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

        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)

    pylab.figure()
    pylab.plot(distances, lateral_accel_plan, label='accel pass')
    pylab.plot(distances, curvature_plan, label='voltage limit pass')
    pylab.xlabel("distance along spline (m)")
    pylab.ylabel("velocity (m/s)")
    pylab.legend()

    pylab.show()


if __name__ == '__main__':
    argv = FLAGS(sys.argv)
    sys.exit(main(argv))