Allow constructing control loops from flatbuffers

The core changes here are to:
* Allow constructing StateFeedbackLoop's from flatbuffers using the
  code in *state_feedback_loop_converters.*
* Add constructors to the single-dof subsystem class to make use of
  this.
* Add code to control_loops.py to generate JSON files with the requisite
  constants (these end up containing identical information to the
  generated .cc files).
* Add interfaces to actually support the new JSON codegen to single-dof
  subsystem classes.
* Convert all of the drivetrains over to generating these. This I mostly
  do so that I can write a test where Iconfirm that the .cc files and
  the JSON files generate exactly the same content.

Change-Id: Iceac48f25ecac96200b7bf992c8f34a15fe6800c
Signed-off-by: James Kuszmaul <jabukuszmaul+collab@gmail.com>
diff --git a/frc971/control_loops/python/angular_system.py b/frc971/control_loops/python/angular_system.py
index c6ffb50..551f3d0 100755
--- a/frc971/control_loops/python/angular_system.py
+++ b/frc971/control_loops/python/angular_system.py
@@ -461,7 +461,8 @@
     loop_writer.AddConstant(
         control_loop.Constant('kFreeSpeed', '%f',
                               angular_systems[0].motor.free_speed))
-    loop_writer.Write(plant_files[0], plant_files[1])
+    loop_writer.Write(plant_files[0], plant_files[1],
+                      None if len(plant_files) < 3 else plant_files[2])
 
     integral_loop_writer = control_loop.ControlLoopWriter(
         'Integral' + name,
@@ -469,4 +470,6 @@
         namespaces=year_namespaces,
         plant_type=plant_type,
         observer_type=observer_type)
-    integral_loop_writer.Write(controller_files[0], controller_files[1])
+    integral_loop_writer.Write(
+        controller_files[0], controller_files[1],
+        None if len(controller_files) < 3 else controller_files[2])
diff --git a/frc971/control_loops/python/angular_system_current.py b/frc971/control_loops/python/angular_system_current.py
index 0dc9595..585f54a 100755
--- a/frc971/control_loops/python/angular_system_current.py
+++ b/frc971/control_loops/python/angular_system_current.py
@@ -501,7 +501,8 @@
     loop_writer.AddConstant(
         control_loop.Constant('kFreeSpeed', '%f',
                               angular_systems[0].motor.free_speed))
-    loop_writer.Write(plant_files[0], plant_files[1])
+    loop_writer.Write(plant_files[0], plant_files[1],
+                      None if len(plant_files) < 3 else plant_files[2])
 
     integral_loop_writer = control_loop.ControlLoopWriter(
         'Integral' + name,
@@ -509,4 +510,6 @@
         namespaces=year_namespaces,
         plant_type=plant_type,
         observer_type=observer_type)
-    integral_loop_writer.Write(controller_files[0], controller_files[1])
+    integral_loop_writer.Write(
+        controller_files[0], controller_files[1],
+        None if len(controller_files) < 3 else controller_files[2])
diff --git a/frc971/control_loops/python/control_loop.py b/frc971/control_loops/python/control_loop.py
index e179193..864a4e4 100644
--- a/frc971/control_loops/python/control_loop.py
+++ b/frc971/control_loops/python/control_loop.py
@@ -1,6 +1,7 @@
 import frc971.control_loops.python.controls as controls
 import numpy
 import os
+import json
 
 
 class Constant(object):
@@ -25,6 +26,16 @@
             (self.comment, typestring, self.name, self.value)
 
 
+def MatrixToJson(matrix):
+    """Returns JSON representation of a numpy matrix."""
+    return {
+        "rows": matrix.shape[0],
+        "cols": matrix.shape[1],
+        "storage_order": "ColMajor",
+        "data": numpy.array(matrix).flatten(order='F').tolist()
+    }
+
+
 class ControlLoopWriter(object):
 
     def __init__(self,
@@ -83,10 +94,12 @@
                 os.path.basename(header_file).upper().replace(
                     '.', '_').replace('/', '_') + '_')
 
-    def Write(self, header_file, cc_file):
+    def Write(self, header_file, cc_file, json_file=None):
         """Writes the loops to the specified files."""
         self.WriteHeader(header_file)
         self.WriteCC(os.path.basename(header_file), cc_file)
+        if json_file is not None:
+            self.WriteJson(json_file)
 
     def _GenericType(self, typename, extra_args=None):
         """Returns a loop template using typename for the type."""
@@ -262,6 +275,19 @@
             fd.write(self._namespace_end)
             fd.write('\n')
 
+    def WriteJson(self, json_file):
+        """Writes a JSON file of the loop constants to the specified json_file."""
+        loops = []
+        for loop in self._loops:
+            loop_json = {}
+            loop_json["plant"] = loop.DumpPlantJson(self._PlantCoeffType())
+            loop_json["controller"] = loop.DumpControllerJson()
+            loop_json["observer"] = loop.DumbObserverJson(
+                self._ObserverCoeffType())
+            loops.append(loop_json)
+        with open(json_file, 'w') as f:
+            f.write(json.dumps(loops))
+
 
 class ControlLoop(object):
 
@@ -389,6 +415,27 @@
         return '%s Make%sPlantCoefficients();\n' % (plant_coefficient_type,
                                                     self._name)
 
+    def DumpPlantJson(self, plant_coefficient_type):
+        result = {
+            "c": MatrixToJson(self.C),
+            "d": MatrixToJson(self.D),
+            "u_min": MatrixToJson(self.U_min),
+            "u_max": MatrixToJson(self.U_max),
+            "u_limit_coefficient": MatrixToJson(self.U_limit_coefficient),
+            "u_limit_constant": MatrixToJson(self.U_limit_constant),
+            "delayed_u": self.delayed_u
+        }
+        if plant_coefficient_type.startswith('StateFeedbackPlant'):
+            result["a"] = MatrixToJson(self.A)
+            result["b"] = MatrixToJson(self.B)
+            result["dt"] = int(self.dt * 1e9)
+        elif plant_coefficient_type.startswith('StateFeedbackHybridPlant'):
+            result["a_continuous"] = MatrixToJson(self.A_continuous)
+            result["b_continuous"] = MatrixToJson(self.B_continuous)
+        else:
+            glog.fatal('Unsupported plant type %s', plant_coefficient_type)
+        return result
+
     def DumpPlant(self, plant_coefficient_type, scalar_type):
         """Writes out a c++ function which will create a PlantCoefficients object.
 
@@ -475,6 +522,10 @@
             num_states, num_inputs, num_outputs, scalar_type,
             self.ControllerFunction())
 
+    def DumpControllerJson(self):
+        result = {"k": MatrixToJson(self.K), "kff": MatrixToJson(self.Kff)}
+        return result
+
     def DumpController(self, scalar_type):
         """Returns a c++ function which will create a Controller object.
 
@@ -511,6 +562,32 @@
         return '%s %s;\n' % (observer_coefficient_type,
                              self.ObserverFunction())
 
+    def GetObserverCoefficients(self):
+        if hasattr(self, 'KalmanGain'):
+            KalmanGain = self.KalmanGain
+            Q = self.Q
+            R = self.R
+        else:
+            KalmanGain = numpy.linalg.inv(self.A) * self.L
+            Q = numpy.zeros(self.A.shape)
+            R = numpy.zeros((self.C.shape[0], self.C.shape[0]))
+        return (KalmanGain, Q, R)
+
+    def DumbObserverJson(self, observer_coefficient_type):
+        result = {"delayed_u": self.delayed_u}
+        if observer_coefficient_type.startswith('StateFeedbackObserver'):
+            KalmanGain, Q, R = self.GetObserverCoefficients()
+            result["kalman_gain"] = MatrixToJson(KalmanGain)
+            result["q"] = MatrixToJson(Q)
+            result["r"] = MatrixToJson(R)
+        elif observer_coefficient_type.startswith('HybridKalman'):
+            result["q_continuous"] = MatrixToJson(self.Q_continuous)
+            result["r_continuous"] = MatrixToJson(self.R_continuous)
+            result["p_steady_state"] = MatrixToJson(self.P_steady_state)
+        else:
+            glog.fatal('Unsupported plant type %s', observer_coefficient_type)
+        return result
+
     def DumpObserver(self, observer_coefficient_type, scalar_type):
         """Returns a c++ function which will create a Observer object.
 
@@ -523,14 +600,7 @@
 
         delayed_u_string = str(self.delayed_u)
         if observer_coefficient_type.startswith('StateFeedbackObserver'):
-            if hasattr(self, 'KalmanGain'):
-                KalmanGain = self.KalmanGain
-                Q = self.Q
-                R = self.R
-            else:
-                KalmanGain = numpy.linalg.inv(self.A) * self.L
-                Q = numpy.zeros(self.A.shape)
-                R = numpy.zeros((self.C.shape[0], self.C.shape[0]))
+            KalmanGain, Q, R = self.GetObserverCoefficients()
             ans.append(self._DumpMatrix('KalmanGain', KalmanGain, scalar_type))
             ans.append(self._DumpMatrix('Q', Q, scalar_type))
             ans.append(self._DumpMatrix('R', R, scalar_type))
diff --git a/frc971/control_loops/python/drivetrain.py b/frc971/control_loops/python/drivetrain.py
index 5ff4404..65b30c0 100644
--- a/frc971/control_loops/python/drivetrain.py
+++ b/frc971/control_loops/python/drivetrain.py
@@ -525,7 +525,8 @@
             "kHighOutputRatio", "%f",
             drivetrain_high_high.G_high * drivetrain_high_high.r))
 
-    dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1])
+    dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1],
+                          drivetrain_files[2])
 
     kf_loop_writer = control_loop.ControlLoopWriter("KFDrivetrain", [
         kf_drivetrain_low_low, kf_drivetrain_low_high, kf_drivetrain_high_low,
@@ -533,7 +534,8 @@
     ],
                                                     namespaces=namespaces,
                                                     scalar_type=scalar_type)
-    kf_loop_writer.Write(kf_drivetrain_files[0], kf_drivetrain_files[1])
+    kf_loop_writer.Write(kf_drivetrain_files[0], kf_drivetrain_files[1],
+                         kf_drivetrain_files[2])
 
 
 def PlotDrivetrainSprint(drivetrain_params):
diff --git a/frc971/control_loops/python/linear_system.py b/frc971/control_loops/python/linear_system.py
index 3bda4e1..e28c422 100755
--- a/frc971/control_loops/python/linear_system.py
+++ b/frc971/control_loops/python/linear_system.py
@@ -422,8 +422,11 @@
                               linear_systems[0].G * linear_systems[0].radius))
     loop_writer.AddConstant(
         control_loop.Constant('kRadius', '%f', linear_systems[0].radius))
-    loop_writer.Write(plant_files[0], plant_files[1])
+    loop_writer.Write(plant_files[0], plant_files[1],
+                      None if len(plant_files) < 3 else plant_files[2])
 
     integral_loop_writer = control_loop.ControlLoopWriter(
         'Integral' + name, integral_linear_systems, namespaces=year_namespaces)
-    integral_loop_writer.Write(controller_files[0], controller_files[1])
+    integral_loop_writer.Write(
+        controller_files[0], controller_files[1],
+        None if len(controller_files) < 3 else controller_files[2])
diff --git a/frc971/control_loops/python/polydrivetrain.py b/frc971/control_loops/python/polydrivetrain.py
index a45d738..c3d47a6 100644
--- a/frc971/control_loops/python/polydrivetrain.py
+++ b/frc971/control_loops/python/polydrivetrain.py
@@ -490,7 +490,8 @@
                                                      namespaces=namespaces,
                                                      scalar_type=scalar_type)
 
-    dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1])
+    dog_loop_writer.Write(drivetrain_files[0], drivetrain_files[1],
+                          drivetrain_files[2])
 
     hybrid_loop_writer = control_loop.ControlLoopWriter(
         "HybridVelocityDrivetrain", [
@@ -504,12 +505,12 @@
         plant_type='StateFeedbackHybridPlant',
         observer_type='HybridKalman')
 
-    hybrid_loop_writer.Write(hybrid_files[0], hybrid_files[1])
+    hybrid_loop_writer.Write(hybrid_files[0], hybrid_files[1], hybrid_files[2])
 
     cim_writer = control_loop.ControlLoopWriter("CIM", [CIM()],
                                                 scalar_type=scalar_type)
 
-    cim_writer.Write(motor_files[0], motor_files[1])
+    cim_writer.Write(motor_files[0], motor_files[1], motor_files[2])
 
 
 def PlotPolyDrivetrainMotions(drivetrain_params):
diff --git a/frc971/control_loops/python/static_zeroing_single_dof_profiled_subsystem_test.py b/frc971/control_loops/python/static_zeroing_single_dof_profiled_subsystem_test.py
index 8080aff..92bed42 100644
--- a/frc971/control_loops/python/static_zeroing_single_dof_profiled_subsystem_test.py
+++ b/frc971/control_loops/python/static_zeroing_single_dof_profiled_subsystem_test.py
@@ -38,13 +38,14 @@
         linear_system.PlotMotion(kIntake, R)
 
     # Write the generated constants out to a file.
-    if len(argv) != 5:
-        glog.fatal('Expected .h file name and .cc file name for the \
+    if len(argv) != 7:
+        glog.fatal(
+            'Expected .h, .cc, and .json filenames and .json file name for the \
             static_zeroing_single_dof_profiled_subsystem_test and integral \
             static_zeroing_single_dof_profiled_subsystem_test.')
     else:
         namespaces = ['frc971', 'control_loops']
-        linear_system.WriteLinearSystem(kIntake, argv[1:3], argv[3:5],
+        linear_system.WriteLinearSystem(kIntake, argv[1:4], argv[4:7],
                                         namespaces)
 
 
diff --git a/frc971/control_loops/python/test_drivetrain/BUILD b/frc971/control_loops/python/test_drivetrain/BUILD
new file mode 100644
index 0000000..63cc2a9
--- /dev/null
+++ b/frc971/control_loops/python/test_drivetrain/BUILD
@@ -0,0 +1,117 @@
+py_binary(
+    name = "drivetrain",
+    srcs = [
+        "drivetrain.py",
+    ],
+    legacy_create_init = False,
+    target_compatible_with = ["@platforms//cpu:x86_64"],
+    deps = [
+        ":python_init",
+        "//frc971/control_loops/python:drivetrain",
+        "@pip//glog",
+        "@pip//python_gflags",
+    ],
+)
+
+py_binary(
+    name = "polydrivetrain",
+    srcs = [
+        "drivetrain.py",
+        "polydrivetrain.py",
+    ],
+    legacy_create_init = False,
+    target_compatible_with = ["@platforms//cpu:x86_64"],
+    deps = [
+        ":python_init",
+        "//frc971/control_loops/python:polydrivetrain",
+        "@pip//glog",
+        "@pip//python_gflags",
+    ],
+)
+
+genrule(
+    name = "genrule_drivetrain",
+    outs = [
+        "drivetrain_dog_motor_plant.h",
+        "drivetrain_dog_motor_plant.cc",
+        "drivetrain_dog_motor_plant.json",
+        "kalman_drivetrain_motor_plant.h",
+        "kalman_drivetrain_motor_plant.cc",
+        "kalman_drivetrain_motor_plant.json",
+    ],
+    cmd = "$(location :drivetrain) $(OUTS)",
+    target_compatible_with = ["@platforms//os:linux"],
+    tools = [
+        ":drivetrain",
+    ],
+)
+
+genrule(
+    name = "genrule_polydrivetrain",
+    outs = [
+        "polydrivetrain_dog_motor_plant.h",
+        "polydrivetrain_dog_motor_plant.cc",
+        "polydrivetrain_dog_motor_plant.json",
+        "polydrivetrain_cim_plant.h",
+        "polydrivetrain_cim_plant.cc",
+        "polydrivetrain_cim_plant.json",
+        "hybrid_velocity_drivetrain.h",
+        "hybrid_velocity_drivetrain.cc",
+        "hybrid_velocity_drivetrain.json",
+    ],
+    cmd = "$(location :polydrivetrain) $(OUTS)",
+    target_compatible_with = ["@platforms//os:linux"],
+    tools = [
+        ":polydrivetrain",
+    ],
+)
+
+cc_library(
+    name = "polydrivetrain_plants",
+    srcs = [
+        "drivetrain_dog_motor_plant.cc",
+        "hybrid_velocity_drivetrain.cc",
+        "kalman_drivetrain_motor_plant.cc",
+        "polydrivetrain_dog_motor_plant.cc",
+    ],
+    hdrs = [
+        "drivetrain_dog_motor_plant.h",
+        "hybrid_velocity_drivetrain.h",
+        "kalman_drivetrain_motor_plant.h",
+        "polydrivetrain_dog_motor_plant.h",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//frc971/control_loops:hybrid_state_feedback_loop",
+        "//frc971/control_loops:state_feedback_loop",
+    ],
+)
+
+cc_test(
+    name = "drivetrain_json_test",
+    srcs = [
+        "drivetrain_json_test.cc",
+    ],
+    data = [
+        "drivetrain_dog_motor_plant.json",
+        "hybrid_velocity_drivetrain.json",
+        "kalman_drivetrain_motor_plant.json",
+        "polydrivetrain_cim_plant.json",
+        "polydrivetrain_dog_motor_plant.json",
+    ],
+    target_compatible_with = ["@platforms//os:linux"],
+    deps = [
+        ":polydrivetrain_plants",
+        "//aos/testing:googletest",
+        "//frc971/control_loops:hybrid_state_feedback_loop_converters",
+    ],
+)
+
+py_library(
+    name = "python_init",
+    srcs = ["__init__.py"],
+    target_compatible_with = ["@platforms//os:linux"],
+    visibility = ["//visibility:public"],
+    deps = ["//y2018/control_loops:python_init"],
+)
diff --git a/frc971/control_loops/python/test_drivetrain/__init__.py b/frc971/control_loops/python/test_drivetrain/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/frc971/control_loops/python/test_drivetrain/__init__.py
diff --git a/frc971/control_loops/python/test_drivetrain/drivetrain.py b/frc971/control_loops/python/test_drivetrain/drivetrain.py
new file mode 100644
index 0000000..0291487
--- /dev/null
+++ b/frc971/control_loops/python/test_drivetrain/drivetrain.py
@@ -0,0 +1,52 @@
+#!/usr/bin/python3
+
+from __future__ import print_function
+from frc971.control_loops.python import drivetrain
+import sys
+
+import gflags
+import glog
+
+FLAGS = gflags.FLAGS
+
+gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
+
+# Values stolen from y2018 (most recent year with shifters).
+kDrivetrain = drivetrain.DrivetrainParams(
+    J=6.0,
+    mass=68.0,
+    robot_radius=0.616 / 2.0,
+    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,
+    q_pos_low=0.12,
+    q_pos_high=0.14,
+    q_vel_low=1.0,
+    q_vel_high=0.95,
+    efficiency=0.70,
+    has_imu=True,
+    force=True,
+    kf_q_voltage=13.0,
+    controller_poles=[0.82, 0.82],
+    robot_cg_offset=0.0,
+)
+
+
+def main(argv):
+    argv = FLAGS(argv)
+    glog.init()
+
+    if FLAGS.plot:
+        drivetrain.PlotDrivetrainMotions(kDrivetrain)
+    elif len(argv) != 7:
+        print("Expected .h, .cc, and .json file names.")
+    else:
+        # Write the generated constants out to a file.
+        drivetrain.WriteDrivetrain(
+            argv[1:4], argv[4:7],
+            ['frc971', 'control_loops', 'python', 'test_drivetrain'],
+            kDrivetrain)
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv))
diff --git a/frc971/control_loops/python/test_drivetrain/drivetrain_json_test.cc b/frc971/control_loops/python/test_drivetrain/drivetrain_json_test.cc
new file mode 100644
index 0000000..b280e99
--- /dev/null
+++ b/frc971/control_loops/python/test_drivetrain/drivetrain_json_test.cc
@@ -0,0 +1,121 @@
+#include "gtest/gtest.h"
+
+#include "aos/util/file.h"
+#include "frc971/control_loops/hybrid_state_feedback_loop_converters.h"
+#include "frc971/control_loops/python/test_drivetrain/drivetrain_dog_motor_plant.h"
+#include "frc971/control_loops/python/test_drivetrain/hybrid_velocity_drivetrain.h"
+#include "frc971/control_loops/python/test_drivetrain/kalman_drivetrain_motor_plant.h"
+#include "frc971/control_loops/python/test_drivetrain/polydrivetrain_dog_motor_plant.h"
+
+namespace frc971::control_loops::drivetrain::testing {
+
+class DrivetrainJsonTest : public ::testing::Test {
+ protected:
+  DrivetrainJsonTest() {}
+
+  aos::FlatbufferDetachedBuffer<fbs::StateFeedbackLoopCoefficientsVector>
+  ReadCoefficients(std::string_view file) {
+    const std::string json =
+        "{\"loops\": " + aos::util::ReadFileToStringOrDie(file) + "}";
+    return aos::JsonToFlatbuffer<fbs::StateFeedbackLoopCoefficientsVector>(
+        json);
+  }
+  aos::FlatbufferDetachedBuffer<fbs::StateFeedbackHybridLoopCoefficientsVector>
+  ReadHybridCoefficients(std::string_view file) {
+    const std::string json =
+        "{\"loops\": " + aos::util::ReadFileToStringOrDie(file) + "}";
+    return aos::JsonToFlatbuffer<
+        fbs::StateFeedbackHybridLoopCoefficientsVector>(json);
+  }
+};
+
+TEST_F(DrivetrainJsonTest, DrivetrainLoop) {
+  StateFeedbackLoop<4, 2, 2> made_loop =
+      python::test_drivetrain::MakeDrivetrainLoop();
+  auto coefs = ReadCoefficients(
+      "frc971/control_loops/python/test_drivetrain/"
+      "drivetrain_dog_motor_plant.json");
+
+  StateFeedbackLoop<4, 2, 2> json_loop =
+      MakeStateFeedbackLoop<4, 2, 2>(*CHECK_NOTNULL(coefs.message().loops()));
+  for (size_t index = 0; index < 4; ++index) {
+    ASSERT_TRUE(coefs.span().size() > 0u);
+    made_loop.set_index(index);
+    json_loop.set_index(index);
+#define COMPARE(matrix)                              \
+  EXPECT_EQ(json_loop.plant().coefficients().matrix, \
+            made_loop.plant().coefficients().matrix);
+    COMPARE(A);
+    COMPARE(B);
+    COMPARE(C);
+    COMPARE(D);
+    COMPARE(U_min);
+    COMPARE(U_max);
+    COMPARE(U_limit_coefficient);
+    COMPARE(U_limit_constant);
+    COMPARE(dt);
+    COMPARE(delayed_u);
+#undef COMPARE
+#define COMPARE(matrix)                                   \
+  EXPECT_EQ(json_loop.controller().coefficients().matrix, \
+            made_loop.controller().coefficients().matrix);
+    COMPARE(K);
+    COMPARE(Kff);
+#undef COMPARE
+#define COMPARE(matrix)                                 \
+  EXPECT_EQ(json_loop.observer().coefficients().matrix, \
+            made_loop.observer().coefficients().matrix);
+    COMPARE(KalmanGain);
+    COMPARE(Q);
+    COMPARE(R);
+    COMPARE(delayed_u);
+#undef COMPARE
+  }
+}
+
+typedef StateFeedbackLoop<2, 2, 2, double, StateFeedbackHybridPlant<2, 2, 2>,
+                          HybridKalman<2, 2, 2>>
+    HybridLoop;
+TEST_F(DrivetrainJsonTest, HybridLoop) {
+  HybridLoop made_loop =
+      python::test_drivetrain::MakeHybridVelocityDrivetrainLoop();
+  auto coefs = ReadHybridCoefficients(
+      "frc971/control_loops/python/test_drivetrain/"
+      "hybrid_velocity_drivetrain.json");
+
+  HybridLoop json_loop = MakeHybridStateFeedbackLoop<2, 2, 2>(
+      *CHECK_NOTNULL(coefs.message().loops()));
+  for (size_t index = 0; index < 4; ++index) {
+    ASSERT_TRUE(coefs.span().size() > 0u);
+    made_loop.set_index(index);
+    json_loop.set_index(index);
+#define COMPARE(matrix)                              \
+  EXPECT_EQ(json_loop.plant().coefficients().matrix, \
+            made_loop.plant().coefficients().matrix);
+    COMPARE(A_continuous);
+    COMPARE(B_continuous);
+    COMPARE(C);
+    COMPARE(D);
+    COMPARE(U_min);
+    COMPARE(U_max);
+    COMPARE(U_limit_coefficient);
+    COMPARE(U_limit_constant);
+    COMPARE(delayed_u);
+#undef COMPARE
+#define COMPARE(matrix)                                   \
+  EXPECT_EQ(json_loop.controller().coefficients().matrix, \
+            made_loop.controller().coefficients().matrix);
+    COMPARE(K);
+    COMPARE(Kff);
+#undef COMPARE
+#define COMPARE(matrix)                                 \
+  EXPECT_EQ(json_loop.observer().coefficients().matrix, \
+            made_loop.observer().coefficients().matrix);
+    COMPARE(Q_continuous);
+    COMPARE(R_continuous);
+    COMPARE(P_steady_state);
+    COMPARE(delayed_u);
+#undef COMPARE
+  }
+}
+}  // namespace frc971::control_loops::drivetrain::testing
diff --git a/frc971/control_loops/python/test_drivetrain/polydrivetrain.py b/frc971/control_loops/python/test_drivetrain/polydrivetrain.py
new file mode 100644
index 0000000..70ccb34
--- /dev/null
+++ b/frc971/control_loops/python/test_drivetrain/polydrivetrain.py
@@ -0,0 +1,35 @@
+#!/usr/bin/python3
+
+import sys
+from frc971.control_loops.python.test_drivetrain import drivetrain
+from frc971.control_loops.python import polydrivetrain
+
+import gflags
+import glog
+
+__author__ = 'Austin Schuh (austin.linux@gmail.com)'
+
+FLAGS = gflags.FLAGS
+
+try:
+    gflags.DEFINE_bool('plot', False, 'If true, plot the loop response.')
+except gflags.DuplicateFlagError:
+    pass
+
+
+def main(argv):
+    if FLAGS.plot:
+        polydrivetrain.PlotPolyDrivetrainMotions(drivetrain.kDrivetrain)
+    elif len(argv) != 10:
+        glog.fatal('Expected .h, .cc, and .json filenames')
+    else:
+        polydrivetrain.WritePolyDrivetrain(
+            argv[1:4], argv[4:7], argv[7:10],
+            ['frc971', 'control_loops', 'python', 'test_drivetrain'],
+            drivetrain.kDrivetrain)
+
+
+if __name__ == '__main__':
+    argv = FLAGS(sys.argv)
+    glog.init()
+    sys.exit(main(argv))