Merge "Fix bug in physics and a few in physics_debug" into main
diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index e796ec6..748e3cb 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -1015,14 +1015,18 @@
     mul_dense_dense(R(add(theta_, result.thetas)),
                     DenseMatrix(2, 1, {result.full.Fwx, result.Fwy}),
                     result.full.F);
-    result.full.torque = force_cross(result.mounting_location, result.full.F);
+
+    DenseMatrix rotated_mounting_location = DenseMatrix(2, 1);
+    mul_dense_dense(R(theta_), result.mounting_location,
+                    rotated_mounting_location);
+    result.full.torque = force_cross(rotated_mounting_location, result.full.F);
 
     result.direct.F = DenseMatrix(2, 1);
     mul_dense_dense(R(add(theta_, result.thetas)),
                     DenseMatrix(2, 1, {result.direct.Fwx, result.Fwy}),
                     result.direct.F);
     result.direct.torque =
-        force_cross(result.mounting_location, result.direct.F);
+        force_cross(rotated_mounting_location, result.direct.F);
 
     VLOG(1);
     VLOG(1) << "full torque = " << result.full.torque->__str__();
diff --git a/frc971/control_loops/swerve/jax_dynamics.py b/frc971/control_loops/swerve/jax_dynamics.py
index 76a00d5..dfadbda 100644
--- a/frc971/control_loops/swerve/jax_dynamics.py
+++ b/frc971/control_loops/swerve/jax_dynamics.py
@@ -189,7 +189,7 @@
 
     F = Rthetaplusthetas @ jax.numpy.array([Fwx, Fwy])
 
-    torque = force_cross(mounting_location, F)
+    torque = force_cross(Rtheta @ mounting_location, F)
 
     X_dot_contribution = jax.numpy.hstack((jax.numpy.zeros(
         (4, )), ) * (module_index) + (jax.numpy.array([
@@ -288,7 +288,7 @@
 
     F = Rthetaplusthetas @ jax.numpy.array([Fwx, Fwy])
 
-    torque = force_cross(mounting_location, F)
+    torque = force_cross(Rtheta @ mounting_location, F)
 
     X_dot_contribution = jax.numpy.hstack((jax.numpy.zeros(
         (2, )), ) * (module_index) + (jax.numpy.array([
diff --git a/frc971/control_loops/swerve/physics_debug.py b/frc971/control_loops/swerve/physics_debug.py
index 9a89cc3..04983d4 100644
--- a/frc971/control_loops/swerve/physics_debug.py
+++ b/frc971/control_loops/swerve/physics_debug.py
@@ -182,12 +182,28 @@
         print(result.t.shape)
         xdot = numpy.zeros(result.y.shape)
         print("shape", xdot.shape)
+
+        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], result.y[:, i]),
+                    (8, 1),
+                ),
+            ))
+
+        U_control = numpy.delete(U_control, 0, 1)
+
         for i in range(xdot.shape[1]):
-            xdot[:, i] = self.swerve_physics(result.y[:, i], self.I)[:, 0]
+            xdot[:, i] = self.swerve_physics(result.y[:, i],
+                                             U_control[:, 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])
+            self.print_state(self.swerve_physics, U_control[:, i], result.y[:,
+                                                                            i])
 
         def ev(fn, Y):
             return [fn(Y[:, i], self.I)[0, 0] for i in range(Y.shape[1])]
@@ -217,7 +233,7 @@
                         label=f'steering_accel{i}')
             axs[i].legend()
 
-        fig, axs = pylab.subplots(3)
+        fig, axs = pylab.subplots(4)
         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")
@@ -226,6 +242,7 @@
 
         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, result.y[18, :], label="theta")
         axs[1].plot(result.t, xdot[20, :], label="ay")
         axs[1].plot(result.t, xdot[21, :], label="alpha")
         axs[1].legend()
@@ -239,23 +256,19 @@
         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], result.y[:, i]),
-                    (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()
 
+        F_plot = numpy.array([
+            self.F[0](result.y[:, i], U_control[:, i])
+            for i in range(result.y[0, :].size)
+        ])
+        axs[3].plot(result.t, F_plot[:, 0], label="force_x0")
+        axs[3].plot(result.t, F_plot[:, 1], label="force_y0")
+        axs[3].legend()
+
         pylab.show()
 
 
diff --git a/frc971/control_loops/swerve/physics_test.py b/frc971/control_loops/swerve/physics_test.py
index af3dc17..6e7ddf4 100644
--- a/frc971/control_loops/swerve/physics_test.py
+++ b/frc971/control_loops/swerve/physics_test.py
@@ -677,6 +677,34 @@
         self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS2, 0], 0.0)
         self.assertAlmostEqual(Xdot[dynamics.VELOCITY_STATE_OMEGAS3, 0], 0.0)
 
+    def test_constant_torque(self):
+        """Tests that the robot exerts the same amount of torque no matter the orientation"""
+        steer_I = numpy.reshape(numpy.array([(i % 2) * 20 for i in range(8)]),
+                                (8, 1))
+
+        X = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
+                               omega=0.0,
+                               module_angles=[
+                                   3 * numpy.pi / 4.0, -3 * numpy.pi / 4.0,
+                                   -numpy.pi / 4.0, numpy.pi / 4.0
+                               ],
+                               drive_wheel_velocity=1.0)
+        Xdot = self.swerve_full_dynamics(X, steer_I, skip_compare=True)
+
+        X_rot = utils.state_vector(velocity=numpy.array([[0.0], [0.0]]),
+                                   omega=0.0,
+                                   theta=numpy.pi,
+                                   module_angles=[
+                                       3 * numpy.pi / 4.0, -3 * numpy.pi / 4.0,
+                                       -numpy.pi / 4.0, numpy.pi / 4.0
+                                   ],
+                                   drive_wheel_velocity=1.0)
+        Xdot_rot = self.swerve_full_dynamics(X_rot, steer_I, skip_compare=True)
+
+        self.assertGreater(Xdot[dynamics.STATE_OMEGA, 0], 0.0)
+        self.assertAlmostEquals(Xdot[dynamics.STATE_OMEGA, 0],
+                                Xdot_rot[dynamics.STATE_OMEGA, 0])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/frc971/control_loops/swerve/swerve_notes.tex b/frc971/control_loops/swerve/swerve_notes.tex
index bf83dcc..380d6c8 100644
--- a/frc971/control_loops/swerve/swerve_notes.tex
+++ b/frc971/control_loops/swerve/swerve_notes.tex
@@ -376,7 +376,7 @@
 \end{gather}
 However, the angular velocity differential equation is more complicated as the torque added per module varies in sign
 \begin{gather}
-    \ddot{\theta} = \frac{\Sigma\left(\harpoon{r} \times \harpoon{F}_{mod}\right) + \tau_{d}}{J_{robot}}
+    \ddot{\theta} = \frac{\Sigma\left(R\left(\theta\right)\harpoon{r} \times \harpoon{F}_{mod}\right) + \tau_{d}}{J_{robot}}
 \end{gather}
 \subsection{Simplified longitudinal dynamics}