[wpimath] Add nearest() method to Pose3d (mirroring Pose2d) (#8010)

2026-07-02 02:51:42 +00:00 · 2025-06-05 22:11:49 -07:00
parent 7a3df6175e
commit 075cc4a20f
5 changed files with 254 additions and 2 deletions
--- a/wpimath/src/main/java/edu/wpi/first/math/geometry/Pose3d.java
+++ b/wpimath/src/main/java/edu/wpi/first/math/geometry/Pose3d.java
@@ -21,6 +21,9 @@ import edu.wpi.first.math.numbers.N4;
 import edu.wpi.first.units.measure.Distance;
 import edu.wpi.first.util.protobuf.ProtobufSerializable;
 import edu.wpi.first.util.struct.StructSerializable;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Objects;

 /** Represents a 3D pose containing translational and rotational elements. */
@@ -396,6 +399,22 @@ public class Pose3d implements Interpolatable<Pose3d>, ProtobufSerializable, Str
    return new Pose2d(m_translation.toTranslation2d(), m_rotation.toRotation2d());
  }

+  /**
+   * Returns the nearest Pose3d from a list of poses. If two or more poses in the list have the same
+   * distance from this pose, return the one with the closest rotation component.
+   *
+   * @param poses The list of poses to find the nearest.
+   * @return The nearest Pose3d from the list.
+   */
+  public Pose3d nearest(List<Pose3d> poses) {
+    return Collections.min(
+        poses,
+        Comparator.comparing(
+                (Pose3d other) -> this.getTranslation().getDistance(other.getTranslation()))
+            .thenComparing(
+                (Pose3d other) -> this.getRotation().minus(other.getRotation()).getAngle()));
+  }
+
  @Override
  public String toString() {
    return String.format("Pose3d(%s, %s)", m_translation, m_rotation);
--- a/wpimath/src/main/native/include/frc/geometry/Pose2d.h
+++ b/wpimath/src/main/native/include/frc/geometry/Pose2d.h
@@ -242,7 +242,11 @@ class WPILIB_DLLEXPORT Pose2d {
  }

  /**
-   * Returns the nearest Pose2d from a collection of poses
+   * Returns the nearest Pose2d from a collection of poses.
+   *
+   * If two or more poses in the collection have the same distance from this
+   * pose, return the one with the closest rotation component.
+   *
   * @param poses The collection of poses.
   * @return The nearest Pose2d from the collection.
   */
@@ -264,7 +268,11 @@ class WPILIB_DLLEXPORT Pose2d {
  }

  /**
-   * Returns the nearest Pose2d from a collection of poses
+   * Returns the nearest Pose2d from a collection of poses.
+   *
+   * If two or more poses in the collection have the same distance from this
+   * pose, return the one with the closest rotation component.
+   *
   * @param poses The collection of poses.
   * @return The nearest Pose2d from the collection.
   */
--- a/wpimath/src/main/native/include/frc/geometry/Pose3d.h
+++ b/wpimath/src/main/native/include/frc/geometry/Pose3d.h
@@ -270,6 +270,56 @@ class WPILIB_DLLEXPORT Pose3d {
    return Pose2d{m_translation.X(), m_translation.Y(), m_rotation.Z()};
  }

+  /**
+   * Returns the nearest Pose3d from a collection of poses.
+   *
+   * If two or more poses in the collection have the same distance from this
+   * pose, return the one with the closest rotation component.
+   *
+   * @param poses The collection of poses.
+   * @return The nearest Pose3d from the collection.
+   */
+  constexpr Pose3d Nearest(std::span<const Pose3d> poses) const {
+    return *std::min_element(
+        poses.begin(), poses.end(), [this](const Pose3d& a, const Pose3d& b) {
+          auto aDistance = this->Translation().Distance(a.Translation());
+          auto bDistance = this->Translation().Distance(b.Translation());
+
+          // If the distances are equal sort by difference in rotation
+          if (aDistance == bDistance) {
+            return gcem::abs(
+                       (this->Rotation() - a.Rotation()).Angle().value()) <
+                   gcem::abs((this->Rotation() - b.Rotation()).Angle().value());
+          }
+          return aDistance < bDistance;
+        });
+  }
+
+  /**
+   * Returns the nearest Pose3d from a collection of poses.
+   *
+   * If two or more poses in the collection have the same distance from this
+   * pose, return the one with the closest rotation component.
+   *
+   * @param poses The collection of poses.
+   * @return The nearest Pose3d from the collection.
+   */
+  constexpr Pose3d Nearest(std::initializer_list<Pose3d> poses) const {
+    return *std::min_element(
+        poses.begin(), poses.end(), [this](const Pose3d& a, const Pose3d& b) {
+          auto aDistance = this->Translation().Distance(a.Translation());
+          auto bDistance = this->Translation().Distance(b.Translation());
+
+          // If the distances are equal sort by difference in rotation
+          if (aDistance == bDistance) {
+            return gcem::abs(
+                       (this->Rotation() - a.Rotation()).Angle().value()) <
+                   gcem::abs((this->Rotation() - b.Rotation()).Angle().value());
+          }
+          return aDistance < bDistance;
+        });
+  }
+
 private:
  Translation3d m_translation;
  Rotation3d m_rotation;
--- a/wpimath/src/test/java/edu/wpi/first/math/geometry/Pose3dTest.java
+++ b/wpimath/src/test/java/edu/wpi/first/math/geometry/Pose3dTest.java
@@ -324,4 +324,71 @@ class Pose3dTest {
          () -> assertFalse(((Double) twist.rz).isNaN()));
    }
  }
+
+  @Test
+  void testNearest() {
+    var origin = Pose3d.kZero;
+
+    // Distance sort
+    // poses are in order of closest to farthest away from the origin at various positions in 3D
+    // space.
+    final var pose1 = new Pose3d(1, 0, 0, Rotation3d.kZero);
+    final var pose2 = new Pose3d(0, 2, 0, Rotation3d.kZero);
+    final var pose3 = new Pose3d(0, 0, 3, Rotation3d.kZero);
+    final var pose4 = new Pose3d(2, 2, 2, Rotation3d.kZero);
+    final var pose5 = new Pose3d(3, 3, 3, Rotation3d.kZero);
+
+    assertEquals(pose3, origin.nearest(List.of(pose5, pose3, pose4)));
+    assertEquals(pose1, origin.nearest(List.of(pose1, pose2, pose3)));
+    assertEquals(pose2, origin.nearest(List.of(pose4, pose2, pose3)));
+
+    // Rotation component sort (when distance is the same)
+    // Use the same translation to avoid distance differences
+    final var translation = new Translation3d(1, 0, 0);
+
+    final var poseA = new Pose3d(translation, Rotation3d.kZero); // No rotation
+    final var poseB = new Pose3d(translation, new Rotation3d(Math.toRadians(30), 0, 0));
+    final var poseC = new Pose3d(translation, new Rotation3d(0, Math.toRadians(45), 0));
+    final var poseD = new Pose3d(translation, new Rotation3d(0, 0, Math.toRadians(90)));
+    final var poseE = new Pose3d(translation, new Rotation3d(Math.toRadians(180), 0, 0));
+
+    assertEquals(
+        poseA, new Pose3d(0, 0, 0, Rotation3d.kZero).nearest(List.of(poseA, poseB, poseD)));
+    assertEquals(
+        poseB,
+        new Pose3d(0, 0, 0, new Rotation3d(Math.toRadians(25), 0, 0))
+            .nearest(List.of(poseB, poseC, poseD)));
+    assertEquals(
+        poseC,
+        new Pose3d(0, 0, 0, new Rotation3d(0, Math.toRadians(50), 0))
+            .nearest(List.of(poseB, poseC, poseD)));
+    assertEquals(
+        poseD,
+        new Pose3d(0, 0, 0, new Rotation3d(0, 0, Math.toRadians(85)))
+            .nearest(List.of(poseA, poseC, poseD)));
+    assertEquals(
+        poseE,
+        new Pose3d(0, 0, 0, new Rotation3d(Math.toRadians(170), 0, 0))
+            .nearest(List.of(poseA, poseD, poseE)));
+
+    // Test with complex 3D rotations (combining roll, pitch, yaw)
+    final var complexPose1 =
+        new Pose3d(
+            translation,
+            new Rotation3d(Math.toRadians(45), Math.toRadians(30), Math.toRadians(60)));
+    final var complexPose2 =
+        new Pose3d(
+            translation,
+            new Rotation3d(Math.toRadians(90), Math.toRadians(45), Math.toRadians(90)));
+    final var complexPose3 =
+        new Pose3d(
+            translation,
+            new Rotation3d(Math.toRadians(10), Math.toRadians(15), Math.toRadians(20)));
+
+    assertEquals(
+        complexPose3,
+        new Pose3d(
+                0, 0, 0, new Rotation3d(Math.toRadians(5), Math.toRadians(10), Math.toRadians(15)))
+            .nearest(List.of(complexPose1, complexPose2, complexPose3)));
+  }
 }
--- a/wpimath/src/test/native/cpp/geometry/Pose3dTest.cpp
+++ b/wpimath/src/test/native/cpp/geometry/Pose3dTest.cpp
@@ -221,3 +221,111 @@ TEST(Pose3dTest, TwistNaN) {
    EXPECT_FALSE(std::isnan(twist.rz.value()));
  }
 }
+
+TEST(Pose3dTest, Nearest) {
+  const Pose3d origin{0_m, 0_m, 0_m, Rotation3d{}};
+
+  // Distance sort
+  // poses are in order of closest to farthest away from the origin at
+  // various positions in 3D space.
+  const Pose3d pose1{1_m, 0_m, 0_m, Rotation3d{}};
+  const Pose3d pose2{0_m, 2_m, 0_m, Rotation3d{}};
+  const Pose3d pose3{0_m, 0_m, 3_m, Rotation3d{}};
+  const Pose3d pose4{2_m, 2_m, 2_m, Rotation3d{}};
+  const Pose3d pose5{3_m, 3_m, 3_m, Rotation3d{}};
+
+  EXPECT_DOUBLE_EQ(pose3.X().value(),
+                   origin.Nearest({pose5, pose3, pose4}).X().value());
+  EXPECT_DOUBLE_EQ(pose3.Y().value(),
+                   origin.Nearest({pose5, pose3, pose4}).Y().value());
+  EXPECT_DOUBLE_EQ(pose3.Z().value(),
+                   origin.Nearest({pose5, pose3, pose4}).Z().value());
+
+  EXPECT_DOUBLE_EQ(pose1.X().value(),
+                   origin.Nearest({pose1, pose2, pose3}).X().value());
+  EXPECT_DOUBLE_EQ(pose1.Y().value(),
+                   origin.Nearest({pose1, pose2, pose3}).Y().value());
+  EXPECT_DOUBLE_EQ(pose1.Z().value(),
+                   origin.Nearest({pose1, pose2, pose3}).Z().value());
+
+  EXPECT_DOUBLE_EQ(pose2.X().value(),
+                   origin.Nearest({pose4, pose2, pose3}).X().value());
+  EXPECT_DOUBLE_EQ(pose2.Y().value(),
+                   origin.Nearest({pose4, pose2, pose3}).Y().value());
+  EXPECT_DOUBLE_EQ(pose2.Z().value(),
+                   origin.Nearest({pose4, pose2, pose3}).Z().value());
+
+  // Rotation component sort (when distance is the same)
+  // Use the same translation to avoid distance differences
+  const Translation3d translation{1_m, 0_m, 0_m};
+
+  const Pose3d poseA{translation, Rotation3d{}};  // No rotation
+  const Pose3d poseB{translation, Rotation3d{30_deg, 0_deg, 0_deg}};
+  const Pose3d poseC{translation, Rotation3d{0_deg, 45_deg, 0_deg}};
+  const Pose3d poseD{translation, Rotation3d{0_deg, 0_deg, 90_deg}};
+  const Pose3d poseE{translation, Rotation3d{180_deg, 0_deg, 0_deg}};
+
+  auto result1 =
+      Pose3d{0_m, 0_m, 0_m, Rotation3d{}}.Nearest({poseA, poseB, poseD});
+  EXPECT_DOUBLE_EQ(poseA.Rotation().X().value(),
+                   result1.Rotation().X().value());
+  EXPECT_DOUBLE_EQ(poseA.Rotation().Y().value(),
+                   result1.Rotation().Y().value());
+  EXPECT_DOUBLE_EQ(poseA.Rotation().Z().value(),
+                   result1.Rotation().Z().value());
+
+  auto result2 =
+      Pose3d{0_m, 0_m, 0_m, Rotation3d{25_deg, 0_deg, 0_deg}}.Nearest(
+          {poseB, poseC, poseD});
+  EXPECT_DOUBLE_EQ(poseB.Rotation().X().value(),
+                   result2.Rotation().X().value());
+  EXPECT_DOUBLE_EQ(poseB.Rotation().Y().value(),
+                   result2.Rotation().Y().value());
+  EXPECT_DOUBLE_EQ(poseB.Rotation().Z().value(),
+                   result2.Rotation().Z().value());
+
+  auto result3 =
+      Pose3d{0_m, 0_m, 0_m, Rotation3d{0_deg, 50_deg, 0_deg}}.Nearest(
+          {poseB, poseC, poseD});
+  EXPECT_DOUBLE_EQ(poseC.Rotation().X().value(),
+                   result3.Rotation().X().value());
+  EXPECT_DOUBLE_EQ(poseC.Rotation().Y().value(),
+                   result3.Rotation().Y().value());
+  EXPECT_DOUBLE_EQ(poseC.Rotation().Z().value(),
+                   result3.Rotation().Z().value());
+
+  auto result4 =
+      Pose3d{0_m, 0_m, 0_m, Rotation3d{0_deg, 0_deg, 85_deg}}.Nearest(
+          {poseA, poseC, poseD});
+  EXPECT_DOUBLE_EQ(poseD.Rotation().X().value(),
+                   result4.Rotation().X().value());
+  EXPECT_DOUBLE_EQ(poseD.Rotation().Y().value(),
+                   result4.Rotation().Y().value());
+  EXPECT_DOUBLE_EQ(poseD.Rotation().Z().value(),
+                   result4.Rotation().Z().value());
+
+  auto result5 =
+      Pose3d{0_m, 0_m, 0_m, Rotation3d{170_deg, 0_deg, 0_deg}}.Nearest(
+          {poseA, poseD, poseE});
+  EXPECT_DOUBLE_EQ(poseE.Rotation().X().value(),
+                   result5.Rotation().X().value());
+  EXPECT_DOUBLE_EQ(poseE.Rotation().Y().value(),
+                   result5.Rotation().Y().value());
+  EXPECT_DOUBLE_EQ(poseE.Rotation().Z().value(),
+                   result5.Rotation().Z().value());
+
+  // Test with complex 3D rotations (combining roll, pitch, yaw)
+  const Pose3d complexPose1{translation, Rotation3d{45_deg, 30_deg, 60_deg}};
+  const Pose3d complexPose2{translation, Rotation3d{90_deg, 45_deg, 90_deg}};
+  const Pose3d complexPose3{translation, Rotation3d{10_deg, 15_deg, 20_deg}};
+
+  auto complexResult =
+      Pose3d{0_m, 0_m, 0_m, Rotation3d{5_deg, 10_deg, 15_deg}}.Nearest(
+          {complexPose1, complexPose2, complexPose3});
+  EXPECT_DOUBLE_EQ(complexPose3.Rotation().X().value(),
+                   complexResult.Rotation().X().value());
+  EXPECT_DOUBLE_EQ(complexPose3.Rotation().Y().value(),
+                   complexResult.Rotation().Y().value());
+  EXPECT_DOUBLE_EQ(complexPose3.Rotation().Z().value(),
+                   complexResult.Rotation().Z().value());
+}