[wpimath] Add Sleipnir Java bindings

Author: calcmogul

benchmark/src/main/java/wpilib/robot/CartPoleBenchmark.java

@@ -0,0 +1,128 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package wpilib.robot;
+
+import static org.wpilib.math.autodiff.NumericalIntegration.rk4;
+import static org.wpilib.math.autodiff.Variable.cos;
+import static org.wpilib.math.autodiff.Variable.sin;
+import static org.wpilib.math.autodiff.VariableMatrix.solve;
+import static org.wpilib.math.optimization.Constraints.eq;
+import static org.wpilib.math.optimization.Constraints.ge;
+import static org.wpilib.math.optimization.Constraints.le;
+
+import org.ejml.simple.SimpleMatrix;
+import org.wpilib.math.autodiff.Variable;
+import org.wpilib.math.autodiff.VariableMatrix;
+import org.wpilib.math.optimization.Problem;
+import org.wpilib.math.optimization.solver.Options;
+import org.wpilib.math.util.MathUtil;
+
+public final class CartPoleBenchmark {
+  private CartPoleBenchmark() {
+    // Utility class.
+  }
+
+  @SuppressWarnings("LocalVariableName")
+  private static VariableMatrix cartPoleDynamics(VariableMatrix x, VariableMatrix u) {
+    final double m_c = 5.0; // Cart mass (kg)
+    final double m_p = 0.5; // Pole mass (kg)
+    final double l = 0.5; // Pole length (m)
+    final double g = 9.806; // Acceleration due to gravity (m/s²)
+
+    var q = x.segment(0, 2);
+    var qdot = x.segment(2, 2);
+    var theta = q.get(1);
+    var thetadot = qdot.get(1);
+
+    //        [ m_c + m_p  m_p l cosθ]
+    // M(q) = [m_p l cosθ    m_p l²  ]
+    var M =
+        new VariableMatrix(
+            new Variable[][] {
+              {new Variable(m_c + m_p), cos(theta).times(m_p * l)},
+              {cos(theta).times(m_p * l), new Variable(m_p * Math.pow(l, 2))}
+            });
+
+    //           [0  −m_p lθ̇ sinθ]
+    // C(q, q̇) = [0       0      ]
+    var C =
+        new VariableMatrix(
+            new Variable[][] {
+              {new Variable(0), thetadot.times(-m_p * l).times(sin(theta))},
+              {new Variable(0), new Variable(0)}
+            });
+
+    //          [     0      ]
+    // τ_g(q) = [-m_p gl sinθ]
+    var tau_g =
+        new VariableMatrix(new Variable[][] {{new Variable(0)}, {sin(theta).times(-m_p * g * l)}});
+
+    //     [1]
+    // B = [0]
+    var B = new VariableMatrix(new double[][] {{1}, {0}});
+
+    // q̈ = M⁻¹(q)(τ_g(q) − C(q, q̇)q̇ + Bu)
+    var qddot = new VariableMatrix(4);
+    qddot.segment(0, 2).set(qdot);
+    qddot.segment(2, 2).set(solve(M, tau_g.minus(C.times(qdot)).plus(B.times(u))));
+    return qddot;
+  }
+
+  /** Cart-pole benchmark. */
+  public static void cartPole() {
+    final double T = 5.0; // s
+    final double dt = 0.05; // s
+    final int N = (int) (T / dt);
+
+    final double u_max = 20.0; // N
+    final double d_max = 2.0; // m
+
+    final var x_initial = new SimpleMatrix(new double[][] {{0.0}, {0.0}, {0.0}, {0.0}});
+    final var x_final = new SimpleMatrix(new double[][] {{1.0}, {Math.PI}, {0.0}, {0.0}});
+
+    var problem = new Problem();
+
+    // x = [q, q̇]ᵀ = [x, θ, ẋ, θ̇]ᵀ
+    var X = problem.decisionVariable(4, N + 1);
+
+    // Initial guess
+    for (int k = 0; k < N + 1; ++k) {
+      X.get(0, k).setValue(MathUtil.lerp(x_initial.get(0), x_final.get(0), (double) k / N));
+      X.get(1, k).setValue(MathUtil.lerp(x_initial.get(1), x_final.get(1), (double) k / N));
+    }
+
+    // u = f_x
+    var U = problem.decisionVariable(1, N);
+
+    // Initial conditions
+    problem.subjectTo(eq(X.col(0), x_initial));
+
+    // Final conditions
+    problem.subjectTo(eq(X.col(N), x_final));
+
+    // Cart position constraints
+    problem.subjectTo(ge(X.row(0), 0.0));
+    problem.subjectTo(le(X.row(0), d_max));
+
+    // Input constraints
+    problem.subjectTo(ge(U, -u_max));
+    problem.subjectTo(le(U, u_max));
+
+    // Dynamics constraints - RK4 integration
+    for (int k = 0; k < N; ++k) {
+      problem.subjectTo(
+          eq(X.col(k + 1), rk4(CartPoleBenchmark::cartPoleDynamics, X.col(k), U.col(k), dt)));
+    }
+
+    // Minimize sum squared inputs
+    var J = new Variable(0.0);
+    for (int k = 0; k < N; ++k) {
+      J = J.plus(U.col(k).T().times(U.col(k)).get(0));
+    }
+    problem.minimize(J);
+
+    problem.solve(new Options().withDiagnostics(true));
+  }
+}

benchmark/src/main/java/wpilib/robot/Main.java

@@ -37,6 +37,13 @@ public static void main(String... args) throws RunnerException {
     new Runner(opt).run();
   }
 
+  @Benchmark
+  @BenchmarkMode(Mode.AverageTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public void cartPole() {
+    CartPoleBenchmark.cartPole();
+  }
+
   @Benchmark
   @BenchmarkMode(Mode.AverageTime)
   @OutputTimeUnit(TimeUnit.MICROSECONDS)

wpimath/BUILD.bazel

@@ -236,6 +236,9 @@ wpilib_cc_shared_library(
 wpilib_jni_java_library(
     name = "wpimath-java",
     srcs = [":generated_java"] + glob(["src/main/java/**/*.java"]),
+    javacopts = [
+        "-Xep:UnicodeInCode:OFF",
+    ],
     maven_artifact_name = "wpimath-java",
     maven_group_id = "org.wpilib.wpimath",
     native_libs = [":wpimathjni"],
@@ -288,6 +291,7 @@ wpilib_java_junit5_test(
         "//wpiunits:wpiunits-java",
         "//wpiutil:wpiutil-java",
         "@maven//:org_ejml_ejml_core",
+        "@maven//:org_ejml_ejml_ddense",
         "@maven//:org_ejml_ejml_simple",
         "@maven//:us_hebi_quickbuf_quickbuf_runtime",
     ],

wpimath/CMakeLists.txt

@@ -7,14 +7,18 @@ include(DownloadAndCheck)
 
 file(
     GLOB wpimath_jni_src
-    src/main/native/cpp/jni/ArmFeedforwardJNI.cpp
     src/main/native/cpp/jni/DAREJNI.cpp
     src/main/native/cpp/jni/EigenJNI.cpp
-    src/main/native/cpp/jni/Ellipse2dJNI.cpp
     src/main/native/cpp/jni/Exceptions.cpp
     src/main/native/cpp/jni/StateSpaceUtilJNI.cpp
     src/main/native/cpp/jni/Transform3dJNI.cpp
     src/main/native/cpp/jni/Twist3dJNI.cpp
+    src/main/native/cpp/jni/autodiff/GradientJNI.cpp
+    src/main/native/cpp/jni/autodiff/HessianJNI.cpp
+    src/main/native/cpp/jni/autodiff/JacobianJNI.cpp
+    src/main/native/cpp/jni/autodiff/VariableJNI.cpp
+    src/main/native/cpp/jni/autodiff/VariableMatrixJNI.cpp
+    src/main/native/cpp/jni/optimization/ProblemJNI.cpp
 )
 
 # Java bindings

wpimath/src/main/java/org/wpilib/math/autodiff/ExpressionType.java

@@ -0,0 +1,47 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package org.wpilib.math.autodiff;
+
+/**
+ * Expression type.
+ *
+ * <p>Used for autodiff caching.
+ */
+public enum ExpressionType {
+  /** There is no expression. */
+  NONE(0),
+  /** The expression is a constant. */
+  CONSTANT(1),
+  /** The expression is composed of linear and lower-order operators. */
+  LINEAR(2),
+  /** The expression is composed of quadratic and lower-order operators. */
+  QUADRATIC(3),
+  /** The expression is composed of nonlinear and lower-order operators. */
+  NONLINEAR(4);
+
+  /** ExpressionType value. */
+  public final int value;
+
+  ExpressionType(int value) {
+    this.value = value;
+  }
+
+  /**
+   * Converts integer to its corresponding enum value.
+   *
+   * @param x The integer.
+   * @return The enum value.
+   */
+  public static ExpressionType fromInt(int x) {
+    return switch (x) {
+      case 0 -> ExpressionType.NONE;
+      case 1 -> ExpressionType.CONSTANT;
+      case 2 -> ExpressionType.LINEAR;
+      case 3 -> ExpressionType.QUADRATIC;
+      case 4 -> ExpressionType.NONLINEAR;
+      default -> null;
+    };
+  }
+}

wpimath/src/main/java/org/wpilib/math/autodiff/Gradient.java

@@ -0,0 +1,78 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package org.wpilib.math.autodiff;
+
+import org.ejml.simple.SimpleMatrix;
+
+/**
+ * This class calculates the gradient of a variable with respect to a vector of variables.
+ *
+ * <p>The gradient is only recomputed if the variable expression is quadratic or higher order.
+ */
+public class Gradient implements AutoCloseable {
+  private long m_handle;
+  private int m_rows;
+
+  /**
+   * Constructs a Gradient object.
+   *
+   * @param variable Variable of which to compute the gradient.
+   * @param wrt Variable with respect to which to compute the gradient.
+   */
+  public Gradient(Variable variable, Variable wrt) {
+    this(variable, new VariableMatrix(wrt));
+  }
+
+  /**
+   * Constructs a Gradient object.
+   *
+   * @param variable Variable of which to compute the gradient.
+   * @param wrt Vector of variables with respect to which to compute the gradient.
+   */
+  public Gradient(Variable variable, VariableMatrix wrt) {
+    assert wrt.cols() == 1;
+
+    m_handle = GradientJNI.create(variable.getHandle(), wrt.getHandles());
+    m_rows = wrt.rows();
+  }
+
+  /**
+   * Constructs a Gradient object.
+   *
+   * @param variable Variable of which to compute the gradient.
+   * @param wrt Vector of variables with respect to which to compute the gradient.
+   */
+  public Gradient(Variable variable, VariableBlock wrt) {
+    this(variable, new VariableMatrix(wrt));
+  }
+
+  @Override
+  public void close() {
+    if (m_handle != 0) {
+      GradientJNI.destroy(m_handle);
+      m_handle = 0;
+    }
+  }
+
+  /**
+   * Returns the gradient as a VariableMatrix.
+   *
+   * <p>This is useful when constructing optimization problems with derivatives in them.
+   *
+   * @return The gradient as a VariableMatrix.
+   */
+  public VariableMatrix get() {
+    return new VariableMatrix(m_rows, 1, GradientJNI.get(m_handle));
+  }
+
+  /**
+   * Evaluates the gradient at wrt's value.
+   *
+   * @return The gradient at wrt's value.
+   */
+  public SimpleMatrix value() {
+    return GradientJNI.value(m_handle).toSimpleMatrix(m_rows, 1);
+  }
+}

wpimath/src/main/java/org/wpilib/math/autodiff/GradientJNI.java

@@ -0,0 +1,48 @@
+// Copyright (c) FIRST and other WPILib contributors.
+// Open Source Software; you can modify and/or share it under the terms of
+// the WPILib BSD license file in the root directory of this project.
+
+package org.wpilib.math.autodiff;
+
+import org.wpilib.math.jni.WPIMathJNI;
+
+/** Gradient JNI functions. */
+final class GradientJNI extends WPIMathJNI {
+  private GradientJNI() {
+    // Utility class.
+  }
+
+  /**
+   * Constructs a Gradient object.
+   *
+   * @param variable Variable of which to compute the Gradient.
+   * @param wrt Vector of variables with respect to which to compute the Gradient.
+   */
+  static native long create(long variable, long[] wrt);
+
+  /**
+   * Destructs a Gradient.
+   *
+   * @param handle Gradient handle.
+   */
+  static native void destroy(long handle);
+
+  /**
+   * Returns the Gradient as an array of Variable handles.
+   *
+   * <p>This is useful when constructing optimization problems with derivatives in them.
+   *
+   * @param handle Gradient handle.
+   * @return The Gradient as an array of Variable handles.
+   */
+  static native long[] get(long handle);
+
+  /**
+   * Evaluates the Gradient at wrt's value.
+   *
+   * @param handle Gradient handle.
+   * @return A record containing the triplet row, column and value arrays (int[], int[], and
+   *     double[] respectively).
+   */
+  static native NativeSparseTriplets value(long handle);
+}