Add saturation for current loops

Signed-off-by: Ravago Jones <ravagojones@gmail.com>
Change-Id: If80e6b964453f748ab8846c9c5f46bc94c2058ac
diff --git a/frc971/control_loops/drivetrain/BUILD b/frc971/control_loops/drivetrain/BUILD
index 86c1ec0..42504e1 100644
--- a/frc971/control_loops/drivetrain/BUILD
+++ b/frc971/control_loops/drivetrain/BUILD
@@ -419,7 +419,6 @@
     ],
     target_compatible_with = ["@platforms//os:linux"],
     deps = [
-        ":down_estimator",
         ":drivetrain_goal_fbs",
         ":drivetrain_output_fbs",
         ":drivetrain_position_fbs",
diff --git a/frc971/control_loops/drivetrain/drivetrain.cc b/frc971/control_loops/drivetrain/drivetrain.cc
index 3b16273..be4cf6f 100644
--- a/frc971/control_loops/drivetrain/drivetrain.cc
+++ b/frc971/control_loops/drivetrain/drivetrain.cc
@@ -9,7 +9,6 @@
 #include "Eigen/Dense"
 
 #include "aos/logging/logging.h"
-#include "frc971/control_loops/drivetrain/down_estimator.h"
 #include "frc971/control_loops/drivetrain/drivetrain_config.h"
 #include "frc971/control_loops/drivetrain/drivetrain_goal_generated.h"
 #include "frc971/control_loops/drivetrain/drivetrain_output_generated.h"
diff --git a/frc971/control_loops/hybrid_state_feedback_loop.h b/frc971/control_loops/hybrid_state_feedback_loop.h
index 9d463d8..d8dbe72 100644
--- a/frc971/control_loops/hybrid_state_feedback_loop.h
+++ b/frc971/control_loops/hybrid_state_feedback_loop.h
@@ -30,6 +30,8 @@
         D(other.D),
         U_min(other.U_min),
         U_max(other.U_max),
+        U_limit_coefficient(other.U_limit_coefficient),
+        U_limit_constant(other.U_limit_constant),
         delayed_u(other.delayed_u) {}
 
   StateFeedbackHybridPlantCoefficients(
@@ -40,13 +42,19 @@
       const Eigen::Matrix<Scalar, number_of_outputs, number_of_states> &C,
       const Eigen::Matrix<Scalar, number_of_outputs, number_of_inputs> &D,
       const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_max,
-      const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_min, size_t delayed_u)
+      const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_min,
+      const Eigen::Matrix<Scalar, number_of_inputs, number_of_states>
+          &U_limit_coefficient,
+      const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_limit_constant,
+      size_t delayed_u)
       : A_continuous(A_continuous),
         B_continuous(B_continuous),
         C(C),
         D(D),
         U_min(U_min),
         U_max(U_max),
+        U_limit_coefficient(U_limit_coefficient),
+        U_limit_constant(U_limit_constant),
         delayed_u(delayed_u) {}
 
   const Eigen::Matrix<Scalar, number_of_states, number_of_states> A_continuous;
@@ -55,6 +63,9 @@
   const Eigen::Matrix<Scalar, number_of_outputs, number_of_inputs> D;
   const Eigen::Matrix<Scalar, number_of_inputs, 1> U_min;
   const Eigen::Matrix<Scalar, number_of_inputs, 1> U_max;
+  const Eigen::Matrix<Scalar, number_of_inputs, number_of_states>
+      U_limit_coefficient;
+  const Eigen::Matrix<Scalar, number_of_inputs, 1> U_limit_constant;
 
   const size_t delayed_u;
 };
@@ -109,6 +120,19 @@
     return coefficients().U_max;
   }
   Scalar U_max(int i, int j) const { return U_max()(i, j); }
+  const Eigen::Matrix<Scalar, number_of_inputs, number_of_states>
+      &U_limit_coefficient() const {
+    return coefficients().U_limit_coefficient;
+  }
+  Scalar U_limit_coefficient(int i, int j) const {
+    return U_limit_coefficient()(i, j);
+  }
+  const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_limit_constant() const {
+    return coefficients().U_limit_constant;
+  }
+  Scalar U_limit_constant(int i, int j = 0) const {
+    return U_limit_constant()(i, j);
+  }
 
   const Eigen::Matrix<Scalar, number_of_states, 1> &X() const { return X_; }
   Scalar X(int i) const { return X()(i); }
diff --git a/frc971/control_loops/hybrid_state_feedback_loop_test.cc b/frc971/control_loops/hybrid_state_feedback_loop_test.cc
index 76e2838..b9a0307 100644
--- a/frc971/control_loops/hybrid_state_feedback_loop_test.cc
+++ b/frc971/control_loops/hybrid_state_feedback_loop_test.cc
@@ -20,6 +20,12 @@
   U_max(0, 0) = 12.0;
   Eigen::Matrix<double, 1, 1> U_min;
   U_min(0, 0) = -12.0;
+  Eigen::Matrix<double, 1, 3> U_limit_coefficient;
+  U_limit_coefficient(0, 0) = 0.0;
+  U_limit_coefficient(0, 1) = 0.0;
+  U_limit_coefficient(0, 2) = 0.0;
+  Eigen::Matrix<double, 1, 1> U_limit_constant;
+  U_limit_constant(0, 0) = 12.0;
   Eigen::Matrix<double, 3, 3> A_continuous;
   A_continuous(0, 0) = 0.0;
   A_continuous(0, 1) = 1.0;
@@ -35,7 +41,8 @@
   B_continuous(1, 0) = 443.75;
   B_continuous(2, 0) = 0.0;
   return StateFeedbackHybridPlantCoefficients<3, 1, 1>(
-      A_continuous, B_continuous, C, D, U_max, U_min, false);
+      A_continuous, B_continuous, C, D, U_max, U_min, U_limit_coefficient,
+      U_limit_constant, false);
 }
 
 StateFeedbackControllerCoefficients<3, 1, 1>
@@ -129,6 +136,8 @@
       Eigen::Matrix<double, 7, 4>::Identity(),
       Eigen::Matrix<double, 4, 1>::Constant(1),
       Eigen::Matrix<double, 4, 1>::Constant(-1),
+      Eigen::Matrix<double, 4, 2>::Zero(),
+      Eigen::Matrix<double, 4, 1>::Constant(0),
       frc971::controls::kLoopFrequency, false);
 
   // Build a plant.
diff --git a/frc971/control_loops/python/control_loop.py b/frc971/control_loops/python/control_loop.py
index e836d1a..7e54bdf 100644
--- a/frc971/control_loops/python/control_loop.py
+++ b/frc971/control_loops/python/control_loop.py
@@ -400,20 +400,39 @@
             (plant_coefficient_type, self._name)
         ]
 
+        num_states = self.A.shape[0]
+        num_inputs = self.B.shape[1]
+        num_outputs = self.C.shape[0]
+
         ans.append(self._DumpMatrix('C', self.C, scalar_type))
         ans.append(self._DumpMatrix('D', self.D, scalar_type))
         ans.append(self._DumpMatrix('U_max', self.U_max, scalar_type))
         ans.append(self._DumpMatrix('U_min', self.U_min, scalar_type))
 
+        if not hasattr(self, 'U_limit_coefficient'):
+            self.U_limit_coefficient = numpy.matrix(
+                numpy.zeros((num_inputs, num_states)))
+
+        if not hasattr(self, 'U_limit_constant'):
+            self.U_limit_constant = self.U_max
+
+        ans.append(
+            self._DumpMatrix('U_limit_coefficient', self.U_limit_coefficient,
+                             scalar_type))
+        ans.append(
+            self._DumpMatrix('U_limit_constant', self.U_limit_constant,
+                             scalar_type))
+
         delayed_u_string = str(self.delayed_u)
         if plant_coefficient_type.startswith('StateFeedbackPlant'):
             ans.append(self._DumpMatrix('A', self.A, scalar_type))
             ans.append(self._DumpMatrix('B', self.B, scalar_type))
             ans.append('  const std::chrono::nanoseconds dt(%d);\n' %
                        (self.dt * 1e9))
-            ans.append('  return %s'
-                       '(A, B, C, D, U_max, U_min, dt, %s);\n' %
-                       (plant_coefficient_type, delayed_u_string))
+            ans.append(
+                '  return %s'
+                '(A, B, C, D, U_max, U_min, U_limit_coefficient, U_limit_constant, dt, %s);\n'
+                % (plant_coefficient_type, delayed_u_string))
         elif plant_coefficient_type.startswith('StateFeedbackHybridPlant'):
             ans.append(
                 self._DumpMatrix('A_continuous', self.A_continuous,
@@ -423,8 +442,8 @@
                                  scalar_type))
             ans.append(
                 '  return %s'
-                '(A_continuous, B_continuous, C, D, U_max, U_min, %s);\n' %
-                (plant_coefficient_type, delayed_u_string))
+                '(A_continuous, B_continuous, C, D, U_max, U_min, U_limit_coefficient, U_limit_constant, %s);\n'
+                % (plant_coefficient_type, delayed_u_string))
         else:
             glog.fatal('Unsupported plant type %s', plant_coefficient_type)
 
diff --git a/frc971/control_loops/state_feedback_loop.h b/frc971/control_loops/state_feedback_loop.h
index 65e0ba1..295beca 100644
--- a/frc971/control_loops/state_feedback_loop.h
+++ b/frc971/control_loops/state_feedback_loop.h
@@ -42,6 +42,8 @@
         D(other.D),
         U_min(other.U_min),
         U_max(other.U_max),
+        U_limit_coefficient(other.U_limit_coefficient),
+        U_limit_constant(other.U_limit_constant),
         dt(other.dt),
         delayed_u(other.delayed_u) {}
 
@@ -52,6 +54,9 @@
       const Eigen::Matrix<Scalar, number_of_outputs, number_of_inputs> &D,
       const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_max,
       const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_min,
+      const Eigen::Matrix<Scalar, number_of_inputs, number_of_states>
+          &U_limit_coefficient,
+      const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_limit_constant,
       const std::chrono::nanoseconds dt, size_t delayed_u)
       : A(A),
         B(B),
@@ -59,6 +64,8 @@
         D(D),
         U_min(U_min),
         U_max(U_max),
+        U_limit_coefficient(U_limit_coefficient),
+        U_limit_constant(U_limit_constant),
         dt(dt),
         delayed_u(delayed_u) {}
 
@@ -68,6 +75,9 @@
   const Eigen::Matrix<Scalar, number_of_outputs, number_of_inputs> D;
   const Eigen::Matrix<Scalar, number_of_inputs, 1> U_min;
   const Eigen::Matrix<Scalar, number_of_inputs, 1> U_max;
+  const Eigen::Matrix<Scalar, number_of_inputs, number_of_states>
+      U_limit_coefficient;
+  const Eigen::Matrix<Scalar, number_of_inputs, 1> U_limit_constant;
   const std::chrono::nanoseconds dt;
 
   // If true, this adds a single cycle output delay model to the plant.  This is
@@ -134,6 +144,19 @@
     return coefficients().U_max;
   }
   Scalar U_max(int i, int j = 0) const { return U_max()(i, j); }
+  const Eigen::Matrix<Scalar, number_of_inputs, number_of_states>
+      &U_limit_coefficient() const {
+    return coefficients().U_limit_coefficient;
+  }
+  Scalar U_limit_coefficient(int i, int j) const {
+    return U_limit_coefficient()(i, j);
+  }
+  const Eigen::Matrix<Scalar, number_of_inputs, 1> &U_limit_constant() const {
+    return coefficients().U_limit_constant;
+  }
+  Scalar U_limit_constant(int i, int j = 0) const {
+    return U_limit_constant()(i, j);
+  }
 
   const std::chrono::nanoseconds dt() const { return coefficients().dt; }
 
@@ -549,11 +572,24 @@
   // If U is outside the hardware range, limit it before the plant tries to use
   // it.
   virtual void CapU() {
+    // TODO(Ravago): this runs before the state update step, so it's limiting
+    // the future control based on the old state. This gets even worse when we
+    // have delayed_u.
+    Eigen::Matrix<Scalar, number_of_inputs, 1> U_max_computed =
+        plant().U_limit_coefficient() * plant().X() +
+        plant().U_limit_constant();
+    Eigen::Matrix<Scalar, number_of_inputs, 1> U_min_computed =
+        plant().U_limit_coefficient() * plant().X() -
+        plant().U_limit_constant();
+
+    U_max_computed = U_max_computed.cwiseMin(plant().U_max());
+    U_min_computed = U_min_computed.cwiseMax(plant().U_min());
+
     for (int i = 0; i < kNumInputs; ++i) {
-      if (U(i, 0) > plant().U_max(i, 0)) {
-        U_(i, 0) = plant().U_max(i, 0);
-      } else if (U(i, 0) < plant().U_min(i, 0)) {
-        U_(i, 0) = plant().U_min(i, 0);
+      if (U(i, 0) > U_max_computed(i, 0)) {
+        U_(i, 0) = U_max_computed(i, 0);
+      } else if (U(i, 0) < U_min_computed(i, 0)) {
+        U_(i, 0) = U_min_computed(i, 0);
       }
     }
   }