Add debug script and break out test utils

Signed-off-by: justinT21 <jjturcot@gmail.com>
Change-Id: I5913be5bf7aed9ef9f274d3ad4ced61b1738ca71
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index 29dbf76..3755ab4 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -170,7 +170,7 @@
 )
 
 cc_library(
-    name = "dynamics",
+    name = "eigen_dynamics",
     srcs = ["dynamics.cc"],
     hdrs = ["dynamics.h"],
     deps = [
@@ -196,14 +196,49 @@
 py_test(
     name = "physics_test",
     srcs = [
-        "bigcaster_dynamics.py",
-        "dynamics.py",
-        "nocaster_dynamics.py",
         "physics_test.py",
     ],
     target_compatible_with = ["@platforms//cpu:x86_64"],
     deps = [
+        ":physics_test_utils",
         "@pip//casadi",
         "@pip//numpy",
     ],
 )
+
+py_library(
+    name = "dynamics",
+    srcs = [
+        "bigcaster_dynamics.py",
+        "dynamics.py",
+        "nocaster_dynamics.py",
+    ],
+    deps = [
+        "@pip//casadi",
+    ],
+)
+
+py_library(
+    name = "physics_test_utils",
+    srcs = [
+        "physics_test_utils.py",
+    ],
+    deps = [
+        ":dynamics",
+        "@pip//casadi",
+        "@pip//numpy",
+    ],
+)
+
+py_binary(
+    name = "physics_debug",
+    srcs = [
+        "physics_debug.py",
+    ],
+    deps = [
+        ":physics_test_utils",
+        "@pip//matplotlib",
+        "@pip//pygobject",
+        "@pip//scipy",
+    ],
+)
diff --git a/frc971/control_loops/swerve/physics_debug.py b/frc971/control_loops/swerve/physics_debug.py
new file mode 100644
index 0000000..db9ea85
--- /dev/null
+++ b/frc971/control_loops/swerve/physics_debug.py
@@ -0,0 +1,265 @@
+#!/usr/bin/python3
+
+import numpy, scipy
+
+from matplotlib import pylab
+from frc971.control_loops.swerve import physics_test_utils as utils
+from frc971.control_loops.swerve import dynamics
+
+
+class PhysicsDebug(object):
+
+    def wrap(self, python_module):
+        self.swerve_physics = utils.wrap(python_module.swerve_physics)
+        self.contact_patch_velocity = [
+            utils.wrap_module(python_module.contact_patch_velocity, i)
+            for i in range(4)
+        ]
+        self.wheel_ground_velocity = [
+            utils.wrap_module(python_module.wheel_ground_velocity, i)
+            for i in range(4)
+        ]
+        self.wheel_slip_velocity = [
+            utils.wrap_module(python_module.wheel_slip_velocity, i)
+            for i in range(4)
+        ]
+        self.wheel_force = [
+            utils.wrap_module(python_module.wheel_force, i) for i in range(4)
+        ]
+        self.module_angular_accel = [
+            utils.wrap_module(python_module.module_angular_accel, i)
+            for i in range(4)
+        ]
+        self.F = [utils.wrap_module(python_module.F, i) for i in range(4)]
+        self.mounting_location = [
+            utils.wrap_module(python_module.mounting_location, i)
+            for i in range(4)
+        ]
+
+        self.slip_angle = [
+            utils.wrap_module(python_module.slip_angle, i) for i in range(4)
+        ]
+        self.slip_ratio = [
+            utils.wrap_module(python_module.slip_ratio, i) for i in range(4)
+        ]
+        self.Ms = [utils.wrap_module(python_module.Ms, i) for i in range(4)]
+
+    def print_state(self, swerve_physics, I, x):
+        xdot = swerve_physics(x, I)
+
+        for i in range(4):
+            print(f"  Slip Angle {i} {self.slip_angle[0](x, I)}")
+            print(f"  Slip Ratio {i} {self.slip_ratio[0](x, I)}")
+
+        print("  Steering angle0", x[0])
+        print("  Steering angle1", x[4])
+        print("  Steering angle2", x[8])
+        print("  Steering angle3", x[12])
+        print("  Steering velocity0", xdot[0])
+        print("  Steering velocity1", xdot[4])
+        print("  Steering velocity2", xdot[8])
+        print("  Steering velocity3", xdot[12])
+        print("  Steering accel0", xdot[2])
+        print("  Steering accel1", xdot[6])
+        print("  Steering accel2", xdot[10])
+        print("  Steering accel3", xdot[14])
+        print("  Drive accel0", xdot[3])
+        print("  Drive accel1", xdot[7])
+        print("  Drive accel2", xdot[11])
+        print("  Drive accel3", xdot[15])
+        print("  Drive velocity0", x[3] * 2 * 0.0254)
+        print("  Drive velocity1", x[7] * 2 * 0.0254)
+        print("  Drive velocity2", x[11] * 2 * 0.0254)
+        print("  Drive velocity3", x[15] * 2 * 0.0254)
+        print("  Theta ", x[18])
+        print("  Omega ", x[21])
+        print("  Alpha", xdot[21])
+        print("  vx", xdot[16])
+        print("  vy", xdot[17])
+        print("  ax", xdot[19])
+        print("  ay", xdot[20])
+        print("  Fdx ", x[22])
+        print("  Fdy ", x[23])
+        print("  Moment_d", xdot[24])
+
+    def plot(self):
+        velocity = numpy.array([[1.0], [0.0]])
+
+        flip = False
+        if flip:
+            module_angles = [0.01 + numpy.pi] * 4
+        else:
+            module_angles = [-0.1, -0.1, 0.1, 0.1]
+
+        X = utils.state_vector(
+            velocity=velocity,
+            drive_wheel_velocity=-1.0 if flip else 1.0,
+            omega=0.0,
+            module_angles=module_angles,
+        )
+
+        self.I = numpy.array([[40.0], [0.0], [40.0], [0.0], [40.0], [0.0],
+                              [40.0], [0.0]])
+
+        def calc_I(t, x):
+            x_goal = numpy.zeros((16, 1))
+
+            Kp_steer = 15.0
+            Kp_drive = 0.0
+            Kd_steer = 7.0
+            Kd_drive = 0.0
+            Idrive = 5.0 if t < 5.0 else 15.0
+
+            return numpy.array([
+                [
+                    Kd_steer * (x_goal[2] - x[2]) + Kp_steer *
+                    (x_goal[0] - x[0])
+                ],
+                [Idrive],
+                [
+                    Kd_steer * (x_goal[6] - x[6]) + Kp_steer *
+                    (x_goal[4] - x[4])
+                ],
+                [Idrive],
+                [
+                    Kd_steer * (x_goal[10] - x[10]) + Kp_steer *
+                    (x_goal[8] - x[8])
+                ],
+                [Idrive],
+                [
+                    Kd_steer * (x_goal[14] - x[14]) + Kp_steer *
+                    (x_goal[12] - x[12])
+                ],
+                [Idrive],
+            ]).flatten()
+            return numpy.array([
+                [
+                    Kd_steer * (x_goal[2] - x[2]) + Kp_steer *
+                    (x_goal[0] - x[0])
+                ],
+                [
+                    Kd_drive * (x_goal[3] - x[3]) + Kp_drive *
+                    (x_goal[1] - x[1])
+                ],
+                [
+                    Kd_steer * (x_goal[6] - x[6]) + Kp_steer *
+                    (x_goal[4] - x[4])
+                ],
+                [
+                    Kd_drive * (x_goal[7] - x[7]) + Kp_drive *
+                    (x_goal[5] - x[5])
+                ],
+                [
+                    Kd_steer * (x_goal[10] - x[10]) + Kp_steer *
+                    (x_goal[8] - x[8])
+                ],
+                [
+                    Kd_drive * (x_goal[11] - x[11]) + Kp_drive *
+                    (x_goal[9] - x[9])
+                ],
+                [
+                    Kd_steer * (x_goal[14] - x[14]) + Kp_steer *
+                    (x_goal[12] - x[12])
+                ],
+                [
+                    Kd_drive * (x_goal[15] - x[15]) + Kp_drive *
+                    (x_goal[13] - x[13])
+                ],
+            ]).flatten()
+
+        t_eval = numpy.arange(0, 10.0, 0.005)
+        self.wrap(dynamics)
+        result = scipy.integrate.solve_ivp(
+            lambda t, x: self.swerve_physics(x, calc_I(t, x)).flatten(),
+            [0, t_eval[-1]],
+            X.flatten(),
+            t_eval=t_eval,
+        )
+        print("Function evaluations", result.nfev)
+
+        # continue
+        print(result.y.shape)
+        print(result.t.shape)
+        xdot = numpy.zeros(result.y.shape)
+        print("shape", xdot.shape)
+        for i in range(xdot.shape[1]):
+            xdot[:, i] = self.swerve_physics(result.y[:, i], self.I)[:, 0]
+
+        for i in range(2):
+            print(f"For t {i * 0.005}")
+            self.print_state(self.swerve_physics, self.I, result.y[:, i])
+
+        def ev(fn, Y):
+            return [fn(Y[:, i], self.I)[0, 0] for i in range(Y.shape[1])]
+
+        fig, axs = pylab.subplots(2)
+        axs[0].plot(result.t, result.y[0, :], label="steer0")
+        axs[0].plot(result.t, result.y[4, :], label="steer1")
+        axs[0].plot(result.t, result.y[8, :], label="steer2")
+        axs[0].plot(result.t, result.y[12, :], label="steer3")
+        axs[0].legend()
+        axs[1].plot(result.t, result.y[2, :], label="steer_velocity0")
+        axs[1].plot(result.t, result.y[6, :], label="steer_velocity1")
+        axs[1].plot(result.t, result.y[10, :], label="steer_velocity2")
+        axs[1].plot(result.t, result.y[14, :], label="steer_velocity3")
+        axs[1].legend()
+
+        fig, axs = pylab.subplots(2)
+        for i in range(len(axs)):
+            axs[i].plot(result.t,
+                        ev(self.slip_angle[i], result.y),
+                        label=f"slip_angle{i}")
+            axs[i].plot(result.t,
+                        ev(self.slip_ratio[i], result.y),
+                        label=f"slip_ratio{i}")
+            axs[i].plot(result.t,
+                        xdot[2 + i * 4, :],
+                        label=f'steering_accel{i}')
+            axs[i].legend()
+
+        fig, axs = pylab.subplots(3)
+        axs[0].plot(result.t, result.y[3, :], label="drive_velocity0")
+        axs[0].plot(result.t, result.y[7, :], label="drive_velocity1")
+        axs[0].plot(result.t, result.y[11, :], label="drive_velocity2")
+        axs[0].plot(result.t, result.y[15, :], label="drive_velocity3")
+        axs[0].legend()
+
+        axs[1].plot(result.t, result.y[20, :], label="vy")
+        axs[1].plot(result.t, result.y[21, :], label="omega")
+        axs[1].plot(result.t, xdot[20, :], label="ay")
+        axs[1].plot(result.t, xdot[21, :], label="alpha")
+        axs[1].legend()
+
+        axs[2].plot(result.t, result.y[19, :], label="vx")
+        axs[2].plot(result.t, xdot[19, :], label="ax")
+
+        axs[2].plot(result.t,
+                    numpy.hypot(result.y[19, :], result.y[20, :]),
+                    label="speed")
+        axs[2].legend()
+
+        pylab.figure()
+        U_control = numpy.zeros((8, 1))
+
+        for i in range(numpy.shape(result.t)[0]):
+            U_control = numpy.hstack((
+                U_control,
+                numpy.reshape(
+                    calc_I(result.t[i], numpy.reshape(result.y[:, i],
+                                                      (25, 1))),
+                    (8, 1),
+                ),
+            ))
+
+        U_control = numpy.delete(U_control, 0, 1)
+
+        pylab.plot(result.t, U_control[0, :], label="Is0")
+        pylab.plot(result.t, U_control[1, :], label="Id0")
+        pylab.legend()
+
+        pylab.show()
+
+
+if __name__ == "__main__":
+    debug = PhysicsDebug()
+    debug.plot()
diff --git a/frc971/control_loops/swerve/physics_test.py b/frc971/control_loops/swerve/physics_test.py
index ecfc1ce..c7d60e9 100644
--- a/frc971/control_loops/swerve/physics_test.py
+++ b/frc971/control_loops/swerve/physics_test.py
@@ -9,101 +9,46 @@
 from frc971.control_loops.swerve import bigcaster_dynamics
 from frc971.control_loops.swerve import dynamics
 from frc971.control_loops.swerve import nocaster_dynamics
-
-
-def state_vector(velocity=numpy.array([[1.0], [0.0]]),
-                 dx=0.0,
-                 dy=0.0,
-                 theta=0.0,
-                 omega=0.0,
-                 module_omega=0.0,
-                 module_angle=0.0,
-                 drive_wheel_velocity=None,
-                 module_angles=None):
-    """Returns the state vector with the requested state."""
-    X_initial = numpy.zeros((25, 1))
-    # All the wheels are spinning at the speed needed to hit the velocity in m/s
-    drive_wheel_velocity = (drive_wheel_velocity
-                            or numpy.linalg.norm(velocity))
-
-    X_initial[2, 0] = module_omega
-    X_initial[3, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
-
-    X_initial[6, 0] = module_omega
-    X_initial[7, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
-
-    X_initial[10, 0] = module_omega
-    X_initial[11, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
-
-    X_initial[14, 0] = module_omega
-    X_initial[15, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
-
-    X_initial[0, 0] = module_angle
-    X_initial[4, 0] = module_angle
-    X_initial[8, 0] = module_angle
-    X_initial[12, 0] = module_angle
-
-    if module_angles is not None:
-        assert (len(module_angles) == 4)
-        X_initial[0, 0] = module_angles[0]
-        X_initial[4, 0] = module_angles[1]
-        X_initial[8, 0] = module_angles[2]
-        X_initial[12, 0] = module_angles[3]
-
-    X_initial[18, 0] = theta
-
-    X_initial[19, 0] = velocity[0, 0] + dx
-    X_initial[20, 0] = velocity[1, 0] + dy
-    X_initial[21, 0] = omega
-
-    return X_initial
-
-
-def wrap(fn):
-    evaluated_fn = fn(casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
-    return lambda X, U: numpy.array(evaluated_fn(X, U))
-
-
-def wrap_module(fn, i):
-    evaluated_fn = fn(i, casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
-    return lambda X, U: numpy.array(evaluated_fn(X, U))
+from frc971.control_loops.swerve import physics_test_utils as utils
 
 
 class TestSwervePhysics(unittest.TestCase):
     I = numpy.zeros((8, 1))
 
     def wrap(self, python_module):
-        self.swerve_physics = wrap(python_module.swerve_physics)
+        self.swerve_physics = utils.wrap(python_module.swerve_physics)
         self.contact_patch_velocity = [
-            wrap_module(python_module.contact_patch_velocity, i)
+            utils.wrap_module(python_module.contact_patch_velocity, i)
             for i in range(4)
         ]
         self.wheel_ground_velocity = [
-            wrap_module(python_module.wheel_ground_velocity, i)
+            utils.wrap_module(python_module.wheel_ground_velocity, i)
             for i in range(4)
         ]
         self.wheel_slip_velocity = [
-            wrap_module(python_module.wheel_slip_velocity, i) for i in range(4)
-        ]
-        self.wheel_force = [
-            wrap_module(python_module.wheel_force, i) for i in range(4)
-        ]
-        self.module_angular_accel = [
-            wrap_module(python_module.module_angular_accel, i)
+            utils.wrap_module(python_module.wheel_slip_velocity, i)
             for i in range(4)
         ]
-        self.F = [wrap_module(python_module.F, i) for i in range(4)]
+        self.wheel_force = [
+            utils.wrap_module(python_module.wheel_force, i) for i in range(4)
+        ]
+        self.module_angular_accel = [
+            utils.wrap_module(python_module.module_angular_accel, i)
+            for i in range(4)
+        ]
+        self.F = [utils.wrap_module(python_module.F, i) for i in range(4)]
         self.mounting_location = [
-            wrap_module(python_module.mounting_location, i) for i in range(4)
+            utils.wrap_module(python_module.mounting_location, i)
+            for i in range(4)
         ]
 
         self.slip_angle = [
-            wrap_module(python_module.slip_angle, i) for i in range(4)
+            utils.wrap_module(python_module.slip_angle, i) for i in range(4)
         ]
         self.slip_ratio = [
-            wrap_module(python_module.slip_ratio, i) for i in range(4)
+            utils.wrap_module(python_module.slip_ratio, i) for i in range(4)
         ]
-        self.Ms = [wrap_module(python_module.Ms, i) for i in range(4)]
+        self.Ms = [utils.wrap_module(python_module.Ms, i) for i in range(4)]
 
     def setUp(self):
         self.wrap(dynamics)
@@ -111,10 +56,8 @@
     def test_contact_patch_velocity(self):
         """Tests that the contact patch velocity makes sense."""
         for i in range(4):
-            contact_patch_velocity = wrap_module(
-                dynamics.contact_patch_velocity, i)
-            wheel_ground_velocity = wrap_module(dynamics.wheel_ground_velocity,
-                                                i)
+            contact_patch_velocity = self.contact_patch_velocity[i]
+            wheel_ground_velocity = self.wheel_ground_velocity[i]
 
             # No angular velocity should result in just linear motion.
             for velocity in [
@@ -126,7 +69,8 @@
             ]:
                 for theta in [0.0, 1.0, -1.0, 100.0]:
                     patch_velocity = contact_patch_velocity(
-                        state_vector(velocity=velocity, theta=theta), self.I)
+                        utils.state_vector(velocity=velocity, theta=theta),
+                        self.I)
 
                     assert_array_equal(patch_velocity, velocity)
 
@@ -137,33 +81,41 @@
                 for omega in [0.65, -0.1]:
                     # Point each module to the center to make the math easier.
                     patch_velocity = contact_patch_velocity(
-                        state_vector(velocity=numpy.array([[0.0], [0.0]]),
-                                     theta=theta,
-                                     omega=omega,
-                                     module_angle=module_center_of_mass_angle),
-                        self.I)
+                        utils.state_vector(
+                            velocity=numpy.array([[0.0], [0.0]]),
+                            theta=theta,
+                            omega=omega,
+                            module_angle=module_center_of_mass_angle,
+                        ),
+                        self.I,
+                    )
 
                     assert_array_almost_equal(
                         patch_velocity,
                         (dynamics.ROBOT_WIDTH / numpy.sqrt(2.0) -
-                         dynamics.CASTER) * omega * numpy.array([[
-                             -numpy.sin(theta + module_center_of_mass_angle)
-                         ], [numpy.cos(theta + module_center_of_mass_angle)]]))
+                         dynamics.CASTER) * omega * numpy.array([
+                             [-numpy.sin(theta + module_center_of_mass_angle)],
+                             [numpy.cos(theta + module_center_of_mass_angle)],
+                         ]),
+                    )
 
                     # Point the wheel along +x, rotate it by theta, then spin it.
                     # Confirm the velocities come out right.
                     patch_velocity = contact_patch_velocity(
-                        state_vector(
+                        utils.state_vector(
                             velocity=numpy.array([[0.0], [0.0]]),
                             theta=-module_center_of_mass_angle,
                             module_omega=omega,
-                            module_angle=(theta +
-                                          module_center_of_mass_angle)),
-                        self.I)
+                            module_angle=(theta + module_center_of_mass_angle),
+                        ),
+                        self.I,
+                    )
 
                     assert_array_almost_equal(
-                        patch_velocity, -dynamics.CASTER * omega *
-                        numpy.array([[-numpy.sin(theta)], [numpy.cos(theta)]]))
+                        patch_velocity,
+                        -dynamics.CASTER * omega *
+                        numpy.array([[-numpy.sin(theta)], [numpy.cos(theta)]]),
+                    )
 
             # Now, test that the rotation back to wheel coordinates works.
             # The easiest way to do this is to point the wheel in a direction,
@@ -172,11 +124,16 @@
                 for module_angle in [0.0, 1.0, -5.0]:
                     wheel_patch_velocity = numpy.array(
                         wheel_ground_velocity(
-                            state_vector(velocity=numpy.array(
-                                [[numpy.cos(robot_angle + module_angle)],
-                                 [numpy.sin(robot_angle + module_angle)]]),
-                                         theta=robot_angle,
-                                         module_angle=module_angle), self.I))
+                            utils.state_vector(
+                                velocity=numpy.array([
+                                    [numpy.cos(robot_angle + module_angle)],
+                                    [numpy.sin(robot_angle + module_angle)],
+                                ]),
+                                theta=robot_angle,
+                                module_angle=module_angle,
+                            ),
+                            self.I,
+                        ))
 
                     assert_array_almost_equal(wheel_patch_velocity,
                                               numpy.array([[1], [0]]))
@@ -193,21 +150,25 @@
                     # We have redefined the angle to be the sin of the angle.
                     # That way, when the module flips directions, the slip angle also flips
                     # directions to keep it stable.
-                    computed_angle = self.slip_angle[i](state_vector(
-                        velocity=velocity,
-                        module_angle=numpy.pi * wrap + theta), self.I)[0, 0]
+                    computed_angle = self.slip_angle[i](
+                        utils.state_vector(velocity=velocity,
+                                           module_angle=numpy.pi * wrap +
+                                           theta),
+                        self.I,
+                    )[0, 0]
 
                     expected = numpy.sin(numpy.pi * wrap + theta)
 
                     self.assertAlmostEqual(
                         expected,
                         computed_angle,
-                        msg=f"Trying wrap {wrap} theta {theta}")
+                        msg=f"Trying wrap {wrap} theta {theta}",
+                    )
 
     def test_wheel_torque(self):
         """Tests that the per module self aligning forces have the right signs."""
         # Point all the modules in a little bit.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[-0.001, -0.001, 0.001, 0.001],
         )
@@ -234,7 +195,7 @@
 
         # Now, make the bot want to go left by going to the other side.
         # The wheels will be going too fast based on our calcs, so they should be decelerating.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[0.01, 0.01, 0.01, 0.01],
         )
@@ -262,7 +223,7 @@
         self.assertAlmostEqual(xdot_left[21, 0], 0.0)
 
         # And now do it to the right too.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[-0.01, -0.01, -0.01, -0.01],
         )
@@ -293,7 +254,7 @@
         """Tests that the per module self aligning forces have the right signs when going backwards."""
         self.wrap(nocaster_dynamics)
         # Point all the modules in a little bit, going backwards.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[
                 numpy.pi - 0.001,
@@ -326,7 +287,7 @@
 
         # Now, make the bot want to go left by going to the other side.
         # The wheels will be going too fast based on our calcs, so they should be decelerating.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[numpy.pi + 0.01] * 4,
             drive_wheel_velocity=-1,
@@ -355,7 +316,7 @@
         self.assertAlmostEqual(xdot_left[21, 0], 0.0)
 
         # And now do it to the right too.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             drive_wheel_velocity=-1,
             module_angles=[-0.01 + numpy.pi] * 4,
@@ -387,7 +348,7 @@
         """Tests that the per module self aligning forces have the right signs when going backwards with a lot of caster."""
         self.wrap(bigcaster_dynamics)
         # Point all the modules in a little bit, going backwards.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[
                 numpy.pi - 0.001,
@@ -420,7 +381,7 @@
 
         # Now, make the bot want to go left by going to the other side.
         # The wheels will be going too fast based on our calcs, so they should be decelerating.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             module_angles=[numpy.pi + 0.01] * 4,
             drive_wheel_velocity=-1,
@@ -449,7 +410,7 @@
         self.assertAlmostEqual(xdot_left[21, 0], 0.0)
 
         # And now do it to the right too.
-        X = state_vector(
+        X = utils.state_vector(
             velocity=numpy.array([[1.0], [0.0]]),
             drive_wheel_velocity=-1,
             module_angles=[-0.01 + numpy.pi] * 4,
@@ -480,9 +441,9 @@
     def test_wheel_forces(self):
         """Tests that the per module forces have the right signs."""
         for i in range(4):
-            wheel_force = wrap_module(dynamics.wheel_force, i)
+            wheel_force = self.wheel_force[i]
 
-            X = state_vector()
+            X = utils.state_vector()
             robot_equal = wheel_force(X, self.I)
             xdot_equal = self.swerve_physics(X, self.I)
             self.assertEqual(robot_equal[0, 0], 0.0)
@@ -491,7 +452,7 @@
             self.assertEqual(xdot_equal[3 + 4 * i], 0.0)
 
             # Robot is moving faster than the wheels, it should decelerate.
-            X = state_vector(dx=0.01)
+            X = utils.state_vector(dx=0.01)
             robot_faster = wheel_force(X, self.I)
             xdot_faster = self.swerve_physics(X, self.I)
             self.assertLess(robot_faster[0, 0], -0.1)
@@ -499,7 +460,7 @@
             self.assertGreater(xdot_faster[3 + 4 * i], 0.0)
 
             # Robot is now going slower than the wheels.  It should accelerate.
-            X = state_vector(dx=-0.01)
+            X = utils.state_vector(dx=-0.01)
             robot_slower = wheel_force(X, self.I)
             xdot_slower = self.swerve_physics(X, self.I)
             self.assertGreater(robot_slower[0, 0], 0.1)
@@ -507,15 +468,15 @@
             self.assertLess(xdot_slower[3 + 4 * i], 0.0)
 
             # Positive lateral velocity -> negative force.
-            robot_left = wheel_force(state_vector(dy=0.01), self.I)
+            robot_left = wheel_force(utils.state_vector(dy=0.01), self.I)
             self.assertEqual(robot_left[0, 0], 0.0)
             self.assertLess(robot_left[1, 0], -0.1)
 
             # Negative lateral velocity -> positive force.
-            robot_right = wheel_force(state_vector(dy=-0.01), self.I)
+            robot_right = wheel_force(utils.state_vector(dy=-0.01), self.I)
             self.assertEqual(robot_right[0, 0], 0.0)
             self.assertGreater(robot_right[1, 0], 0.1)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/frc971/control_loops/swerve/physics_test_utils.py b/frc971/control_loops/swerve/physics_test_utils.py
new file mode 100644
index 0000000..30b352f
--- /dev/null
+++ b/frc971/control_loops/swerve/physics_test_utils.py
@@ -0,0 +1,64 @@
+#!/usr/bin/python3
+
+import casadi, numpy
+
+from frc971.control_loops.swerve import dynamics
+
+
+def state_vector(
+    velocity=numpy.array([[1.0], [0.0]]),
+    dx=0.0,
+    dy=0.0,
+    theta=0.0,
+    omega=0.0,
+    module_omega=0.0,
+    module_angle=0.0,
+    drive_wheel_velocity=None,
+    module_angles=None,
+):
+    """Returns the state vector with the requested state."""
+    X_initial = numpy.zeros((25, 1))
+    # All the wheels are spinning at the speed needed to hit the velocity in m/s
+    drive_wheel_velocity = drive_wheel_velocity or numpy.linalg.norm(velocity)
+
+    X_initial[2, 0] = module_omega
+    X_initial[3, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
+
+    X_initial[6, 0] = module_omega
+    X_initial[7, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
+
+    X_initial[10, 0] = module_omega
+    X_initial[11, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
+
+    X_initial[14, 0] = module_omega
+    X_initial[15, 0] = drive_wheel_velocity / (dynamics.WHEEL_RADIUS)
+
+    X_initial[0, 0] = module_angle
+    X_initial[4, 0] = module_angle
+    X_initial[8, 0] = module_angle
+    X_initial[12, 0] = module_angle
+
+    if module_angles is not None:
+        assert len(module_angles) == 4
+        X_initial[0, 0] = module_angles[0]
+        X_initial[4, 0] = module_angles[1]
+        X_initial[8, 0] = module_angles[2]
+        X_initial[12, 0] = module_angles[3]
+
+    X_initial[18, 0] = theta
+
+    X_initial[19, 0] = velocity[0, 0] + dx
+    X_initial[20, 0] = velocity[1, 0] + dy
+    X_initial[21, 0] = omega
+
+    return X_initial
+
+
+def wrap(fn):
+    evaluated_fn = fn(casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
+    return lambda X, U: numpy.array(evaluated_fn(X, U))
+
+
+def wrap_module(fn, i):
+    evaluated_fn = fn(i, casadi.SX.sym("X", 25, 1), casadi.SX.sym("U", 8, 1))
+    return lambda X, U: numpy.array(evaluated_fn(X, U))