frc971/control_loops/python/polydrivetrain.py

#!/usr/bin/python3

import numpy
from frc971.control_loops.python import polytope
import frc971.control_loops.python.drivetrain
from frc971.control_loops.python import control_loop
from frc971.control_loops.python import controls
from frc971.control_loops.python.cim import CIM
from matplotlib import pylab

import glog


def CoerceGoal(region, K, w, R):
    """Intersects a line with a region, and finds the closest point to R.

    Finds a point that is closest to R inside the region, and on the line
    defined by K X = w.  If it is not possible to find a point on the line,
    finds a point that is inside the region and closest to the line.  This
    function assumes that

    Args:
        region: HPolytope, the valid goal region.
        K: numpy.matrix (2 x 1), the matrix for the equation [K1, K2] [x1; x2] = w
        w: float, the offset in the equation above.
        R: numpy.matrix (2 x 1), the point to be closest to.

    Returns:
        numpy.matrix (2 x 1), the point.
    """
    return DoCoerceGoal(region, K, w, R)[0]


def DoCoerceGoal(region, K, w, R):
    if region.IsInside(R):
        return (R, True)

    perpendicular_vector = K.T / numpy.linalg.norm(K)
    parallel_vector = numpy.matrix([[perpendicular_vector[1, 0]],
                                    [-perpendicular_vector[0, 0]]])

    # We want to impose the constraint K * X = w on the polytope H * X <= k.
    # We do this by breaking X up into parallel and perpendicular components to
    # the half plane.  This gives us the following equation.
    #
    #  parallel * (parallel.T \dot X) + perpendicular * (perpendicular \dot X)) = X
    #
    # Then, substitute this into the polytope.
    #
    #  H * (parallel * (parallel.T \dot X) + perpendicular * (perpendicular \dot X)) <= k
    #
    # Substitute K * X = w
    #
    # H * parallel * (parallel.T \dot X) + H * perpendicular * w <= k
    #
    # Move all the knowns to the right side.
    #
    # H * parallel * ([parallel1 parallel2] * X) <= k - H * perpendicular * w
    #
    # Let t = parallel.T \dot X, the component parallel to the surface.
    #
    # H * parallel * t <= k - H * perpendicular * w
    #
    # This is a polytope which we can solve, and use to figure out the range of X
    # that we care about!

    t_poly = polytope.HPolytope(region.H * parallel_vector,
                                region.k - region.H * perpendicular_vector * w)

    vertices = t_poly.Vertices()

    if vertices.shape[0]:
        # The region exists!
        # Find the closest vertex
        min_distance = numpy.infty
        closest_point = None
        for vertex in vertices:
            point = parallel_vector * vertex + perpendicular_vector * w
            length = numpy.linalg.norm(R - point)
            if length < min_distance:
                min_distance = length
                closest_point = point

        return (closest_point, True)
    else:
        # Find the vertex of the space that is closest to the line.
        region_vertices = region.Vertices()
        min_distance = numpy.infty
        closest_point = None
        for vertex in region_vertices:
            point = vertex.T
            length = numpy.abs((perpendicular_vector.T * point)[0, 0])
            if length < min_distance:
                min_distance = length
                closest_point = point

        return (closest_point, False)


class VelocityDrivetrainModel(control_loop.ControlLoop):

    def __init__(self,
                 drivetrain_params,
                 left_low=True,
                 right_low=True,
                 name="VelocityDrivetrainModel"):
        super(VelocityDrivetrainModel, self).__init__(name)
        self._drivetrain = frc971.control_loops.python.drivetrain.Drivetrain(
            left_low=left_low,
            right_low=right_low,
            drivetrain_params=drivetrain_params)
        self.dt = drivetrain_params.dt
        self.A_continuous = numpy.matrix([[
            self._drivetrain.A_continuous[1, 1],
            self._drivetrain.A_continuous[1, 3]
        ],
                                          [
                                              self._drivetrain.A_continuous[3,
                                                                            1],
                                              self._drivetrain.A_continuous[3,
                                                                            3]
                                          ]])

        self.B_continuous = numpy.matrix([[
            self._drivetrain.B_continuous[1, 0],
            self._drivetrain.B_continuous[1, 1]
        ],
                                          [
                                              self._drivetrain.B_continuous[3,
                                                                            0],
                                              self._drivetrain.B_continuous[3,
                                                                            1]
                                          ]])
        self.C = numpy.matrix(numpy.eye(2))
        self.D = numpy.matrix(numpy.zeros((2, 2)))

        self.A, self.B = self.ContinuousToDiscrete(self.A_continuous,
                                                   self.B_continuous, self.dt)

        # FF * X = U (steady state)
        self.FF = self.B.I * (numpy.eye(2) - self.A)

        self.PlaceControllerPoles(drivetrain_params.controller_poles)
        # Build a kalman filter for the velocity.  We don't care what the gains
        # are, but the hybrid kalman filter that we want to write to disk to get
        # access to A_continuous and B_continuous needs this for completeness.
        self.Q_continuous = numpy.matrix([[(0.5**2.0), 0.0], [0.0,
                                                              (0.5**2.0)]])
        self.R_continuous = numpy.matrix([[(0.1**2.0), 0.0], [0.0,
                                                              (0.1**2.0)]])
        self.PlaceObserverPoles(drivetrain_params.observer_poles)
        _, _, self.Q, self.R = controls.kalmd(A_continuous=self.A_continuous,
                                              B_continuous=self.B_continuous,
                                              Q_continuous=self.Q_continuous,
                                              R_continuous=self.R_continuous,
                                              dt=self.dt)

        self.KalmanGain, self.P_steady_state = controls.kalman(A=self.A,
                                                               B=self.B,
                                                               C=self.C,
                                                               Q=self.Q,
                                                               R=self.R)

        self.G_high = self._drivetrain.G_high
        self.G_low = self._drivetrain.G_low
        self.resistance = self._drivetrain.resistance
        self.r = self._drivetrain.r
        self.Kv = self._drivetrain.Kv
        self.Kt = self._drivetrain.Kt

        self.U_max = self._drivetrain.U_max
        self.U_min = self._drivetrain.U_min
        self.wrap_point = numpy.matrix(numpy.zeros((2, 1)))

    @property
    def robot_radius_l(self):
        return self._drivetrain.robot_radius_l

    @property
    def robot_radius_r(self):
        return self._drivetrain.robot_radius_r


class VelocityDrivetrain(object):
    HIGH = 'high'
    LOW = 'low'
    SHIFTING_UP = 'up'
    SHIFTING_DOWN = 'down'

    def __init__(self, drivetrain_params, name='VelocityDrivetrain'):
        self.drivetrain_low_low = VelocityDrivetrainModel(
            left_low=True,
            right_low=True,
            name=name + 'LowLow',
            drivetrain_params=drivetrain_params)
        self.drivetrain_low_high = VelocityDrivetrainModel(
            left_low=True,
            right_low=False,
            name=name + 'LowHigh',
            drivetrain_params=drivetrain_params)
        self.drivetrain_high_low = VelocityDrivetrainModel(
            left_low=False,
            right_low=True,
            name=name + 'HighLow',
            drivetrain_params=drivetrain_params)
        self.drivetrain_high_high = VelocityDrivetrainModel(
            left_low=False,
            right_low=False,
            name=name + 'HighHigh',
            drivetrain_params=drivetrain_params)

        # X is [lvel, rvel]
        self.X = numpy.matrix([[0.0], [0.0]])

        self.U_poly = polytope.HPolytope(
            numpy.matrix([[1, 0], [-1, 0], [0, 1], [0, -1]]),
            numpy.matrix([[12], [12], [12], [12]]))

        self.U_max = numpy.matrix([[12.0], [12.0]])
        self.U_min = numpy.matrix([[-12.0000000000], [-12.0000000000]])

        self.dt = 0.00505

        self.R = numpy.matrix([[0.0], [0.0]])

        self.U_ideal = numpy.matrix([[0.0], [0.0]])

        # ttrust is the comprimise between having full throttle negative inertia,
        # and having no throttle negative inertia.  A value of 0 is full throttle
        # inertia.  A value of 1 is no throttle negative inertia.
        self.ttrust = 1.0

        self.left_gear = VelocityDrivetrain.LOW
        self.right_gear = VelocityDrivetrain.LOW
        self.left_shifter_position = 0.0
        self.right_shifter_position = 0.0
        self.left_cim = CIM()
        self.right_cim = CIM()

    def IsInGear(self, gear):
        return gear is VelocityDrivetrain.HIGH or gear is VelocityDrivetrain.LOW

    def MotorRPM(self, shifter_position, velocity):
        if shifter_position > 0.5:
            return (velocity / self.CurrentDrivetrain().G_high /
                    self.CurrentDrivetrain().r)
        else:
            return (velocity / self.CurrentDrivetrain().G_low /
                    self.CurrentDrivetrain().r)

    def CurrentDrivetrain(self):
        if self.left_shifter_position > 0.5:
            if self.right_shifter_position > 0.5:
                return self.drivetrain_high_high
            else:
                return self.drivetrain_high_low
        else:
            if self.right_shifter_position > 0.5:
                return self.drivetrain_low_high
            else:
                return self.drivetrain_low_low

    def SimShifter(self, gear, shifter_position):
        if gear is VelocityDrivetrain.HIGH or gear is VelocityDrivetrain.SHIFTING_UP:
            shifter_position = min(shifter_position + 0.5, 1.0)
        else:
            shifter_position = max(shifter_position - 0.5, 0.0)

        if shifter_position == 1.0:
            gear = VelocityDrivetrain.HIGH
        elif shifter_position == 0.0:
            gear = VelocityDrivetrain.LOW

        return gear, shifter_position

    def ComputeGear(self,
                    wheel_velocity,
                    should_print=False,
                    current_gear=False,
                    gear_name=None):
        high_omega = (wheel_velocity / self.CurrentDrivetrain().G_high /
                      self.CurrentDrivetrain().r)
        low_omega = (wheel_velocity / self.CurrentDrivetrain().G_low /
                     self.CurrentDrivetrain().r)
        #print gear_name, "Motor Energy Difference.", 0.5 * 0.000140032647 * (low_omega * low_omega - high_omega * high_omega), "joules"
        high_torque = ((12.0 - high_omega / self.CurrentDrivetrain().Kv) *
                       self.CurrentDrivetrain().Kt /
                       self.CurrentDrivetrain().resistance)
        low_torque = ((12.0 - low_omega / self.CurrentDrivetrain().Kv) *
                      self.CurrentDrivetrain().Kt /
                      self.CurrentDrivetrain().resistance)
        high_power = high_torque * high_omega
        low_power = low_torque * low_omega
        #if should_print:
        #  print gear_name, "High omega", high_omega, "Low omega", low_omega
        #  print gear_name, "High torque", high_torque, "Low torque", low_torque
        #  print gear_name, "High power", high_power, "Low power", low_power

        # Shift algorithm improvements.
        # TODO(aschuh):
        # It takes time to shift.  Shifting down for 1 cycle doesn't make sense
        # because you will end up slower than without shifting.  Figure out how
        # to include that info.
        # If the driver is still in high gear, but isn't asking for the extra power
        # from low gear, don't shift until he asks for it.
        goal_gear_is_high = high_power > low_power
        #goal_gear_is_high = True

        if not self.IsInGear(current_gear):
            glog.debug('%s Not in gear.', gear_name)
            return current_gear
        else:
            is_high = current_gear is VelocityDrivetrain.HIGH
            if is_high != goal_gear_is_high:
                if goal_gear_is_high:
                    glog.debug('%s Shifting up.', gear_name)
                    return VelocityDrivetrain.SHIFTING_UP
                else:
                    glog.debug('%s Shifting down.', gear_name)
                    return VelocityDrivetrain.SHIFTING_DOWN
            else:
                return current_gear

    def FilterVelocity(self, throttle):
        # Invert the plant to figure out how the velocity filter would have to work
        # out in order to filter out the forwards negative inertia.
        # This math assumes that the left and right power and velocity are equal.

        # The throttle filter should filter such that the motor in the highest gear
        # should be controlling the time constant.
        # Do this by finding the index of FF that has the lowest value, and computing
        # the sums using that index.
        FF_sum = self.CurrentDrivetrain().FF.sum(axis=1)
        min_FF_sum_index = numpy.argmin(FF_sum)
        min_FF_sum = FF_sum[min_FF_sum_index, 0]
        min_K_sum = self.CurrentDrivetrain().K[min_FF_sum_index, :].sum()
        # Compute the FF sum for high gear.
        high_min_FF_sum = self.drivetrain_high_high.FF[0, :].sum()

        # U = self.K[0, :].sum() * (R - x_avg) + self.FF[0, :].sum() * R
        # throttle * 12.0 = (self.K[0, :].sum() + self.FF[0, :].sum()) * R
        #                   - self.K[0, :].sum() * x_avg

        # R = (throttle * 12.0 + self.K[0, :].sum() * x_avg) /
        #     (self.K[0, :].sum() + self.FF[0, :].sum())

        # U = (K + FF) * R - K * X
        # (K + FF) ^-1 * (U + K * X) = R

        # Scale throttle by min_FF_sum / high_min_FF_sum.  This will make low gear
        # have the same velocity goal as high gear, and so that the robot will hold
        # the same speed for the same throttle for all gears.
        adjusted_ff_voltage = numpy.clip(
            throttle * 12.0 * min_FF_sum / high_min_FF_sum, -12.0, 12.0)
        return ((adjusted_ff_voltage + self.ttrust * min_K_sum *
                 (self.X[0, 0] + self.X[1, 0]) / 2.0) /
                (self.ttrust * min_K_sum + min_FF_sum))

    def Update(self, throttle, steering):
        # Shift into the gear which sends the most power to the floor.
        # This is the same as sending the most torque down to the floor at the
        # wheel.

        self.left_gear = self.right_gear = True
        if True:
            self.left_gear = self.ComputeGear(self.X[0, 0],
                                              should_print=True,
                                              current_gear=self.left_gear,
                                              gear_name="left")
            self.right_gear = self.ComputeGear(self.X[1, 0],
                                               should_print=True,
                                               current_gear=self.right_gear,
                                               gear_name="right")
            if self.IsInGear(self.left_gear):
                self.left_cim.X[0,
                                0] = self.MotorRPM(self.left_shifter_position,
                                                   self.X[0, 0])

            if self.IsInGear(self.right_gear):
                self.right_cim.X[0, 0] = self.MotorRPM(
                    self.right_shifter_position, self.X[0, 0])

        if self.IsInGear(self.left_gear) and self.IsInGear(self.right_gear):
            # Filter the throttle to provide a nicer response.
            fvel = self.FilterVelocity(throttle)

            # Constant radius means that angualar_velocity / linear_velocity = constant.
            # Compute the left and right velocities.
            steering_velocity = numpy.abs(fvel) * steering
            left_velocity = fvel - steering_velocity
            right_velocity = fvel + steering_velocity

            # Write this constraint in the form of K * R = w
            # angular velocity / linear velocity = constant
            # (left - right) / (left + right) = constant
            # left - right = constant * left + constant * right

            # (fvel - steering * numpy.abs(fvel) - fvel - steering * numpy.abs(fvel)) /
            #  (fvel - steering * numpy.abs(fvel) + fvel + steering * numpy.abs(fvel)) =
            #       constant
            # (- 2 * steering * numpy.abs(fvel)) / (2 * fvel) = constant
            # (-steering * sign(fvel)) = constant
            # (-steering * sign(fvel)) * (left + right) = left - right
            # (steering * sign(fvel) + 1) * left + (steering * sign(fvel) - 1) * right = 0

            equality_k = numpy.matrix([[
                1 + steering * numpy.sign(fvel),
                -(1 - steering * numpy.sign(fvel))
            ]])
            equality_w = 0.0

            self.R[0, 0] = left_velocity
            self.R[1, 0] = right_velocity

            # Construct a constraint on R by manipulating the constraint on U
            # Start out with H * U <= k
            # U = FF * R + K * (R - X)
            # H * (FF * R + K * R - K * X) <= k
            # H * (FF + K) * R <= k + H * K * X
            R_poly = polytope.HPolytope(
                self.U_poly.H *
                (self.CurrentDrivetrain().K + self.CurrentDrivetrain().FF),
                self.U_poly.k +
                self.U_poly.H * self.CurrentDrivetrain().K * self.X)

            # Limit R back inside the box.
            self.boxed_R = CoerceGoal(R_poly, equality_k, equality_w, self.R)

            FF_volts = self.CurrentDrivetrain().FF * self.boxed_R
            self.U_ideal = self.CurrentDrivetrain().K * (self.boxed_R -
                                                         self.X) + FF_volts
        else:
            glog.debug('Not all in gear')
            if not self.IsInGear(self.left_gear) and not self.IsInGear(
                    self.right_gear):
                # TODO(austin): Use battery volts here.
                R_left = self.MotorRPM(self.left_shifter_position, self.X[0,
                                                                          0])
                self.U_ideal[0, 0] = numpy.clip(
                    self.left_cim.K * (R_left - self.left_cim.X) +
                    R_left / self.left_cim.Kv, self.left_cim.U_min,
                    self.left_cim.U_max)
                self.left_cim.Update(self.U_ideal[0, 0])

                R_right = self.MotorRPM(self.right_shifter_position, self.X[1,
                                                                            0])
                self.U_ideal[1, 0] = numpy.clip(
                    self.right_cim.K * (R_right - self.right_cim.X) +
                    R_right / self.right_cim.Kv, self.right_cim.U_min,
                    self.right_cim.U_max)
                self.right_cim.Update(self.U_ideal[1, 0])
            else:
                assert False

        self.U = numpy.clip(self.U_ideal, self.U_min, self.U_max)

        # TODO(austin): Model the robot as not accelerating when you shift...
        # This hack only works when you shift at the same time.
        if self.IsInGear(self.left_gear) and self.IsInGear(self.right_gear):
            self.X = self.CurrentDrivetrain(
            ).A * self.X + self.CurrentDrivetrain().B * self.U

        self.left_gear, self.left_shifter_position = self.SimShifter(
            self.left_gear, self.left_shifter_position)
        self.right_gear, self.right_shifter_position = self.SimShifter(
            self.right_gear, self.right_shifter_position)

        glog.debug('U is %s %s', str(self.U[0, 0]), str(self.U[1, 0]))
        glog.debug('Left shifter %s %d Right shifter %s %d', self.left_gear,
                   self.left_shifter_position, self.right_gear,
                   self.right_shifter_position)


def WritePolyDrivetrain(drivetrain_files,
                        motor_files,
                        hybrid_files,
                        year_namespace,
                        drivetrain_params,
                        scalar_type='double'):
    vdrivetrain = VelocityDrivetrain(drivetrain_params)
    hybrid_vdrivetrain = VelocityDrivetrain(drivetrain_params,
                                            name="HybridVelocityDrivetrain")
    if isinstance(year_namespace, list):
        namespaces = year_namespace
    else:
        namespaces = [year_namespace, 'control_loops', 'drivetrain']
    dog_loop_writer = control_loop.ControlLoopWriter("VelocityDrivetrain", [
        vdrivetrain.drivetrain_low_low, vdrivetrain.drivetrain_low_high,
        vdrivetrain.drivetrain_high_low, vdrivetrain.drivetrain_high_high
    ],
                                                     namespaces=namespaces,
                                                     scalar_type=scalar_type)

    dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1],
                          drivetrain_files[2], "velocity_drivetrain_loop")

    hybrid_loop_writer = control_loop.ControlLoopWriter(
        "HybridVelocityDrivetrain", [
            hybrid_vdrivetrain.drivetrain_low_low,
            hybrid_vdrivetrain.drivetrain_low_high,
            hybrid_vdrivetrain.drivetrain_high_low,
            hybrid_vdrivetrain.drivetrain_high_high
        ],
        namespaces=namespaces,
        scalar_type=scalar_type,
        plant_type='StateFeedbackHybridPlant',
        observer_type='HybridKalman')

    hybrid_loop_writer.Write(hybrid_files[0], hybrid_files[1], hybrid_files[2],
                             "hybrid_velocity_drivetrain_loop")

    cim_writer = control_loop.ControlLoopWriter("CIM", [CIM()],
                                                scalar_type=scalar_type)

    cim_writer.Write(motor_files[0], motor_files[1], motor_files[2])


def PlotPolyDrivetrainMotions(drivetrain_params):
    vdrivetrain = VelocityDrivetrain(drivetrain_params)
    vl_plot = []
    vr_plot = []
    ul_plot = []
    ur_plot = []
    radius_plot = []
    t_plot = []
    left_gear_plot = []
    right_gear_plot = []
    vdrivetrain.left_shifter_position = 0.0
    vdrivetrain.right_shifter_position = 0.0
    vdrivetrain.left_gear = VelocityDrivetrain.LOW
    vdrivetrain.right_gear = VelocityDrivetrain.LOW

    glog.debug('K is %s', str(vdrivetrain.CurrentDrivetrain().K))

    if vdrivetrain.left_gear is VelocityDrivetrain.HIGH:
        glog.debug('Left is high')
    else:
        glog.debug('Left is low')
    if vdrivetrain.right_gear is VelocityDrivetrain.HIGH:
        glog.debug('Right is high')
    else:
        glog.debug('Right is low')

    for t in numpy.arange(0, 1.7, vdrivetrain.dt):
        if t < 0.5:
            vdrivetrain.Update(throttle=0.00, steering=1.0)
        elif t < 1.2:
            vdrivetrain.Update(throttle=0.5, steering=1.0)
        else:
            vdrivetrain.Update(throttle=0.00, steering=1.0)
        t_plot.append(t)
        vl_plot.append(vdrivetrain.X[0, 0])
        vr_plot.append(vdrivetrain.X[1, 0])
        ul_plot.append(vdrivetrain.U[0, 0])
        ur_plot.append(vdrivetrain.U[1, 0])
        left_gear_plot.append(
            (vdrivetrain.left_gear is VelocityDrivetrain.HIGH) * 2.0 - 10.0)
        right_gear_plot.append(
            (vdrivetrain.right_gear is VelocityDrivetrain.HIGH) * 2.0 - 10.0)

        fwd_velocity = (vdrivetrain.X[1, 0] + vdrivetrain.X[0, 0]) / 2
        turn_velocity = (vdrivetrain.X[1, 0] - vdrivetrain.X[0, 0])
        if abs(fwd_velocity) < 0.0000001:
            radius_plot.append(turn_velocity)
        else:
            radius_plot.append(turn_velocity / fwd_velocity)

    # TODO(austin):
    # Shifting compensation.

    # Tighten the turn.
    # Closed loop drive.

    pylab.plot(t_plot, vl_plot, label='left velocity')
    pylab.plot(t_plot, vr_plot, label='right velocity')
    pylab.plot(t_plot, ul_plot, label='left voltage')
    pylab.plot(t_plot, ur_plot, label='right voltage')
    pylab.plot(t_plot, radius_plot, label='radius')
    pylab.plot(t_plot, left_gear_plot, label='left gear high')
    pylab.plot(t_plot, right_gear_plot, label='right gear high')
    pylab.legend()
    pylab.show()