Automatically generate the drivetrain constants.

Change-Id: Ib96f7d55734d44d8ac9d918ec90edf3b1fb46cd3
diff --git a/y2014/control_loops/python/drivetrain.py b/y2014/control_loops/python/drivetrain.py
index ac05a57..5a2e478 100755
--- a/y2014/control_loops/python/drivetrain.py
+++ b/y2014/control_loops/python/drivetrain.py
@@ -1,7 +1,7 @@
 #!/usr/bin/python
 
-import control_loop
-import controls
+from frc971.control_loops.python import control_loop
+from frc971.control_loops.python import controls
 import numpy
 import sys
 from matplotlib import pylab
@@ -71,7 +71,7 @@
     # Radius of the wheels, in meters.
     self.r = .04445
     # Resistance of the motor, divided by the number of motors.
-    self.R = 12.0 / self.stall_current / 4
+    self.R = 12.0 / self.stall_current / 2
     # Motor velocity constant
     self.Kv = ((self.free_speed / 60.0 * 2.0 * numpy.pi) /
                (12.0 - self.R * self.free_current))
@@ -120,7 +120,6 @@
     self.D = numpy.matrix([[0, 0],
                            [0, 0]])
 
-    #print "THE NUMBER I WANT" + str(numpy.linalg.inv(self.A_continuous) * -self.B_continuous * numpy.matrix([[12.0], [12.0]]))
     self.A, self.B = self.ContinuousToDiscrete(
         self.A_continuous, self.B_continuous, self.dt)
 
@@ -128,7 +127,7 @@
     self.hp = 0.65
     self.lp = 0.83
     self.PlaceControllerPoles([self.hp, self.lp, self.hp, self.lp])
-    print self.K
+    #print self.K
     q_pos = 0.07
     q_vel = 1.0
     self.Q = numpy.matrix([[(1.0 / (q_pos ** 2.0)), 0.0, 0.0, 0.0],
@@ -139,10 +138,10 @@
     self.R = numpy.matrix([[(1.0 / (12.0 ** 2.0)), 0.0],
                            [0.0, (1.0 / (12.0 ** 2.0))]])
     self.K = controls.dlqr(self.A, self.B, self.Q, self.R)
-    print self.A
-    print self.B
-    print self.K
-    print numpy.linalg.eig(self.A - self.B * self.K)[0]
+    #print self.A
+    #print self.B
+    #print self.K
+    #print numpy.linalg.eig(self.A - self.B * self.K)[0]
 
     self.hlp = 0.3
     self.llp = 0.4
@@ -152,6 +151,66 @@
     self.U_min = numpy.matrix([[-12.0], [-12.0]])
     self.InitializeState()
 
+
+class KFDrivetrain(Drivetrain):
+  def __init__(self, name="KFDrivetrain", left_low=True, right_low=True):
+    super(KFDrivetrain, self).__init__(name, left_low, right_low)
+
+    self.unaugmented_A_continuous = self.A_continuous
+    self.unaugmented_B_continuous = self.B_continuous
+
+    # The states are
+    # The practical voltage applied to the wheels is
+    #   V_left = U_left + left_voltage_error
+    #
+    # [left position, left velocity, right position, right velocity,
+    #  left voltage error, right voltage error, angular_error]
+    self.A_continuous = numpy.matrix(numpy.zeros((7, 7)))
+    self.B_continuous = numpy.matrix(numpy.zeros((7, 2)))
+    self.A_continuous[0:4,0:4] = self.unaugmented_A_continuous
+    self.A_continuous[0:4,4:6] = self.unaugmented_B_continuous
+    self.B_continuous[0:4,0:2] = self.unaugmented_B_continuous
+
+    self.A, self.B = self.ContinuousToDiscrete(
+        self.A_continuous, self.B_continuous, self.dt)
+
+    self.C = numpy.matrix([[1, 0, 0, 0, 0, 0, self.rb],
+                           [0, 0, 1, 0, 0, 0, -self.rb],
+                           [0, -2.0 / self.rb, 0, 2.0 / self.rb, 0, 0, 0]])
+
+    self.D = numpy.matrix([[0, 0],
+                           [0, 0],
+                           [0, 0]])
+
+    q_pos = 0.08
+    q_vel = 0.40
+    q_voltage = 6.0
+    q_gyro = 0.1
+
+    self.Q = numpy.matrix([[(q_pos ** 2.0), 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                           [0.0, (q_vel ** 2.0), 0.0, 0.0, 0.0, 0.0, 0.0],
+                           [0.0, 0.0, (q_pos ** 2.0), 0.0, 0.0, 0.0, 0.0],
+                           [0.0, 0.0, 0.0, (q_vel ** 2.0), 0.0, 0.0, 0.0],
+                           [0.0, 0.0, 0.0, 0.0, (q_voltage ** 2.0), 0.0, 0.0],
+                           [0.0, 0.0, 0.0, 0.0, 0.0, (q_voltage ** 2.0), 0.0],
+                           [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, (q_gyro ** 2.0)]])
+
+    r_pos = 0.05
+    r_gyro = 0.001
+    self.R = numpy.matrix([[(r_pos ** 2.0), 0.0, 0.0],
+                           [0.0, (r_pos ** 2.0), 0.0],
+                           [0.0, 0.0, (r_gyro ** 2.0)]])
+
+    # Solving for kf gains.
+    self.KalmanGain, self.Q_steady = controls.kalman(
+        A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
+
+    self.L = self.A * self.KalmanGain
+
+    # We need a nothing controller for the autogen code to be happy.
+    self.K = numpy.matrix(numpy.zeros((self.B.shape[1], self.A.shape[0])))
+
+
 def main(argv):
   # Simulate the response of the system to a step input.
   drivetrain = Drivetrain()
@@ -218,22 +277,37 @@
   #pylab.show()
 
   # Write the generated constants out to a file.
-  print "Output one"
   drivetrain_low_low = Drivetrain(name="DrivetrainLowLow", left_low=True, right_low=True)
   drivetrain_low_high = Drivetrain(name="DrivetrainLowHigh", left_low=True, right_low=False)
   drivetrain_high_low = Drivetrain(name="DrivetrainHighLow", left_low=False, right_low=True)
   drivetrain_high_high = Drivetrain(name="DrivetrainHighHigh", left_low=False, right_low=False)
 
+  kf_drivetrain_low_low = KFDrivetrain(name="KFDrivetrainLowLow", left_low=True, right_low=True)
+  kf_drivetrain_low_high = KFDrivetrain(name="KFDrivetrainLowHigh", left_low=True, right_low=False)
+  kf_drivetrain_high_low = KFDrivetrain(name="KFDrivetrainHighLow", left_low=False, right_low=True)
+  kf_drivetrain_high_high = KFDrivetrain(name="KFDrivetrainHighHigh", left_low=False, right_low=False)
+
   if len(argv) != 5:
     print "Expected .h file name and .cc file name"
   else:
+    namespaces = ['y2014', 'control_loops', 'drivetrain']
     dog_loop_writer = control_loop.ControlLoopWriter(
         "Drivetrain", [drivetrain_low_low, drivetrain_low_high,
-                       drivetrain_high_low, drivetrain_high_high])
+                       drivetrain_high_low, drivetrain_high_high],
+        namespaces = namespaces)
     if argv[1][-3:] == '.cc':
       dog_loop_writer.Write(argv[2], argv[1])
     else:
       dog_loop_writer.Write(argv[1], argv[2])
 
+    kf_loop_writer = control_loop.ControlLoopWriter(
+        "KFDrivetrain", [kf_drivetrain_low_low, kf_drivetrain_low_high,
+                       kf_drivetrain_high_low, kf_drivetrain_high_high],
+        namespaces = namespaces)
+    if argv[3][-3:] == '.cc':
+      kf_loop_writer.Write(argv[4], argv[3])
+    else:
+      kf_loop_writer.Write(argv[3], argv[4])
+
 if __name__ == '__main__':
   sys.exit(main(sys.argv))