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/BUILD b/y2022/control_loops/python/BUILD
index 286b874..17f7b2a 100644
--- a/y2022/control_loops/python/BUILD
+++ b/y2022/control_loops/python/BUILD
@@ -55,6 +55,8 @@
     ],
     target_compatible_with = ["@platforms//cpu:x86_64"],
     deps = [
+        "//aos/util:py_trapezoid_profile",
+        "//frc971/control_loops/python:angular_system",
         "//frc971/control_loops/python:controls",
         "@matplotlib_repo//:matplotlib3",
     ],
diff --git a/y2022/control_loops/python/catapult.py b/y2022/control_loops/python/catapult.py
index 491286e..ce7496a 100755
--- a/y2022/control_loops/python/catapult.py
+++ b/y2022/control_loops/python/catapult.py
@@ -40,27 +40,35 @@
 M_cup = (1750 * 0.0254 * 0.04 * 2 * math.pi * (ball_diameter / 2.)**2.0)
 J_cup = M_cup * lever**2.0 + M_cup * (ball_diameter / 2.)**2.0
 
-print("J ball", ball_mass * lever * lever)
-print("J bar", J_bar)
-print("bar mass", M_bar)
-print("J cup", J_cup)
-print("cup mass", M_cup)
 
 J = (J_ball + J_bar + J_cup * 1.5)
-print("J", J)
+JEmpty = (J_bar + J_cup * 1.5)
 
-kCatapult = catapult_lib.CatapultParams(
+kCatapultWithBall = catapult_lib.CatapultParams(
     name='Catapult',
     motor=AddResistance(control_loop.NMotor(control_loop.Falcon(), 2), 0.03),
     G=G,
     J=J,
-    lever=lever,
-    q_pos=0.01,
-    q_vel=10.0,
-    q_voltage=4.0,
-    r_pos=0.01,
-    controller_poles=[.93],
-    dt=0.00505)
+    radius=lever,
+    q_pos=2.8,
+    q_vel=20.0,
+    kalman_q_pos=0.12,
+    kalman_q_vel=1.0,
+    kalman_q_voltage=1.5,
+    kalman_r_position=0.05)
+
+kCatapultEmpty = catapult_lib.CatapultParams(
+    name='Catapult',
+    motor=AddResistance(control_loop.NMotor(control_loop.Falcon(), 2), 0.03),
+    G=G,
+    J=JEmpty,
+    radius=lever,
+    q_pos=2.8,
+    q_vel=20.0,
+    kalman_q_pos=0.12,
+    kalman_q_vel=1.0,
+    kalman_q_voltage=1.5,
+    kalman_r_position=0.05)
 
 # Ideas for adjusting the cost function:
 #
@@ -252,7 +260,7 @@
 
 
 def CatapultProblem():
-    c = catapult_lib.Catapult(kCatapult)
+    c = catapult_lib.Catapult(kCatapultWithBall)
 
     kHorizon = 40
 
@@ -300,54 +308,69 @@
 
 
 def main(argv):
-    # Do all our math with a lower voltage so we have headroom.
-    U = numpy.matrix([[9.0]])
-
-    prob = osqp.OSQP()
-
-    kHorizon = 40
-    catapult = catapult_lib.Catapult(kCatapult)
-    X_initial = numpy.matrix([[0.0], [0.0]])
-    X_final = numpy.matrix([[2.0], [25.0]])
-    P, q, c = quadratic_cost(catapult, X_initial, X_final, kHorizon)
-    A = numpy.identity(kHorizon)
-    l = numpy.zeros((kHorizon, 1))
-    u = numpy.ones((kHorizon, 1)) * 12.0
-
-    prob.setup(scipy.sparse.csr_matrix(P),
-               q,
-               scipy.sparse.csr_matrix(A),
-               l,
-               u,
-               warm_start=True)
-
-    result = prob.solve()
-    # Check solver status
-    if result.info.status != 'solved':
-        raise ValueError('OSQP did not solve the problem!')
-
-    # Apply first control input to the plant
-    print(result.x)
-
     if FLAGS.plot:
-        print(
+        # Do all our math with a lower voltage so we have headroom.
+        U = numpy.matrix([[9.0]])
+
+        prob = osqp.OSQP()
+
+        kHorizon = 40
+        catapult = catapult_lib.Catapult(kCatapultWithBall)
+        X_initial = numpy.matrix([[0.0], [0.0]])
+        X_final = numpy.matrix([[2.0], [25.0]])
+        P, q, c = quadratic_cost(catapult, X_initial, X_final, kHorizon)
+        A = numpy.identity(kHorizon)
+        l = numpy.zeros((kHorizon, 1))
+        u = numpy.ones((kHorizon, 1)) * 12.0
+
+        prob.setup(scipy.sparse.csr_matrix(P),
+                   q,
+                   scipy.sparse.csr_matrix(A),
+                   l,
+                   u,
+                   warm_start=True)
+
+        result = prob.solve()
+        # Check solver status
+        if result.info.status != 'solved':
+            raise ValueError('OSQP did not solve the problem!')
+
+        # Apply first control input to the plant
+        print(result.x)
+
+        glog.debug("J ball", ball_mass * lever * lever)
+        glog.debug("J bar", J_bar)
+        glog.debug("bar mass", M_bar)
+        glog.debug("J cup", J_cup)
+        glog.debug("cup mass", M_cup)
+        glog.debug("J", J)
+        glog.debug("J Empty", JEmpty)
+
+        glog.debug(
             "For G:", G, " max speed ",
-            catapult_lib.MaxSpeed(params=kCatapult,
+            catapult_lib.MaxSpeed(params=kCatapultWithBall,
                                   U=U,
                                   final_position=math.pi / 2.0))
 
         CatapultProblem()
+
+        catapult_lib.PlotStep(params=kCatapultWithBall,
+                              R=numpy.matrix([[1.0], [0.0]]))
+        catapult_lib.PlotKick(params=kCatapultWithBall,
+                              R=numpy.matrix([[1.0], [0.0]]))
         return 0
 
-        catapult_lib.PlotShot(kCatapult, U, final_position=math.pi / 4.0)
+        catapult_lib.PlotShot(kCatapultWithBall,
+                              U,
+                              final_position=math.pi / 4.0)
 
         gs = []
         speed = []
         for i in numpy.linspace(0.01, 0.15, 150):
-            kCatapult.G = i
-            gs.append(kCatapult.G)
+            kCatapultWithBall.G = i
+            gs.append(kCatapultWithBall.G)
             speed.append(
-                catapult_lib.MaxSpeed(params=kCatapult,
+                catapult_lib.MaxSpeed(params=kCatapultWithBall,
                                       U=U,
                                       final_position=math.pi / 2.0))
         pylab.plot(gs, speed, label="max_speed")
@@ -359,10 +382,11 @@
         )
     else:
         namespaces = ['y2022', 'control_loops', 'superstructure', 'catapult']
-        catapult_lib.WriteCatapult(kCatapult, argv[1:3], argv[3:5], namespaces)
+        catapult_lib.WriteCatapult([kCatapultWithBall, kCatapultEmpty], argv[1:3], argv[3:5], namespaces)
     return 0
 
 
 if __name__ == '__main__':
     argv = FLAGS(sys.argv)
+    glog.init()
     sys.exit(main(argv))
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(