commit 5dac229f5da196f1191e16d713e81d0adf1b46ed
author Austin Schuh <austin.linux@gmail.com> Sat Oct 19 13:56:58 2024 -0700
committer Austin Schuh <austin.linux@gmail.com> Sat Oct 26 14:58:42 2024 -0700
tree 834ea9fa702dd3fb62ffb13de8d1d9b8d24166be
parent a9550c0b62dcfccb17eaa340f947acf8bd486fd8

Add JAX cost, and test to verify correctness

Turns out there's a bug in the casadi too, so fix that while we are
here.  We were using the unrotated radius for computing the cost
function torque instead of the rotated radius.

Split the MPC code out into a separate file as well to make it easier to
pull into other things.

Change-Id: Iee6c9999fa8b6a91d6963af4edba5cf92a085e9b
Signed-off-by: Austin Schuh <austin.linux@gmail.com>

frc971/control_loops/swerve/generate_physics.cc

diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index 7ea3db5..de52d79 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -70,6 +70,7 @@
 // State per module.
 struct Module {
   DenseMatrix mounting_location;
+  DenseMatrix rotated_mounting_location;
 
   RCP<const Symbol> Is;
 
@@ -938,6 +939,14 @@
         },
         &result_py);
 
+    DefineVector2dVelocityFunction(
+        "rotated_mounting_location",
+        "Returns the mounting location of wheel in field aligned coordinates",
+        [](const Module &m, int dimension) {
+          return ccode(*m.rotated_mounting_location.get(dimension, 0));
+        },
+        &result_py);
+
     DefineScalarFunction(
         "Ms", "Returns the self aligning moment of the ith wheel",
         [this](const Module &m) {
@@ -1189,17 +1198,18 @@
                     DenseMatrix(2, 1, {result.full.Fwx, result.Fwy}),
                     result.full.F);
 
-    DenseMatrix rotated_mounting_location = DenseMatrix(2, 1);
+    result.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.rotated_mounting_location);
+    result.full.torque =
+        force_cross(result.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(rotated_mounting_location, result.direct.F);
+        force_cross(result.rotated_mounting_location, result.direct.F);
 
     VLOG(1);
     VLOG(1) << "full torque = " << result.full.torque->__str__();