diff --git a/wpimath/src/main/java/edu/wpi/first/math/geometry/Pose3d.java b/wpimath/src/main/java/edu/wpi/first/math/geometry/Pose3d.java index 16de1dc33b..f08618b566 100644 --- 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, 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 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); diff --git a/wpimath/src/main/native/include/frc/geometry/Pose2d.h b/wpimath/src/main/native/include/frc/geometry/Pose2d.h index 2d8825caa5..d488e312a7 100644 --- 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. */ diff --git a/wpimath/src/main/native/include/frc/geometry/Pose3d.h b/wpimath/src/main/native/include/frc/geometry/Pose3d.h index 110faf7355..743ccfbc04 100644 --- 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 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 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; diff --git a/wpimath/src/test/java/edu/wpi/first/math/geometry/Pose3dTest.java b/wpimath/src/test/java/edu/wpi/first/math/geometry/Pose3dTest.java index f8fadea6a5..583e125fc0 100644 --- 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))); + } } diff --git a/wpimath/src/test/native/cpp/geometry/Pose3dTest.cpp b/wpimath/src/test/native/cpp/geometry/Pose3dTest.cpp index ee3f8fa53a..7a84718ea1 100644 --- 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()); +}