Update Catapult to mostly just be an AngularSystem

The code was mostly duplicated, just the analysis/plotting was
different.  Update the python file to match.  This lets us wrap position
loops around it all to return the catapult to the starting position.

Change-Id: I7a60f1d07f812c91cca767d6c737f8ae6c21e8d2
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/y2022/control_loops/python/catapult_lib.py b/y2022/control_loops/python/catapult_lib.py
index 2606c44..d6040d1 100644
--- a/y2022/control_loops/python/catapult_lib.py
+++ b/y2022/control_loops/python/catapult_lib.py
@@ -1,180 +1,33 @@
+from frc971.control_loops.python import angular_system
 from frc971.control_loops.python import control_loop
 from frc971.control_loops.python import controls
+from aos.util.trapezoid_profile import TrapezoidProfile
 import numpy
 from matplotlib import pylab
 
 import gflags
 import glog
 
-
-class CatapultParams(object):
-    def __init__(self,
-                 name,
-                 motor,
-                 G,
-                 J,
-                 lever,
-                 q_pos,
-                 q_vel,
-                 q_voltage,
-                 r_pos,
-                 controller_poles,
-                 dt=0.00505):
-        self.name = name
-        self.motor = motor
-        self.G = G
-        self.J = J
-        self.lever = lever
-        self.q_pos = q_pos
-        self.q_vel = q_vel
-        self.q_voltage = q_voltage
-        self.r_pos = r_pos
-        self.dt = dt
-        self.controller_poles = controller_poles
+CatapultParams = angular_system.AngularSystemParams
 
 
-class VelocityCatapult(control_loop.HybridControlLoop):
+# TODO(austin): This is mostly the same as angular_system.  Can we either wrap an angular_system or assign it?
+class Catapult(angular_system.AngularSystem):
     def __init__(self, params, name="Catapult"):
-        super(VelocityCatapult, self).__init__(name=name)
-        self.params = params
-        # Set Motor
-        self.motor = self.params.motor
-        # Gear ratio
-        self.G = self.params.G
-        # Moment of inertia of the catapult wheel in kg m^2
-        self.J = self.params.J + self.motor.motor_inertia / (self.G**2.0)
-        # Control loop time step
-        self.dt = self.params.dt
-
-        # State feedback matrices
-        # [angular velocity]
-        self.A_continuous = numpy.matrix([[
-            -self.motor.Kt / self.motor.Kv /
-            (self.J * self.G * self.G * self.motor.resistance)
-        ]])
-        self.B_continuous = numpy.matrix(
-            [[self.motor.Kt / (self.J * self.G * self.motor.resistance)]])
-        self.C = numpy.matrix([[1]])
-        self.D = numpy.matrix([[0]])
-
-        self.A, self.B = self.ContinuousToDiscrete(self.A_continuous,
-                                                   self.B_continuous, self.dt)
-
-        self.PlaceControllerPoles(self.params.controller_poles)
-
-        # Generated controller not used.
-        self.PlaceObserverPoles([0.3])
-
-        self.U_max = numpy.matrix([[12.0]])
-        self.U_min = numpy.matrix([[-12.0]])
-
-        qff_vel = 8.0
-        self.Qff = numpy.matrix([[1.0 / (qff_vel**2.0)]])
-
-        self.Kff = controls.TwoStateFeedForwards(self.B, self.Qff)
-
-        glog.debug('K: %s', str(self.K))
-        glog.debug('Poles: %s',
-                   str(numpy.linalg.eig(self.A - self.B * self.K)[0]))
-
-
-class Catapult(VelocityCatapult):
-    def __init__(self, params, name="Catapult"):
-        super(Catapult, self).__init__(params, name=name)
-
-        self.A_continuous_unaugmented = self.A_continuous
-        self.B_continuous_unaugmented = self.B_continuous
-
-        self.A_continuous = numpy.matrix(numpy.zeros((2, 2)))
-        self.A_continuous[1:2, 1:2] = self.A_continuous_unaugmented
-        self.A_continuous[0, 1] = 1
-
-        self.B_continuous = numpy.matrix(numpy.zeros((2, 1)))
-        self.B_continuous[1:2, 0] = self.B_continuous_unaugmented
-
-        # State feedback matrices
-        # [position, angular velocity]
-        self.C = numpy.matrix([[1, 0]])
-        self.D = numpy.matrix([[0]])
-
-        self.A, self.B = self.ContinuousToDiscrete(self.A_continuous,
-                                                   self.B_continuous, self.dt)
-
-        rpl = 0.45
-        ipl = 0.07
-        self.PlaceObserverPoles([rpl + 1j * ipl, rpl - 1j * ipl])
-
-        self.K_unaugmented = self.K
-        self.K = numpy.matrix(numpy.zeros((1, 2)))
-        self.K[0, 1:2] = self.K_unaugmented
-        self.Kff_unaugmented = self.Kff
-        self.Kff = numpy.matrix(numpy.zeros((1, 2)))
-        self.Kff[0, 1:2] = self.Kff_unaugmented
+        super(Catapult, self).__init__(params, name)
+        # Signal that we have a single cycle output delay to compensate for in
+        # our observer.
+        self.delayed_u = True
 
         self.InitializeState()
 
 
-class IntegralCatapult(Catapult):
+class IntegralCatapult(angular_system.IntegralAngularSystem):
     def __init__(self, params, name="IntegralCatapult"):
         super(IntegralCatapult, self).__init__(params, name=name)
-
-        self.A_continuous_unaugmented = self.A_continuous
-        self.B_continuous_unaugmented = self.B_continuous
-
-        self.A_continuous = numpy.matrix(numpy.zeros((3, 3)))
-        self.A_continuous[0:2, 0:2] = self.A_continuous_unaugmented
-        self.A_continuous[0:2, 2] = self.B_continuous_unaugmented
-
-        self.B_continuous = numpy.matrix(numpy.zeros((3, 1)))
-        self.B_continuous[0:2, 0] = self.B_continuous_unaugmented
-
-        # states
-        # [position, velocity, voltage_error]
-        self.C_unaugmented = self.C
-        self.C = numpy.matrix(numpy.zeros((1, 3)))
-        self.C[0:1, 0:2] = self.C_unaugmented
-
-        glog.debug('A_continuous %s' % str(self.A_continuous))
-        glog.debug('B_continuous %s' % str(self.B_continuous))
-        glog.debug('C %s' % str(self.C))
-
-        self.A, self.B = self.ContinuousToDiscrete(self.A_continuous,
-                                                   self.B_continuous, self.dt)
-
-        glog.debug('A %s' % str(self.A))
-        glog.debug('B %s' % str(self.B))
-
-        q_pos = self.params.q_pos
-        q_vel = self.params.q_vel
-        q_voltage = self.params.q_voltage
-        self.Q_continuous = numpy.matrix([[(q_pos**2.0), 0.0,
-                                           0.0], [0.0, (q_vel**2.0), 0.0],
-                                          [0.0, 0.0, (q_voltage**2.0)]])
-
-        r_pos = self.params.r_pos
-        self.R_continuous = numpy.matrix([[(r_pos**2.0)]])
-
-        _, _, 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)
-
-        glog.debug('Q_discrete %s' % (str(self.Q)))
-        glog.debug('R_discrete %s' % (str(self.R)))
-
-        self.KalmanGain, self.P_steady_state = controls.kalman(
-            A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
-        self.L = self.A * self.KalmanGain
-
-        self.K_unaugmented = self.K
-        self.K = numpy.matrix(numpy.zeros((1, 3)))
-        self.K[0, 0:2] = self.K_unaugmented
-        self.K[0, 2] = 1
-        self.Kff_unaugmented = self.Kff
-        self.Kff = numpy.matrix(numpy.zeros((1, 3)))
-        self.Kff[0, 0:2] = self.Kff_unaugmented
+        # Signal that we have a single cycle output delay to compensate for in
+        # our observer.
+        self.delayed_u = True
 
         self.InitializeState()
 
@@ -201,7 +54,7 @@
     while True:
         X_hat = catapult.X
         if catapult.X[0, 0] > final_position:
-            return catapult.X[1, 0] * params.lever
+            return catapult.X[1, 0] * params.radius
 
         if observer_catapult is not None:
             X_hat = observer_catapult.X_hat
@@ -210,13 +63,13 @@
 
         if observer_catapult is not None:
             observer_catapult.Y = catapult.Y
-            observer_catapult.CorrectHybridObserver(U)
+            observer_catapult.CorrectObserver(U)
 
         applied_U = U.copy()
         catapult.Update(applied_U)
 
         if observer_catapult is not None:
-            observer_catapult.PredictHybridObserver(U, catapult.dt)
+            observer_catapult.PredictObserver(U)
 
 
 def PlotShot(params, U, final_position):
@@ -262,7 +115,7 @@
             X_hat = observer_catapult.X_hat
             x_hat.append(observer_catapult.X_hat[0, 0])
             w_hat.append(observer_catapult.X_hat[1, 0])
-            v_hat.append(observer_catapult.X_hat[1, 0] * params.lever)
+            v_hat.append(observer_catapult.X_hat[1, 0] * params.radius)
 
         U[0, 0] = numpy.clip(U[0, 0], -vbat, vbat)
         x.append(catapult.X[0, 0])
@@ -277,13 +130,13 @@
 
         if observer_catapult is not None:
             observer_catapult.Y = catapult.Y
-            observer_catapult.CorrectHybridObserver(U)
+            observer_catapult.CorrectObserver(U)
             offset.append(observer_catapult.X_hat[2, 0])
 
         catapult.Update(U)
 
         if observer_catapult is not None:
-            observer_catapult.PredictHybridObserver(U, catapult.dt)
+            observer_catapult.PredictObserver(U)
 
         t.append(initial_t + i * catapult.dt)
         u.append(U[0, 0])
@@ -307,6 +160,95 @@
 
     pylab.show()
 
+
+RunTest = angular_system.RunTest
+
+
+def PlotStep(params, R, plant_params=None):
+    """Plots a step move to the goal.
+
+    Args:
+      params: CatapultParams for the controller and observer
+      plant_params: CatapultParams for the plant.  Defaults to params if
+        plant_params is None.
+      R: numpy.matrix(2, 1), the goal"""
+    plant = Catapult(plant_params or params, params.name)
+    controller = IntegralCatapult(params, params.name)
+    observer = IntegralCatapult(params, params.name)
+
+    # Test moving the system.
+    initial_X = numpy.matrix([[0.0], [0.0]])
+    augmented_R = numpy.matrix(numpy.zeros((3, 1)))
+    augmented_R[0:2, :] = R
+    RunTest(plant,
+            end_goal=augmented_R,
+            controller=controller,
+            observer=observer,
+            duration=2.0,
+            use_profile=False,
+            kick_time=1.0,
+            kick_magnitude=0.0)
+
+
+def PlotKick(params, R, plant_params=None):
+    """Plots a step motion with a kick at 1.0 seconds.
+
+    Args:
+      params: CatapultParams for the controller and observer
+      plant_params: CatapultParams for the plant.  Defaults to params if
+        plant_params is None.
+      R: numpy.matrix(2, 1), the goal"""
+    plant = Catapult(plant_params or params, params.name)
+    controller = IntegralCatapult(params, params.name)
+    observer = IntegralCatapult(params, params.name)
+
+    # Test moving the system.
+    initial_X = numpy.matrix([[0.0], [0.0]])
+    augmented_R = numpy.matrix(numpy.zeros((3, 1)))
+    augmented_R[0:2, :] = R
+    RunTest(plant,
+            end_goal=augmented_R,
+            controller=controller,
+            observer=observer,
+            duration=2.0,
+            use_profile=False,
+            kick_time=1.0,
+            kick_magnitude=2.0)
+
+
+def PlotMotion(params,
+               R,
+               max_velocity=10.0,
+               max_acceleration=70.0,
+               plant_params=None):
+    """Plots a trapezoidal motion.
+
+    Args:
+      params: CatapultParams for the controller and observer
+      plant_params: CatapultParams for the plant.  Defaults to params if
+        plant_params is None.
+      R: numpy.matrix(2, 1), the goal,
+      max_velocity: float, The max velocity of the profile.
+      max_acceleration: float, The max acceleration of the profile.
+    """
+    plant = Catapult(plant_params or params, params.name)
+    controller = IntegralCatapult(params, params.name)
+    observer = IntegralCatapult(params, params.name)
+
+    # Test moving the system.
+    initial_X = numpy.matrix([[0.0], [0.0]])
+    augmented_R = numpy.matrix(numpy.zeros((3, 1)))
+    augmented_R[0:2, :] = R
+    RunTest(plant,
+            end_goal=augmented_R,
+            controller=controller,
+            observer=observer,
+            duration=2.0,
+            use_profile=True,
+            max_velocity=max_velocity,
+            max_acceleration=max_acceleration)
+
+
 def WriteCatapult(params, plant_files, controller_files, year_namespaces):
     """Writes out the constants for a catapult to a file.
 
@@ -325,19 +267,18 @@
     if type(params) is list:
         name = params[0].name
         for index, param in enumerate(params):
-            catapults.append(
-                Catapult(param, param.name + str(index)))
+            catapults.append(Catapult(param, param.name + str(index)))
             integral_catapults.append(
-                IntegralCatapult(param,
-                                      'Integral' + param.name + str(index)))
+                IntegralCatapult(param, 'Integral' + param.name + str(index)))
     else:
         name = params.name
         catapults.append(Catapult(params, params.name))
         integral_catapults.append(
             IntegralCatapult(params, 'Integral' + params.name))
 
-    loop_writer = control_loop.ControlLoopWriter(
-        name, catapults, namespaces=year_namespaces)
+    loop_writer = control_loop.ControlLoopWriter(name,
+                                                 catapults,
+                                                 namespaces=year_namespaces)
     loop_writer.AddConstant(
         control_loop.Constant('kOutputRatio', '%f', catapults[0].G))
     loop_writer.AddConstant(