Add JNI for Pose3d exp and log

wpimath/src/main/java/edu/wpi/first/math/WPIMathJNI.java

@@ -76,6 +76,30 @@ public static native void dare(
    */
   public static native void pow(double[] src, int rows, double exponent, double[] dst);
 
+  /**
+   * Obtain a Pose3d from a (constant curvature) velocity.
+   *
+   * <p>A Pose3d is represented by [t_x, t_y, t_z, q_w, q_x, q_y, q_z] and a Twist3d is represented
+   * by [dx, dy, dz, rx, ry, rz].
+   *
+   * @param start The starting Pose3d as a double array.
+   * @param twist The twist to apply as a double array.
+   * @return The new pose as a double array.
+   */
+  public static native double[] expPose3d(double[] start, double[] twist);
+
+  /**
+   * Returns a Twist3d that maps this pose to the end pose.
+   *
+   * <p>A Pose3d is represented by [t_x, t_y, t_z, q_w, q_x, q_y, q_z] and a Twist3d is represented
+   * by [dx, dy, dz, rx, ry, rz].
+   *
+   * @param start The starting Pose3d as a double array.
+   * @param end The ending Pose3d as a double array.
+   * @return The twist that maps start to end as a double array.
+   */
+  public static native double[] logPose3d(double[] start, double[] end);
+
   /**
    * Returns true if (A, B) is a stabilizable pair.
    *

wpimath/src/main/native/cpp/jni/WPIMathJNI.cpp

@@ -14,6 +14,7 @@
 #include "Eigen/QR"
 #include "edu_wpi_first_math_WPIMathJNI.h"
 #include "frc/DARE.h"
+#include "frc/geometry/Pose3d.h"
 #include "frc/trajectory/TrajectoryUtil.h"
 #include "unsupported/Eigen/MatrixFunctions"
 
@@ -57,6 +58,55 @@ bool IsStabilizable(const Eigen::Ref<const Eigen::MatrixXd>& A,
 
 }  // namespace
 
+frc::Pose3d Pose3dFromJDoubleArray(JNIEnv* env, jdoubleArray arr) {
+  jdouble* nativeArray = env->GetDoubleArrayElements(arr, nullptr);
+  frc::Translation3d translation{units::meter_t{nativeArray[0]},
+                                 units::meter_t{nativeArray[1]},
+                                 units::meter_t{nativeArray[2]}};
+  frc::Quaternion quaternion{nativeArray[3], nativeArray[4], nativeArray[5],
+                             nativeArray[6]};
+  env->ReleaseDoubleArrayElements(arr, nativeArray, 0);
+  return {translation, frc::Rotation3d{quaternion}};
+}
+
+frc::Twist3d Twist3dFromJDoubleArray(JNIEnv* env, jdoubleArray arr) {
+  jdouble* nativeArray = env->GetDoubleArrayElements(arr, nullptr);
+  frc::Twist3d twist = {
+      units::meter_t{nativeArray[0]},  units::meter_t{nativeArray[1]},
+      units::meter_t{nativeArray[2]},  units::radian_t{nativeArray[3]},
+      units::radian_t{nativeArray[4]}, units::radian_t{nativeArray[5]}};
+  env->ReleaseDoubleArrayElements(arr, nativeArray, 0);
+  return twist;
+}
+
+jdoubleArray MakeJDoubleArray(JNIEnv* env, const frc::Pose3d& pose) {
+  const frc::Quaternion& quaternion = pose.Rotation().GetQuaternion();
+  jdoubleArray results = env->NewDoubleArray(7);
+  jdouble* buffer = env->GetDoubleArrayElements(results, nullptr);
+  buffer[0] = pose.X().value();
+  buffer[1] = pose.Y().value();
+  buffer[2] = pose.Z().value();
+  buffer[3] = quaternion.W();
+  buffer[4] = quaternion.X();
+  buffer[5] = quaternion.Y();
+  buffer[6] = quaternion.Z();
+  env->ReleaseDoubleArrayElements(results, buffer, 0);
+  return results;
+}
+
+jdoubleArray MakeJDoubleArray(JNIEnv* env, const frc::Twist3d& twist) {
+  jdoubleArray results = env->NewDoubleArray(6);
+  jdouble* buffer = env->GetDoubleArrayElements(results, nullptr);
+  buffer[0] = twist.dx.value();
+  buffer[1] = twist.dy.value();
+  buffer[2] = twist.dz.value();
+  buffer[3] = twist.rx.value();
+  buffer[4] = twist.ry.value();
+  buffer[5] = twist.rz.value();
+  env->ReleaseDoubleArrayElements(results, buffer, 0);
+  return results;
+}
+
 std::vector<double> GetElementsFromTrajectory(
     const frc::Trajectory& trajectory) {
   std::vector<double> elements;
@@ -187,6 +237,40 @@ Java_edu_wpi_first_math_WPIMathJNI_pow
   env->SetDoubleArrayRegion(dst, 0, rows * rows, Apow.data());
 }
 
+/*
+ * Class:     edu_wpi_first_math_WPIMathJNI
+ * Method:    expPose3d
+ * Signature: ([D[D)[D
+ */
+JNIEXPORT jdoubleArray JNICALL
+Java_edu_wpi_first_math_WPIMathJNI_expPose3d
+  (JNIEnv* env, jclass, jdoubleArray startPose, jdoubleArray twist)
+{
+  frc::Pose3d startPoseCpp = Pose3dFromJDoubleArray(env, startPose);
+  frc::Twist3d twistCpp = Twist3dFromJDoubleArray(env, twist);
+
+  frc::Pose3d result = startPoseCpp.Exp(twistCpp);
+
+  return MakeJDoubleArray(env, result);
+}
+
+/*
+ * Class:     edu_wpi_first_math_WPIMathJNI
+ * Method:    logPose3d
+ * Signature: ([D[D)[D
+ */
+JNIEXPORT jdoubleArray JNICALL
+Java_edu_wpi_first_math_WPIMathJNI_logPose3d
+  (JNIEnv* env, jclass, jdoubleArray start, jdoubleArray end)
+{
+  frc::Pose3d startCpp = Pose3dFromJDoubleArray(env, start);
+  frc::Pose3d endCpp = Pose3dFromJDoubleArray(env, end);
+
+  frc::Twist3d result = startCpp.Log(endCpp);
+
+  return MakeJDoubleArray(env, result);
+}
+
 /*
  * Class:     edu_wpi_first_math_WPIMathJNI
  * Method:    isStabilizable

wpimath/src/test/java/edu/wpi/first/math/WPIMathJNITest.java

@@ -5,12 +5,63 @@
 package edu.wpi.first.math;
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
+import edu.wpi.first.math.geometry.Pose3d;
+import edu.wpi.first.math.geometry.Quaternion;
+import edu.wpi.first.math.geometry.Rotation3d;
+import edu.wpi.first.math.geometry.Twist3d;
 import org.junit.jupiter.api.Test;
 
 public class WPIMathJNITest {
+  private static double[] toDoubleArray(Pose3d pose) {
+    Quaternion quaternion = pose.getRotation().getQuaternion();
+    return new double[] {
+      pose.getX(),
+      pose.getY(),
+      pose.getZ(),
+      quaternion.getW(),
+      quaternion.getX(),
+      quaternion.getY(),
+      quaternion.getZ()
+    };
+  }
+
+  private static double[] toDoubleArray(Twist3d twist) {
+    return new double[] {twist.dx, twist.dy, twist.dz, twist.rx, twist.ry, twist.rz};
+  }
+
+  private static Pose3d pose3dFromDoubleArray(double[] arr) {
+    return new Pose3d(
+        arr[0], arr[1], arr[2], new Rotation3d(new Quaternion(arr[3], arr[4], arr[5], arr[6])));
+  }
+
+  private static Twist3d twist3dFromDoubleArray(double[] arr) {
+    return new Twist3d(arr[0], arr[1], arr[2], arr[3], arr[4], arr[5]);
+  }
+
   @Test
   public void testLink() {
     assertDoesNotThrow(WPIMathJNI::forceLoad);
   }
+
+  @Test
+  public void testExpPose3d() {
+    final var start = new Pose3d(1, 0, 0, new Rotation3d(0, -Math.PI / 2, 0));
+    final var twist = new Twist3d(0, 1, 0, Math.PI, 0, 0);
+    final var expected = start.exp(twist);
+    final var result =
+        pose3dFromDoubleArray(WPIMathJNI.expPose3d(toDoubleArray(start), toDoubleArray(twist)));
+    assertEquals(expected, result);
+  }
+
+  @Test
+  public void testLogPose3d() {
+    final var start = new Pose3d(1, 0, 0, new Rotation3d(0, -Math.PI / 2, 0));
+    final var end = new Pose3d(0, 1, 0, new Rotation3d(Math.PI / 2, 0, 0));
+    final var expected = start.log(end);
+    final var result =
+        twist3dFromDoubleArray(WPIMathJNI.logPose3d(toDoubleArray(start), toDoubleArray(end)));
+    assertEquals(expected, result);
+  }
 }