Add osqp based python solver

This gets us in a form we can solve quickly on the roboRIO.  Now, all we
need to do is to convert it to C++.

Change-Id: I8416d9f89274d4d288402d3ffd676ed22efc7017
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/y2022/control_loops/python/catapult.py b/y2022/control_loops/python/catapult.py
index ae6079a..491286e 100755
--- a/y2022/control_loops/python/catapult.py
+++ b/y2022/control_loops/python/catapult.py
@@ -3,6 +3,7 @@
 from frc971.control_loops.python import control_loop
 from frc971.control_loops.python import controls
 import numpy
+import osqp
 import math
 import scipy.optimize
 import sys
@@ -49,7 +50,7 @@
 print("J", J)
 
 kCatapult = catapult_lib.CatapultParams(
-    name='Finisher',
+    name='Catapult',
     motor=AddResistance(control_loop.NMotor(control_loop.Falcon(), 2), 0.03),
     G=G,
     J=J,
@@ -61,8 +62,6 @@
     controller_poles=[.93],
     dt=0.00505)
 
-catapult = catapult_lib.Catapult(kCatapult)
-
 # Ideas for adjusting the cost function:
 #
 # Penalize battery current?
@@ -76,8 +75,86 @@
 # We really want our cost function to be robust so that we can tolerate the
 # battery not delivering like we want at the end.
 
+def quadratic_cost(catapult, X_initial, X_final, horizon):
+    Q_final = numpy.matrix([[10000.0, 0.0], [0.0, 10000.0]])
 
-def mpc_cost(X_initial, X_final, u_matrix):
+    As = numpy.vstack([catapult.A**(n + 1) for n in range(0, horizon)])
+    Af = catapult.A**horizon
+
+    Bs = numpy.matrix(numpy.zeros((2 * horizon, horizon)))
+    for n in range(0, horizon):
+        for m in range(0, n + 1):
+            Bs[n * 2:(n * 2) + 2, m] = catapult.A**(n - m) * catapult.B
+
+    Bf = Bs[horizon * 2 - 2:, :]
+
+    P_final = 2.0 * Bf.transpose() * Q_final * Bf
+    q_final = (2.0 * (Af * X_initial - X_final).transpose() * Q_final *
+         Bf).transpose()
+
+    constant_final = (Af * X_initial - X_final).transpose() * Q_final * (
+        Af * X_initial - X_final)
+
+    m = numpy.matrix([[catapult.A[1, 1] ** (n + 1)] for n in range(horizon)])
+    M = Bs[1:horizon * 2:2, :]
+
+    W = numpy.matrix(numpy.identity(horizon) - numpy.eye(horizon, horizon, -1)) / catapult.dt
+    w = -numpy.matrix(numpy.eye(horizon, 1, 0)) / catapult.dt
+
+
+    Pi = numpy.diag([
+        (0.01 ** 2.0) + (0.02 * max(0.0, 20 - (horizon - n)) / 20.0) ** 2.0 for n in range(horizon)
+    ])
+
+    P_accel = 2.0 * M.transpose() * W.transpose() * Pi * W * M
+    q_accel = 2.0 * (((W * m + w) * X_initial[1, 0]).transpose() * Pi * W * M).transpose()
+    constant_accel = ((W * m + w) * X_initial[1, 0]).transpose() * Pi * (
+        (W * m + w) * X_initial[1, 0])
+
+    return ((P_accel + P_final), (q_accel + q_final), (constant_accel + constant_final))
+
+
+def new_cost(catapult, X_initial, X_final, u):
+    u_matrix = numpy.matrix(u).transpose()
+    Q_final = numpy.matrix([[10000.0, 0.0], [0.0, 10000.0]])
+
+    As = numpy.vstack([catapult.A**(n + 1) for n in range(0, len(u))])
+    Af = catapult.A**len(u)
+
+    Bs = numpy.matrix(numpy.zeros((2 * len(u), len(u))))
+    for n in range(0, len(u)):
+        for m in range(0, n + 1):
+            Bs[n * 2:(n * 2) + 2, m] = catapult.A**(n - m) * catapult.B
+
+    Bf = Bs[len(u) * 2 - 2:, :]
+
+    P_final = 2.0 * Bf.transpose() * Q_final * Bf
+    q_final = (2.0 * (Af * X_initial - X_final).transpose() * Q_final *
+         Bf).transpose()
+
+    constant_final = (Af * X_initial - X_final).transpose() * Q_final * (
+        Af * X_initial - X_final)
+
+    m = numpy.matrix([[catapult.A[1, 1] ** (n + 1)] for n in range(len(u))])
+    M = Bs[1:len(u) * 2:2, :]
+
+    W = numpy.matrix(numpy.identity(len(u)) - numpy.eye(len(u), len(u), -1)) / catapult.dt
+    w = -numpy.matrix(numpy.eye(len(u), 1, 0)) * X_initial[1, 0] / catapult.dt
+
+    accel = W * (M * u_matrix + m * X_initial[1, 0]) + w
+
+    Pi = numpy.diag([
+        (0.01 ** 2.0) + (0.02 * max(0.0, 20 - (len(u) - n)) / 20.0) ** 2.0 for n in range(len(u))
+    ])
+
+    P_accel = 2.0 * M.transpose() * W.transpose() * Pi * W * M
+    q_accel = 2.0 * ((W * m * X_initial[1, 0] + w).transpose() * Pi * W * M).transpose()
+    constant_accel = (W * m * X_initial[1, 0] +
+                      w).transpose() * Pi * (W * m * X_initial[1, 0] + w)
+
+
+def mpc_cost(catapult, X_initial, X_final, u_matrix):
+
     X = X_initial.copy()
     cost = 0.0
     last_u = u_matrix[0]
@@ -118,7 +195,7 @@
     return cost
 
 
-def SolveCatapult(X_initial, X_final, u):
+def SolveCatapult(catapult, X_initial, X_final, u):
     """ Solves for the optimal action given a seed, state, and target """
     def vbat_constraint(z, i):
         return 12.0 - z[i]
@@ -126,10 +203,31 @@
     def forward(z, i):
         return z[i]
 
-    def mpc_cost2(u_matrix):
-        return mpc_cost(X_initial, X_final, u_matrix)
+    P, q, c = quadratic_cost(catapult, X_initial, X_final, len(u))
 
-    constraints = [{
+
+    def mpc_cost2(u_solver):
+        u_matrix = numpy.matrix(u_solver).transpose()
+        cost = mpc_cost(catapult, X_initial, X_final, u_solver)
+        return cost
+
+
+    def mpc_cost3(u_solver):
+        u_matrix = numpy.matrix(u_solver).transpose()
+        return (0.5 * u_matrix.transpose() * P * u_matrix +
+                q.transpose() * u_matrix + c)[0, 0]
+
+    # If we provide scipy with the analytical jacobian and hessian, it solves
+    # more accurately and a *lot* faster.
+    def jacobian(u):
+        u_matrix = numpy.matrix(u).transpose()
+        return numpy.array(P * u_matrix + q)
+
+    def hessian(u):
+        return numpy.array(P)
+
+    constraints = []
+    constraints += [{
         'type': 'ineq',
         'fun': vbat_constraint,
         'args': (i, )
@@ -141,9 +239,12 @@
         'args': (i, )
     } for i in numpy.arange(len(u))]
 
-    result = scipy.optimize.minimize(mpc_cost2,
+    result = scipy.optimize.minimize(mpc_cost3,
                                      u,
+                                     jac=jacobian,
+                                     hess=hessian,
                                      method='SLSQP',
+                                     tol=1e-12,
                                      constraints=constraints)
     print(result)
 
@@ -171,7 +272,8 @@
     # Iteratively solve our MPC and simulate it.
     u_samples = []
     for i in range(kHorizon):
-        u_horizon = SolveCatapult(X, X_final, u)
+        u_horizon = SolveCatapult(c, X, X_final, u)
+
         u_samples.append(u_horizon[0])
         v_prior = X[1, 0]
         X = c.A * X + c.B * numpy.matrix([[u_horizon[0]]])
@@ -200,15 +302,43 @@
 def main(argv):
     # Do all our math with a lower voltage so we have headroom.
     U = numpy.matrix([[9.0]])
-    print(
-        "For G:", G, " max speed ",
-        catapult_lib.MaxSpeed(params=kCatapult,
-                              U=U,
-                              final_position=math.pi / 2.0))
 
-    CatapultProblem()
+    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(
+            "For G:", G, " max speed ",
+            catapult_lib.MaxSpeed(params=kCatapult,
+                                  U=U,
+                                  final_position=math.pi / 2.0))
+
+        CatapultProblem()
+        return 0
+
         catapult_lib.PlotShot(kCatapult, U, final_position=math.pi / 4.0)
 
         gs = []
@@ -222,7 +352,15 @@
                                       final_position=math.pi / 2.0))
         pylab.plot(gs, speed, label="max_speed")
         pylab.show()
-        return 0
+
+    if len(argv) != 5:
+        glog.fatal(
+            'Expected .h file name and .cc file name for the catapult and integral catapult.'
+        )
+    else:
+        namespaces = ['y2022', 'control_loops', 'superstructure', 'catapult']
+        catapult_lib.WriteCatapult(kCatapult, argv[1:3], argv[3:5], namespaces)
+    return 0
 
 
 if __name__ == '__main__':