commit 76534f36c8be2864e543bd53080ece7d672f7898
author Austin Schuh <austin.linux@gmail.com> Mon Sep 02 13:52:45 2024 -0700
committer Austin Schuh <austin.linux@gmail.com> Sat Sep 14 16:58:25 2024 -0700
tree 8713b1d80407bb1a3cd4849b4bc6b9b7f9b10426
parent 3f3d7ca4eb810d91f66ba267447b41ea3ac963df

Add JAX version of the physics, and tests to confirm it matches

This sets us up to very efficiently compute dynamics and jit.

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

frc971/control_loops/swerve/BUILD

diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index c7355b5..26828e4 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -197,14 +197,51 @@
     ],
 )
 
+py_library(
+    name = "jax_dynamics",
+    srcs = [
+        "jax_dynamics.py",
+    ],
+    deps = [
+        ":dynamics",
+        "//frc971/control_loops/python:controls",
+        "@pip//jax",
+    ],
+)
+
 py_test(
-    name = "physics_test",
+    name = "physics_test_cpu",
     srcs = [
         "physics_test.py",
     ],
+    env = {
+        "JAX_PLATFORMS": "cpu",
+    },
+    main = "physics_test.py",
     target_compatible_with = ["@platforms//cpu:x86_64"],
     deps = [
         ":dynamics",
+        ":jax_dynamics",
+        ":physics_test_utils",
+        "@pip//casadi",
+        "@pip//numpy",
+        "@pip//scipy",
+    ],
+)
+
+py_test(
+    name = "physics_test_gpu",
+    srcs = [
+        "physics_test.py",
+    ],
+    env = {
+        "JAX_PLATFORMS": "cuda",
+    },
+    main = "physics_test.py",
+    target_compatible_with = ["@platforms//cpu:x86_64"],
+    deps = [
+        ":dynamics",
+        ":jax_dynamics",
         ":physics_test_utils",
         "@pip//casadi",
         "@pip//numpy",