[wpimath] Add nearest to Pose2d and Translation2d (#4882)

Co-authored-by: David Vo <auscompgeek@users.noreply.github.com>

wpimath/src/main/java/edu/wpi/first/math/geometry/Pose2d.java

@@ -9,6 +9,9 @@ import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import edu.wpi.first.math.interpolation.Interpolatable;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Objects;
 
 /** Represents a 2D pose containing translational and rotational elements. */
@@ -238,6 +241,23 @@ public class Pose2d implements Interpolatable<Pose2d> {
     return new Twist2d(translationPart.getX(), translationPart.getY(), dtheta);
   }
 
+  /**
+   * Returns the nearest Pose2d 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 Pose2d from the list.
+   */
+  public Pose2d nearest(List<Pose2d> poses) {
+    return Collections.min(
+        poses,
+        Comparator.comparing(
+                (Pose2d other) -> this.getTranslation().getDistance(other.getTranslation()))
+            .thenComparing(
+                (Pose2d other) ->
+                    Math.abs(this.getRotation().minus(other.getRotation()).getRadians())));
+  }
+
   @Override
   public String toString() {
     return String.format("Pose2d(%s, %s)", m_translation, m_rotation);

wpimath/src/main/java/edu/wpi/first/math/geometry/Translation2d.java

@@ -10,6 +10,9 @@ import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import edu.wpi.first.math.MathUtil;
 import edu.wpi.first.math.interpolation.Interpolatable;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
 import java.util.Objects;
 
 /**
@@ -185,6 +188,16 @@ public class Translation2d implements Interpolatable<Translation2d> {
     return new Translation2d(m_x / scalar, m_y / scalar);
   }
 
+  /**
+   * Returns the nearest Translation2d from a list of translations.
+   *
+   * @param translations The list of translations.
+   * @return The nearest Translation2d from the list.
+   */
+  public Translation2d nearest(List<Translation2d> translations) {
+    return Collections.min(translations, Comparator.comparing(this::getDistance));
+  }
+
   @Override
   public String toString() {
     return String.format("Translation2d(X: %.2f, Y: %.2f)", m_x, m_y);

wpimath/src/main/native/cpp/geometry/Pose2d.cpp

@@ -8,6 +8,8 @@
 
 #include <wpi/json.h>
 
+#include "frc/MathUtil.h"
+
 using namespace frc;
 
 Transform2d Pose2d::operator-(const Pose2d& other) const {
@@ -67,6 +69,36 @@ Twist2d Pose2d::Log(const Pose2d& end) const {
   return {translationPart.X(), translationPart.Y(), units::radian_t{dtheta}};
 }
 
+Pose2d Pose2d::Nearest(std::span<const Pose2d> poses) const {
+  return *std::min_element(
+      poses.begin(), poses.end(), [this](const Pose2d& a, const Pose2d& 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 std::abs((this->Rotation() - a.Rotation()).Radians().value()) <
+                 std::abs((this->Rotation() - b.Rotation()).Radians().value());
+        }
+        return aDistance < bDistance;
+      });
+}
+
+Pose2d Pose2d::Nearest(std::initializer_list<Pose2d> poses) const {
+  return *std::min_element(
+      poses.begin(), poses.end(), [this](const Pose2d& a, const Pose2d& 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 std::abs((this->Rotation() - a.Rotation()).Radians().value()) <
+                 std::abs((this->Rotation() - b.Rotation()).Radians().value());
+        }
+        return aDistance < bDistance;
+      });
+}
+
 void frc::to_json(wpi::json& json, const Pose2d& pose) {
   json = wpi::json{{"translation", pose.Translation()},
                    {"rotation", pose.Rotation()}};

wpimath/src/main/native/cpp/geometry/Translation2d.cpp

@@ -23,6 +23,22 @@ bool Translation2d::operator==(const Translation2d& other) const {
          units::math::abs(m_y - other.m_y) < 1E-9_m;
 }
 
+Translation2d Translation2d::Nearest(
+    std::span<const Translation2d> translations) const {
+  return *std::min_element(translations.begin(), translations.end(),
+                           [this](Translation2d a, Translation2d b) {
+                             return this->Distance(a) < this->Distance(b);
+                           });
+}
+
+Translation2d Translation2d::Nearest(
+    std::initializer_list<Translation2d> translations) const {
+  return *std::min_element(translations.begin(), translations.end(),
+                           [this](Translation2d a, Translation2d b) {
+                             return this->Distance(a) < this->Distance(b);
+                           });
+}
+
 void frc::to_json(wpi::json& json, const Translation2d& translation) {
   json =
       wpi::json{{"x", translation.X().value()}, {"y", translation.Y().value()}};

wpimath/src/main/native/include/frc/geometry/Pose2d.h

@@ -4,6 +4,9 @@
 
 #pragma once
 
+#include <initializer_list>
+#include <span>
+
 #include <wpi/SymbolExports.h>
 
 #include "Transform2d.h"
@@ -176,6 +179,20 @@ class WPILIB_DLLEXPORT Pose2d {
    */
   Twist2d Log(const Pose2d& end) const;
 
+  /**
+   * Returns the nearest Pose2d from a collection of poses
+   * @param poses The collection of poses.
+   * @return The nearest Pose2d from the collection.
+   */
+  Pose2d Nearest(std::span<const Pose2d> poses) const;
+
+  /**
+   * Returns the nearest Pose2d from a collection of poses
+   * @param poses The collection of poses.
+   * @return The nearest Pose2d from the collection.
+   */
+  Pose2d Nearest(std::initializer_list<Pose2d> poses) const;
+
  private:
   Translation2d m_translation;
   Rotation2d m_rotation;

wpimath/src/main/native/include/frc/geometry/Translation2d.h

@@ -4,6 +4,9 @@
 
 #pragma once
 
+#include <initializer_list>
+#include <span>
+
 #include <wpi/SymbolExports.h>
 
 #include "Rotation2d.h"
@@ -170,6 +173,21 @@ class WPILIB_DLLEXPORT Translation2d {
    */
   bool operator==(const Translation2d& other) const;
 
+  /**
+   * Returns the nearest Translation2d from a collection of translations
+   * @param translations The collection of translations.
+   * @return The nearest Translation2d from the collection.
+   */
+  Translation2d Nearest(std::span<const Translation2d> translations) const;
+
+  /**
+   * Returns the nearest Translation2d from a collection of translations
+   * @param translations The collection of translations.
+   * @return The nearest Translation2d from the collection.
+   */
+  Translation2d Nearest(
+      std::initializer_list<Translation2d> translations) const;
+
  private:
   units::meter_t m_x = 0_m;
   units::meter_t m_y = 0_m;