Add single speed robot option to DrivetrainParams

It just forces high and low to have the same ratio, but then we can
change this in the future easier.

Change-Id: I7011703d5f0800638bc285d9da75b03342bc602c
diff --git a/frc971/control_loops/python/drivetrain.py b/frc971/control_loops/python/drivetrain.py
index b1244f3..3f9123f 100644
--- a/frc971/control_loops/python/drivetrain.py
+++ b/frc971/control_loops/python/drivetrain.py
@@ -13,12 +13,15 @@
                  mass,
                  robot_radius,
                  wheel_radius,
-                 G_high,
-                 G_low,
-                 q_pos_low=0.12,
-                 q_pos_high=0.14,
-                 q_vel_low=1.0,
-                 q_vel_high=0.95,
+                 G=None,
+                 G_high=None,
+                 G_low=None,
+                 q_pos=None,
+                 q_pos_low=None,
+                 q_pos_high=None,
+                 q_vel=None,
+                 q_vel_low=None,
+                 q_vel_high=None,
                  efficiency=0.60,
                  has_imu=False,
                  force=False,
@@ -37,11 +40,14 @@
           robot_radius: float, Radius of the robot, in meters (requires tuning by
             hand).
           wheel_radius: float, Radius of the wheels, in meters.
+          G: float, Gear ratio for a single speed.
           G_high: float, Gear ratio for high gear.
           G_low: float, Gear ratio for low gear.
           dt: float, Control loop time step.
+          q_pos: float, q position for a single speed.
           q_pos_low: float, q position low gear.
           q_pos_high: float, q position high gear.
+          q_vel: float, q velocity for a single speed
           q_vel_low: float, q velocity low gear.
           q_vel_high: float, q velocity high gear.
           efficiency: float, gear box effiency.
@@ -54,6 +60,29 @@
           observer_poles: array, An array of poles. (See control_loop.py)
           robot_cg_offset: offset in meters of CG from robot center to left side
         """
+        if G is not None:
+            assert (G_high is None)
+            assert (G_low is None)
+            G_high = G
+            G_low = G
+        assert (G_high is not None)
+        assert (G_low is not None)
+
+        if q_pos is not None:
+            assert (q_pos_low is None)
+            assert (q_pos_high is None)
+            q_pos_low = q_pos
+            q_pos_high = q_pos
+        assert (q_pos_low is not None)
+        assert (q_pos_high is not None)
+
+        if q_vel is not None:
+            assert (q_vel_low is None)
+            assert (q_vel_high is None)
+            q_vel_low = q_vel
+            q_vel_high = q_vel
+        assert (q_vel_low is not None)
+        assert (q_vel_high is not None)
 
         self.J = J
         self.mass = mass
@@ -334,9 +363,9 @@
             A=self.A, B=self.B, C=self.C, Q=self.Q, R=self.R)
 
         if not self.has_imu:
-          self.KalmanGain = numpy.hstack((self.KalmanGain, numpy.matrix(numpy.zeros((7, 1)))))
-          self.C = numpy.vstack((self.C, numpy.matrix(numpy.zeros((1, 7)))))
-          self.D = numpy.vstack((self.D, numpy.matrix(numpy.zeros((1, 2)))))
+            self.KalmanGain = numpy.hstack((self.KalmanGain, numpy.matrix(numpy.zeros((7, 1)))))
+            self.C = numpy.vstack((self.C, numpy.matrix(numpy.zeros((1, 7)))))
+            self.D = numpy.vstack((self.D, numpy.matrix(numpy.zeros((1, 2)))))
 
         self.L = self.A * self.KalmanGain
 
@@ -414,9 +443,9 @@
         drivetrain_params=drivetrain_params)
 
     if isinstance(year_namespace, list):
-      namespaces = year_namespace
+        namespaces = year_namespace
     else:
-      namespaces = [year_namespace, 'control_loops', 'drivetrain']
+        namespaces = [year_namespace, 'control_loops', 'drivetrain']
     dog_loop_writer = control_loop.ControlLoopWriter(
         "Drivetrain", [
             drivetrain_low_low, drivetrain_low_high, drivetrain_high_low,
diff --git a/motors/seems_reasonable/drivetrain.py b/motors/seems_reasonable/drivetrain.py
index dfa2f16..f4a5094 100644
--- a/motors/seems_reasonable/drivetrain.py
+++ b/motors/seems_reasonable/drivetrain.py
@@ -15,8 +15,11 @@
     wheel_radius=0.127 / 2.0 * 120.0 / 118.0,
     G_low=46.0 / 60.0 * 20.0 / 48.0 * 14.0 / 62.0,
     G_high=62.0 / 44.0 * 20.0 / 48.0 * 14.0 / 62.0,
-    controller_poles=[0.82, 0.82],
-)
+    q_pos_low=0.12,
+    q_pos_high=0.14,
+    q_vel_low=1.0,
+    q_vel_high=0.95,
+    controller_poles=[0.82, 0.82])
 
 
 def main(argv):
diff --git a/y2017/control_loops/python/drivetrain.py b/y2017/control_loops/python/drivetrain.py
index 74b114d..50a5a1c 100755
--- a/y2017/control_loops/python/drivetrain.py
+++ b/y2017/control_loops/python/drivetrain.py
@@ -10,29 +10,29 @@
 
 gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
 
-kDrivetrain = drivetrain.DrivetrainParams(J = 6.0,
-                                          mass = 52,
-                                          robot_radius = 0.59055 / 2.0,
-                                          wheel_radius = 0.08255 / 2.0,
-                                          G_high = 11.0 / 60.0,
-                                          G_low = 11.0 / 60.0,
-                                          q_pos_low = 0.12,
-                                          q_pos_high = 0.14,
-                                          q_vel_low = 1.0,
-                                          q_vel_high = 0.95,
-                                          has_imu = False)
+kDrivetrain = drivetrain.DrivetrainParams(
+    J=6.0,
+    mass=52,
+    robot_radius=0.59055 / 2.0,
+    wheel_radius=0.08255 / 2.0,
+    G=11.0 / 60.0,
+    q_pos_low=0.12,
+    q_pos_high=0.14,
+    q_vel_low=1.0,
+    q_vel_high=0.95,
+    has_imu=False)
 
 def main(argv):
-  argv = FLAGS(argv)
-  glog.init()
+    argv = FLAGS(argv)
+    glog.init()
 
-  if FLAGS.plot:
-    drivetrain.PlotDrivetrainMotions(kDrivetrain)
-  elif len(argv) != 5:
-    print "Expected .h file name and .cc file name"
-  else:
-    # Write the generated constants out to a file.
-    drivetrain.WriteDrivetrain(argv[1:3], argv[3:5], 'y2017', kDrivetrain)
+    if FLAGS.plot:
+        drivetrain.PlotDrivetrainMotions(kDrivetrain)
+    elif len(argv) != 5:
+        print "Expected .h file name and .cc file name"
+    else:
+        # Write the generated constants out to a file.
+        drivetrain.WriteDrivetrain(argv[1:3], argv[3:5], 'y2017', kDrivetrain)
 
 if __name__ == '__main__':
-  sys.exit(main(sys.argv))
+    sys.exit(main(sys.argv))