[wpimath] Add support for swerve joystick normalization (#4516)

wpimath/src/main/java/edu/wpi/first/math/kinematics/SwerveDriveKinematics.java

@@ -237,4 +237,48 @@ public class SwerveDriveKinematics {
       }
     }
   }
+
+  /**
+   * Renormalizes the wheel speeds if any individual speed is above the specified maximum, as well
+   * as getting rid of joystick saturation at edges of joystick.
+   *
+   * <p>Sometimes, after inverse kinematics, the requested speed from one or more modules may be
+   * above the max attainable speed for the driving motor on that module. To fix this issue, one can
+   * reduce all the wheel speeds to make sure that all requested module speeds are at-or-below the
+   * absolute threshold, while maintaining the ratio of speeds between modules.
+   *
+   * @param moduleStates Reference to array of module states. The array will be mutated with the
+   *     normalized speeds!
+   * @param currentChassisSpeed The current speed of the robot
+   * @param attainableMaxModuleSpeedMetersPerSecond The absolute max speed that a module can reach
+   * @param attainableMaxTranslationalSpeedMetersPerSecond The absolute max speed that your robot
+   *     can reach while translating
+   * @param attainableMaxRotationalVelocityRadiansPerSecond The absolute max speed the robot can
+   *     reach while rotating
+   */
+  public static void desaturateWheelSpeeds(
+      SwerveModuleState[] moduleStates,
+      ChassisSpeeds currentChassisSpeed,
+      double attainableMaxModuleSpeedMetersPerSecond,
+      double attainableMaxTranslationalSpeedMetersPerSecond,
+      double attainableMaxRotationalVelocityRadiansPerSecond) {
+    double realMaxSpeed = Collections.max(Arrays.asList(moduleStates)).speedMetersPerSecond;
+
+    if (attainableMaxTranslationalSpeedMetersPerSecond == 0
+        || attainableMaxRotationalVelocityRadiansPerSecond == 0
+        || realMaxSpeed == 0) {
+      return;
+    }
+    double translationalK =
+        Math.hypot(currentChassisSpeed.vxMetersPerSecond, currentChassisSpeed.vyMetersPerSecond)
+            / attainableMaxTranslationalSpeedMetersPerSecond;
+    double rotationalK =
+        Math.abs(currentChassisSpeed.omegaRadiansPerSecond)
+            / attainableMaxRotationalVelocityRadiansPerSecond;
+    double k = Math.max(translationalK, rotationalK);
+    double scale = Math.min(k * attainableMaxModuleSpeedMetersPerSecond / realMaxSpeed, 1);
+    for (SwerveModuleState moduleState : moduleStates) {
+      moduleState.speedMetersPerSecond *= scale;
+    }
+  }
 }

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

@@ -215,6 +215,34 @@ class SwerveDriveKinematics {
       wpi::array<SwerveModuleState, NumModules>* moduleStates,
       units::meters_per_second_t attainableMaxSpeed);
 
+  /**
+   * Renormalizes the wheel speeds if any individual speed is above the
+   * specified maximum, as well as getting rid of joystick saturation at edges
+   * of joystick.
+   *
+   * Sometimes, after inverse kinematics, the requested speed
+   * from one or more modules may be above the max attainable speed for the
+   * driving motor on that module. To fix this issue, one can reduce all the
+   * wheel speeds to make sure that all requested module speeds are at-or-below
+   * the absolute threshold, while maintaining the ratio of speeds between
+   * modules.
+   *
+   * @param moduleStates Reference to array of module states. The array will be
+   * mutated with the normalized speeds!
+   * @param currentChassisSpeed Current speed of the robot
+   * @param attainableMaxModuleSpeed The absolute max speed a module can reach
+   * @param attainableMaxRobotTranslationSpeed The absolute max speed the robot
+   * can reach while translating
+   * @param attainableMaxRobotRotationSpeed The absolute max speed the robot can
+   * reach while rotating
+   */
+  static void DesaturateWheelSpeeds(
+      wpi::array<SwerveModuleState, NumModules>* moduleStates,
+      ChassisSpeeds currentChassisSpeed,
+      units::meters_per_second_t attainableMaxModuleSpeed,
+      units::meters_per_second_t attainableMaxRobotTranslationSpeed,
+      units::radians_per_second_t attainableMaxRobotRotationSpeed);
+
  private:
   mutable Matrixd<NumModules * 2, 3> m_inverseKinematics;
   Eigen::HouseholderQR<Matrixd<NumModules * 2, 3>> m_forwardKinematics;

wpimath/src/main/native/include/frc/kinematics/SwerveDriveKinematics.inc

@@ -7,6 +7,7 @@
 #include <algorithm>
 #include <utility>
 
+#include "frc/kinematics/ChassisSpeeds.h"
 #include "frc/kinematics/SwerveDriveKinematics.h"
 #include "units/math.h"
 
@@ -147,4 +148,41 @@ void SwerveDriveKinematics<NumModules>::DesaturateWheelSpeeds(
   }
 }
 
+template <size_t NumModules>
+void SwerveDriveKinematics<NumModules>::DesaturateWheelSpeeds(
+    wpi::array<SwerveModuleState, NumModules>* moduleStates,
+    ChassisSpeeds currentChassisSpeed,
+    units::meters_per_second_t attainableMaxModuleSpeed,
+    units::meters_per_second_t attainableMaxRobotTranslationSpeed,
+    units::radians_per_second_t attainableMaxRobotRotationSpeed) {
+  auto& states = *moduleStates;
+
+  auto realMaxSpeed = std::max_element(states.begin(), states.end(),
+                                       [](const auto& a, const auto& b) {
+                                         return units::math::abs(a.speed) <
+                                                units::math::abs(b.speed);
+                                       })
+                          ->speed;
+
+  if (attainableMaxRobotTranslationSpeed == 0_mps ||
+      attainableMaxRobotRotationSpeed == 0_rad_per_s || realMaxSpeed == 0_mps) {
+    return;
+  }
+
+  auto translationalK =
+      units::math::hypot(currentChassisSpeed.vx, currentChassisSpeed.vy) /
+      attainableMaxRobotTranslationSpeed;
+
+  auto rotationalK = units::math::abs(currentChassisSpeed.omega) /
+                     attainableMaxRobotRotationSpeed;
+
+  auto k = units::math::max(translationalK, rotationalK);
+
+  auto scale = units::math::min(k * attainableMaxModuleSpeed / realMaxSpeed,
+                                units::scalar_t(1));
+  for (auto& module : states) {
+    module.speed = module.speed * scale;
+  }
+}
+
 }  // namespace frc

wpimath/src/test/java/edu/wpi/first/math/kinematics/SwerveDriveKinematicsTest.java

@@ -351,4 +351,24 @@ class SwerveDriveKinematicsTest {
         () -> assertEquals(4.0 * factor, arr[2].speedMetersPerSecond, kEpsilon),
         () -> assertEquals(7.0 * factor, arr[3].speedMetersPerSecond, kEpsilon));
   }
+
+  @Test
+  void testDesaturateSmooth() {
+    SwerveModuleState fl = new SwerveModuleState(5, new Rotation2d());
+    SwerveModuleState fr = new SwerveModuleState(6, new Rotation2d());
+    SwerveModuleState bl = new SwerveModuleState(4, new Rotation2d());
+    SwerveModuleState br = new SwerveModuleState(7, new Rotation2d());
+
+    SwerveModuleState[] arr = {fl, fr, bl, br};
+    SwerveDriveKinematics.desaturateWheelSpeeds(
+        arr, m_kinematics.toChassisSpeeds(arr), 5.5, 5.5, 3.5);
+
+    double factor = 5.5 / 7.0;
+
+    assertAll(
+        () -> assertEquals(5.0 * factor, arr[0].speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(6.0 * factor, arr[1].speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(4.0 * factor, arr[2].speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(7.0 * factor, arr[3].speedMetersPerSecond, kEpsilon));
+  }
 }

wpimath/src/test/native/cpp/kinematics/SwerveDriveKinematicsTest.cpp

@@ -256,3 +256,21 @@ TEST_F(SwerveDriveKinematicsTest, Desaturate) {
   EXPECT_NEAR(arr[2].speed.value(), 4.0 * kFactor, kEpsilon);
   EXPECT_NEAR(arr[3].speed.value(), 7.0 * kFactor, kEpsilon);
 }
+
+TEST_F(SwerveDriveKinematicsTest, DesaturateSmooth) {
+  SwerveModuleState state1{5.0_mps, 0_deg};
+  SwerveModuleState state2{6.0_mps, 0_deg};
+  SwerveModuleState state3{4.0_mps, 0_deg};
+  SwerveModuleState state4{7.0_mps, 0_deg};
+
+  wpi::array<SwerveModuleState, 4> arr{state1, state2, state3, state4};
+  SwerveDriveKinematics<4>::DesaturateWheelSpeeds(
+      &arr, m_kinematics.ToChassisSpeeds(arr), 5.5_mps, 5.5_mps, 3.5_rad_per_s);
+
+  double kFactor = 5.5 / 7.0;
+
+  EXPECT_NEAR(arr[0].speed.value(), 5.0 * kFactor, kEpsilon);
+  EXPECT_NEAR(arr[1].speed.value(), 6.0 * kFactor, kEpsilon);
+  EXPECT_NEAR(arr[2].speed.value(), 4.0 * kFactor, kEpsilon);
+  EXPECT_NEAR(arr[3].speed.value(), 7.0 * kFactor, kEpsilon);
+}