[wpimath] Add 3D odometry and pose estimation (#7119)

wpimath/src/main/native/include/frc/kinematics/DifferentialDriveOdometry.h

@@ -62,11 +62,7 @@ class WPILIB_DLLEXPORT DifferentialDriveOdometry
    */
   void ResetPosition(const Rotation2d& gyroAngle, units::meter_t leftDistance,
                      units::meter_t rightDistance, const Pose2d& pose) {
-    Odometry<DifferentialDriveWheelSpeeds,
-             DifferentialDriveWheelPositions>::ResetPosition(gyroAngle,
-                                                             {leftDistance,
-                                                              rightDistance},
-                                                             pose);
+    Odometry::ResetPosition(gyroAngle, {leftDistance, rightDistance}, pose);
   }
 
   /**
@@ -82,10 +78,7 @@ class WPILIB_DLLEXPORT DifferentialDriveOdometry
    */
   const Pose2d& Update(const Rotation2d& gyroAngle, units::meter_t leftDistance,
                        units::meter_t rightDistance) {
-    return Odometry<DifferentialDriveWheelSpeeds,
-                    DifferentialDriveWheelPositions>::Update(gyroAngle,
-                                                             {leftDistance,
-                                                              rightDistance});
+    return Odometry::Update(gyroAngle, {leftDistance, rightDistance});
   }
 
  private:

wpimath/src/main/native/include/frc/kinematics/DifferentialDriveOdometry3d.h

@@ -0,0 +1,87 @@
+// 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.
+
+#pragma once
+
+#include <wpi/SymbolExports.h>
+
+#include "frc/geometry/Pose2d.h"
+#include "frc/kinematics/DifferentialDriveKinematics.h"
+#include "frc/kinematics/DifferentialDriveWheelPositions.h"
+#include "frc/kinematics/DifferentialDriveWheelSpeeds.h"
+#include "frc/kinematics/Odometry3d.h"
+#include "units/length.h"
+
+namespace frc {
+/**
+ * Class for differential drive odometry. Odometry allows you to track the
+ * robot's position on the field over the course of a match using readings from
+ * 2 encoders and a gyroscope.
+ *
+ * Teams can use odometry during the autonomous period for complex tasks like
+ * path following. Furthermore, odometry can be used for latency compensation
+ * when using computer-vision systems.
+ *
+ * It is important that you reset your encoders to zero before using this class.
+ * Any subsequent pose resets also require the encoders to be reset to zero.
+ */
+class WPILIB_DLLEXPORT DifferentialDriveOdometry3d
+    : public Odometry3d<DifferentialDriveWheelSpeeds,
+                        DifferentialDriveWheelPositions> {
+ public:
+  /**
+   * Constructs a DifferentialDriveOdometry3d object.
+   *
+   * IF leftDistance and rightDistance are unspecified,
+   * You NEED to reset your encoders (to zero).
+   *
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param leftDistance The distance traveled by the left encoder.
+   * @param rightDistance The distance traveled by the right encoder.
+   * @param initialPose The starting position of the robot on the field.
+   */
+  explicit DifferentialDriveOdometry3d(const Rotation3d& gyroAngle,
+                                       units::meter_t leftDistance,
+                                       units::meter_t rightDistance,
+                                       const Pose3d& initialPose = Pose3d{});
+
+  /**
+   * Resets the robot's position on the field.
+   *
+   * IF leftDistance and rightDistance are unspecified,
+   * You NEED to reset your encoders (to zero).
+   *
+   * The gyroscope angle does not need to be reset here on the user's robot
+   * code. The library automatically takes care of offsetting the gyro angle.
+   *
+   * @param pose The position on the field that your robot is at.
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param leftDistance The distance traveled by the left encoder.
+   * @param rightDistance The distance traveled by the right encoder.
+   */
+  void ResetPosition(const Rotation3d& gyroAngle, units::meter_t leftDistance,
+                     units::meter_t rightDistance, const Pose3d& pose) {
+    Odometry3d::ResetPosition(gyroAngle, {leftDistance, rightDistance}, pose);
+  }
+
+  /**
+   * Updates the robot position on the field using distance measurements from
+   * encoders. This method is more numerically accurate than using velocities to
+   * integrate the pose and is also advantageous for teams that are using lower
+   * CPR encoders.
+   *
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param leftDistance The distance traveled by the left encoder.
+   * @param rightDistance The distance traveled by the right encoder.
+   * @return The new pose of the robot.
+   */
+  const Pose3d& Update(const Rotation3d& gyroAngle, units::meter_t leftDistance,
+                       units::meter_t rightDistance) {
+    return Odometry3d::Update(gyroAngle, {leftDistance, rightDistance});
+  }
+
+ private:
+  DifferentialDriveKinematics m_kinematicsImpl{units::meter_t{1}};
+};
+}  // namespace frc

wpimath/src/main/native/include/frc/kinematics/MecanumDriveOdometry3d.h

@@ -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.
+
+#pragma once
+
+#include <wpi/SymbolExports.h>
+#include <wpi/timestamp.h>
+
+#include "frc/geometry/Pose2d.h"
+#include "frc/kinematics/MecanumDriveKinematics.h"
+#include "frc/kinematics/MecanumDriveWheelPositions.h"
+#include "frc/kinematics/MecanumDriveWheelSpeeds.h"
+#include "frc/kinematics/Odometry3d.h"
+#include "units/time.h"
+
+namespace frc {
+
+/**
+ * Class for mecanum drive odometry. Odometry allows you to track the robot's
+ * position on the field over a course of a match using readings from your
+ * mecanum wheel encoders.
+ *
+ * Teams can use odometry during the autonomous period for complex tasks like
+ * path following. Furthermore, odometry can be used for latency compensation
+ * when using computer-vision systems.
+ */
+class WPILIB_DLLEXPORT MecanumDriveOdometry3d
+    : public Odometry3d<MecanumDriveWheelSpeeds, MecanumDriveWheelPositions> {
+ public:
+  /**
+   * Constructs a MecanumDriveOdometry3d object.
+   *
+   * @param kinematics The mecanum drive kinematics for your drivetrain.
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param wheelPositions The current distances measured by each wheel.
+   * @param initialPose The starting position of the robot on the field.
+   */
+  explicit MecanumDriveOdometry3d(
+      MecanumDriveKinematics kinematics, const Rotation3d& gyroAngle,
+      const MecanumDriveWheelPositions& wheelPositions,
+      const Pose3d& initialPose = Pose3d{});
+
+ private:
+  MecanumDriveKinematics m_kinematicsImpl;
+};
+
+}  // namespace frc

wpimath/src/main/native/include/frc/kinematics/Odometry3d.h

@@ -0,0 +1,153 @@
+// 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.
+
+#pragma once
+
+#include <wpi/SymbolExports.h>
+
+#include "frc/geometry/Pose2d.h"
+#include "frc/geometry/Pose3d.h"
+#include "frc/geometry/Rotation2d.h"
+#include "frc/geometry/Rotation3d.h"
+#include "frc/geometry/Translation2d.h"
+#include "frc/geometry/Translation3d.h"
+#include "frc/kinematics/Kinematics.h"
+
+namespace frc {
+
+/**
+ * Class for odometry. Robot code should not use this directly- Instead, use the
+ * particular type for your drivetrain (e.g., DifferentialDriveOdometry3d).
+ * Odometry allows you to track the robot's position on the field over a course
+ * of a match using readings from your wheel encoders.
+ *
+ * Teams can use odometry during the autonomous period for complex tasks like
+ * path following. Furthermore, odometry can be used for latency compensation
+ * when using computer-vision systems.
+ *
+ * @tparam WheelSpeeds Wheel speeds type.
+ * @tparam WheelPositions Wheel positions type.
+ */
+template <typename WheelSpeeds, typename WheelPositions>
+class WPILIB_DLLEXPORT Odometry3d {
+ public:
+  /**
+   * Constructs an Odometry3d object.
+   *
+   * @param kinematics The kinematics for your drivetrain.
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param wheelPositions The current distances measured by each wheel.
+   * @param initialPose The starting position of the robot on the field.
+   */
+  explicit Odometry3d(const Kinematics<WheelSpeeds, WheelPositions>& kinematics,
+                      const Rotation3d& gyroAngle,
+                      const WheelPositions& wheelPositions,
+                      const Pose3d& initialPose = Pose3d{})
+      : m_kinematics(kinematics),
+        m_pose(initialPose),
+        m_previousWheelPositions(wheelPositions) {
+    m_previousAngle = m_pose.Rotation();
+    m_gyroOffset = m_pose.Rotation() - gyroAngle;
+  }
+
+  /**
+   * Resets the robot's position on the field.
+   *
+   * The gyroscope angle does not need to be reset here on the user's robot
+   * code. The library automatically takes care of offsetting the gyro angle.
+   *
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param wheelPositions The current distances measured by each wheel.
+   * @param pose The position on the field that your robot is at.
+   */
+  void ResetPosition(const Rotation3d& gyroAngle,
+                     const WheelPositions& wheelPositions, const Pose3d& pose) {
+    m_pose = pose;
+    m_previousAngle = pose.Rotation();
+    m_gyroOffset = m_pose.Rotation() - gyroAngle;
+    m_previousWheelPositions = wheelPositions;
+  }
+
+  /**
+   * Resets the pose.
+   *
+   * @param pose The pose to reset to.
+   */
+  void ResetPose(const Pose3d& pose) {
+    m_gyroOffset = m_gyroOffset + (pose.Rotation() - m_pose.Rotation());
+    m_pose = pose;
+    m_previousAngle = pose.Rotation();
+  }
+
+  /**
+   * Resets the translation of the pose.
+   *
+   * @param translation The translation to reset to.
+   */
+  void ResetTranslation(const Translation3d& translation) {
+    m_pose = Pose3d{translation, m_pose.Rotation()};
+  }
+
+  /**
+   * Resets the rotation of the pose.
+   *
+   * @param rotation The rotation to reset to.
+   */
+  void ResetRotation(const Rotation3d& rotation) {
+    m_gyroOffset = m_gyroOffset + (rotation - m_pose.Rotation());
+    m_pose = Pose3d{m_pose.Translation(), rotation};
+    m_previousAngle = rotation;
+  }
+
+  /**
+   * Returns the position of the robot on the field.
+   * @return The pose of the robot.
+   */
+  const Pose3d& GetPose() const { return m_pose; }
+
+  /**
+   * Updates the robot's position on the field using forward kinematics and
+   * integration of the pose over time. This method takes in an angle parameter
+   * which is used instead of the angular rate that is calculated from forward
+   * kinematics, in addition to the current distance measurement at each wheel.
+   *
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param wheelPositions The current distances measured by each wheel.
+   *
+   * @return The new pose of the robot.
+   */
+  const Pose3d& Update(const Rotation3d& gyroAngle,
+                       const WheelPositions& wheelPositions) {
+    auto angle = gyroAngle + m_gyroOffset;
+    auto angle_difference =
+        (angle - m_previousAngle).GetQuaternion().ToRotationVector();
+
+    auto twist2d =
+        m_kinematics.ToTwist2d(m_previousWheelPositions, wheelPositions);
+    Twist3d twist{twist2d.dx,
+                  twist2d.dy,
+                  0_m,
+                  units::radian_t{angle_difference(0)},
+                  units::radian_t{angle_difference(1)},
+                  units::radian_t{angle_difference(2)}};
+
+    auto newPose = m_pose.Exp(twist);
+
+    m_previousWheelPositions = wheelPositions;
+    m_previousAngle = angle;
+    m_pose = {newPose.Translation(), angle};
+
+    return m_pose;
+  }
+
+ private:
+  const Kinematics<WheelSpeeds, WheelPositions>& m_kinematics;
+  Pose3d m_pose;
+
+  WheelPositions m_previousWheelPositions;
+  Rotation3d m_previousAngle;
+  Rotation3d m_gyroOffset;
+};
+
+}  // namespace frc

wpimath/src/main/native/include/frc/kinematics/SwerveDriveOdometry.h

@@ -45,9 +45,8 @@ class SwerveDriveOdometry
       SwerveDriveKinematics<NumModules> kinematics, const Rotation2d& gyroAngle,
       const wpi::array<SwerveModulePosition, NumModules>& modulePositions,
       const Pose2d& initialPose = Pose2d{})
-      : Odometry<wpi::array<SwerveModuleState, NumModules>,
-                 wpi::array<SwerveModulePosition, NumModules>>(
-            m_kinematicsImpl, gyroAngle, modulePositions, initialPose),
+      : SwerveDriveOdometry::Odometry(m_kinematicsImpl, gyroAngle,
+                                      modulePositions, initialPose),
         m_kinematicsImpl(kinematics) {
     wpi::math::MathSharedStore::ReportUsage(
         wpi::math::MathUsageId::kOdometry_SwerveDrive, 1);

wpimath/src/main/native/include/frc/kinematics/SwerveDriveOdometry3d.h

@@ -0,0 +1,70 @@
+// 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.
+
+#pragma once
+
+#include <chrono>
+#include <cstddef>
+#include <ctime>
+
+#include <wpi/SymbolExports.h>
+#include <wpi/timestamp.h>
+
+#include "Odometry3d.h"
+#include "SwerveDriveKinematics.h"
+#include "SwerveModulePosition.h"
+#include "SwerveModuleState.h"
+#include "frc/geometry/Pose2d.h"
+#include "units/time.h"
+
+namespace frc {
+
+/**
+ * Class for swerve drive odometry. Odometry allows you to track the robot's
+ * position on the field over a course of a match using readings from your
+ * swerve drive encoders and swerve azimuth encoders.
+ *
+ * Teams can use odometry during the autonomous period for complex tasks like
+ * path following. Furthermore, odometry can be used for latency compensation
+ * when using computer-vision systems.
+ */
+template <size_t NumModules>
+class SwerveDriveOdometry3d
+    : public Odometry3d<wpi::array<SwerveModuleState, NumModules>,
+                        wpi::array<SwerveModulePosition, NumModules>> {
+ public:
+  /**
+   * Constructs a SwerveDriveOdometry3d object.
+   *
+   * @param kinematics The swerve drive kinematics for your drivetrain.
+   * @param gyroAngle The angle reported by the gyroscope.
+   * @param modulePositions The wheel positions reported by each module.
+   * @param initialPose The starting position of the robot on the field.
+   */
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif  // defined(__GNUC__) && !defined(__clang__)
+  SwerveDriveOdometry3d(
+      SwerveDriveKinematics<NumModules> kinematics, const Rotation3d& gyroAngle,
+      const wpi::array<SwerveModulePosition, NumModules>& modulePositions,
+      const Pose3d& initialPose = Pose3d{})
+      : SwerveDriveOdometry3d::Odometry3d(m_kinematicsImpl, gyroAngle,
+                                          modulePositions, initialPose),
+        m_kinematicsImpl(kinematics) {
+    wpi::math::MathSharedStore::ReportUsage(
+        wpi::math::MathUsageId::kOdometry_SwerveDrive, 1);
+  }
+#if defined(__GNUC__) && !defined(__clang__)
+#pragma GCC diagnostic pop
+#endif  // defined(__GNUC__) && !defined(__clang__)
+
+ private:
+  SwerveDriveKinematics<NumModules> m_kinematicsImpl;
+};
+
+extern template class EXPORT_TEMPLATE_DECLARE(WPILIB_DLLEXPORT)
+    SwerveDriveOdometry3d<4>;
+
+}  // namespace frc