[wpimath] Add cosineScale method to SwerveModuleState and instance optimize (#7114)

wpilibcExamples/src/main/cpp/examples/SwerveBot/cpp/SwerveModule.cpp

@@ -46,29 +46,29 @@ frc::SwerveModulePosition SwerveModule::GetPosition() const {
           units::radian_t{m_turningEncoder.GetDistance()}};
 }
 
-void SwerveModule::SetDesiredState(
-    const frc::SwerveModuleState& referenceState) {
+void SwerveModule::SetDesiredState(frc::SwerveModuleState& referenceState) {
   frc::Rotation2d encoderRotation{
       units::radian_t{m_turningEncoder.GetDistance()}};
 
   // Optimize the reference state to avoid spinning further than 90 degrees
-  auto state =
-      frc::SwerveModuleState::Optimize(referenceState, encoderRotation);
+  referenceState.Optimize(encoderRotation);
 
   // Scale speed by cosine of angle error. This scales down movement
   // perpendicular to the desired direction of travel that can occur when
   // modules change directions. This results in smoother driving.
-  state.speed *= (state.angle - encoderRotation).Cos();
+  referenceState.CosineScale(encoderRotation);
 
   // Calculate the drive output from the drive PID controller.
   const auto driveOutput = m_drivePIDController.Calculate(
-      m_driveEncoder.GetRate(), state.speed.value());
+      m_driveEncoder.GetRate(), referenceState.speed.value());
 
-  const auto driveFeedforward = m_driveFeedforward.Calculate(state.speed);
+  const auto driveFeedforward =
+      m_driveFeedforward.Calculate(referenceState.speed);
 
   // Calculate the turning motor output from the turning PID controller.
   const auto turnOutput = m_turningPIDController.Calculate(
-      units::radian_t{m_turningEncoder.GetDistance()}, state.angle.Radians());
+      units::radian_t{m_turningEncoder.GetDistance()},
+      referenceState.angle.Radians());
 
   const auto turnFeedforward = m_turnFeedforward.Calculate(
       m_turningPIDController.GetSetpoint().velocity);

wpilibcExamples/src/main/cpp/examples/SwerveBot/include/SwerveModule.h

@@ -25,7 +25,7 @@ class SwerveModule {
                int turningEncoderChannelA, int turningEncoderChannelB);
   frc::SwerveModuleState GetState() const;
   frc::SwerveModulePosition GetPosition() const;
-  void SetDesiredState(const frc::SwerveModuleState& state);
+  void SetDesiredState(frc::SwerveModuleState& state);
 
  private:
   static constexpr double kWheelRadius = 0.0508;

wpilibcExamples/src/main/cpp/examples/SwerveControllerCommand/cpp/subsystems/SwerveModule.cpp

@@ -53,27 +53,26 @@ frc::SwerveModulePosition SwerveModule::GetPosition() {
           units::radian_t{m_turningEncoder.GetDistance()}};
 }
 
-void SwerveModule::SetDesiredState(
-    const frc::SwerveModuleState& referenceState) {
+void SwerveModule::SetDesiredState(frc::SwerveModuleState& referenceState) {
   frc::Rotation2d encoderRotation{
       units::radian_t{m_turningEncoder.GetDistance()}};
 
   // Optimize the reference state to avoid spinning further than 90 degrees
-  auto state =
-      frc::SwerveModuleState::Optimize(referenceState, encoderRotation);
+  referenceState.Optimize(encoderRotation);
 
   // Scale speed by cosine of angle error. This scales down movement
   // perpendicular to the desired direction of travel that can occur when
   // modules change directions. This results in smoother driving.
-  state.speed *= (state.angle - encoderRotation).Cos();
+  referenceState.CosineScale(encoderRotation);
 
   // Calculate the drive output from the drive PID controller.
   const auto driveOutput = m_drivePIDController.Calculate(
-      m_driveEncoder.GetRate(), state.speed.value());
+      m_driveEncoder.GetRate(), referenceState.speed.value());
 
   // Calculate the turning motor output from the turning PID controller.
   auto turnOutput = m_turningPIDController.Calculate(
-      units::radian_t{m_turningEncoder.GetDistance()}, state.angle.Radians());
+      units::radian_t{m_turningEncoder.GetDistance()},
+      referenceState.angle.Radians());
 
   // Set the motor outputs.
   m_driveMotor.Set(driveOutput);

wpilibcExamples/src/main/cpp/examples/SwerveControllerCommand/include/subsystems/SwerveModule.h

@@ -27,7 +27,7 @@ class SwerveModule {
 
   frc::SwerveModulePosition GetPosition();
 
-  void SetDesiredState(const frc::SwerveModuleState& state);
+  void SetDesiredState(frc::SwerveModuleState& state);
 
   void ResetEncoders();
 

wpilibcExamples/src/main/cpp/examples/SwerveDrivePoseEstimator/cpp/SwerveModule.cpp

@@ -46,29 +46,29 @@ frc::SwerveModulePosition SwerveModule::GetPosition() const {
           units::radian_t{m_turningEncoder.GetDistance()}};
 }
 
-void SwerveModule::SetDesiredState(
-    const frc::SwerveModuleState& referenceState) {
+void SwerveModule::SetDesiredState(frc::SwerveModuleState& referenceState) {
   frc::Rotation2d encoderRotation{
       units::radian_t{m_turningEncoder.GetDistance()}};
 
   // Optimize the reference state to avoid spinning further than 90 degrees
-  auto state =
-      frc::SwerveModuleState::Optimize(referenceState, encoderRotation);
+  referenceState.Optimize(encoderRotation);
 
   // Scale speed by cosine of angle error. This scales down movement
   // perpendicular to the desired direction of travel that can occur when
   // modules change directions. This results in smoother driving.
-  state.speed *= (state.angle - encoderRotation).Cos();
+  referenceState.CosineScale(encoderRotation);
 
   // Calculate the drive output from the drive PID controller.
   const auto driveOutput = m_drivePIDController.Calculate(
-      m_driveEncoder.GetRate(), state.speed.value());
+      m_driveEncoder.GetRate(), referenceState.speed.value());
 
-  const auto driveFeedforward = m_driveFeedforward.Calculate(state.speed);
+  const auto driveFeedforward =
+      m_driveFeedforward.Calculate(referenceState.speed);
 
   // Calculate the turning motor output from the turning PID controller.
   const auto turnOutput = m_turningPIDController.Calculate(
-      units::radian_t{m_turningEncoder.GetDistance()}, state.angle.Radians());
+      units::radian_t{m_turningEncoder.GetDistance()},
+      referenceState.angle.Radians());
 
   const auto turnFeedforward = m_turnFeedforward.Calculate(
       m_turningPIDController.GetSetpoint().velocity);

wpilibcExamples/src/main/cpp/examples/SwerveDrivePoseEstimator/include/SwerveModule.h

@@ -25,7 +25,7 @@ class SwerveModule {
                int turningEncoderChannelA, int turningEncoderChannelB);
   frc::SwerveModuleState GetState() const;
   frc::SwerveModulePosition GetPosition() const;
-  void SetDesiredState(const frc::SwerveModuleState& state);
+  void SetDesiredState(frc::SwerveModuleState& state);
 
  private:
   static constexpr auto kWheelRadius = 0.0508_m;

wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/examples/swervebot/SwerveModule.java

@@ -115,24 +115,27 @@ public class SwerveModule {
     var encoderRotation = new Rotation2d(m_turningEncoder.getDistance());
 
     // Optimize the reference state to avoid spinning further than 90 degrees
-    SwerveModuleState state = SwerveModuleState.optimize(desiredState, encoderRotation);
+    desiredState.optimize(encoderRotation);
 
     // Scale speed by cosine of angle error. This scales down movement perpendicular to the desired
     // direction of travel that can occur when modules change directions. This results in smoother
     // driving.
-    state.speedMetersPerSecond *= state.angle.minus(encoderRotation).getCos();
+    desiredState.cosineScale(encoderRotation);
 
     // Calculate the drive output from the drive PID controller.
 
     final double driveOutput =
-        m_drivePIDController.calculate(m_driveEncoder.getRate(), state.speedMetersPerSecond);
+        m_drivePIDController.calculate(m_driveEncoder.getRate(), desiredState.speedMetersPerSecond);
 
     final double driveFeedforward =
-        m_driveFeedforward.calculate(MetersPerSecond.of(state.speedMetersPerSecond)).in(Volts);
+        m_driveFeedforward
+            .calculate(MetersPerSecond.of(desiredState.speedMetersPerSecond))
+            .in(Volts);
 
     // Calculate the turning motor output from the turning PID controller.
     final double turnOutput =
-        m_turningPIDController.calculate(m_turningEncoder.getDistance(), state.angle.getRadians());
+        m_turningPIDController.calculate(
+            m_turningEncoder.getDistance(), desiredState.angle.getRadians());
 
     final double turnFeedforward =
         m_turnFeedforward

wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/examples/swervecontrollercommand/subsystems/SwerveModule.java

@@ -108,20 +108,21 @@ public class SwerveModule {
     var encoderRotation = new Rotation2d(m_turningEncoder.getDistance());
 
     // Optimize the reference state to avoid spinning further than 90 degrees
-    SwerveModuleState state = SwerveModuleState.optimize(desiredState, encoderRotation);
+    desiredState.optimize(encoderRotation);
 
     // Scale speed by cosine of angle error. This scales down movement perpendicular to the desired
     // direction of travel that can occur when modules change directions. This results in smoother
     // driving.
-    state.speedMetersPerSecond *= state.angle.minus(encoderRotation).getCos();
+    desiredState.cosineScale(encoderRotation);
 
     // Calculate the drive output from the drive PID controller.
     final double driveOutput =
-        m_drivePIDController.calculate(m_driveEncoder.getRate(), state.speedMetersPerSecond);
+        m_drivePIDController.calculate(m_driveEncoder.getRate(), desiredState.speedMetersPerSecond);
 
     // Calculate the turning motor output from the turning PID controller.
     final double turnOutput =
-        m_turningPIDController.calculate(m_turningEncoder.getDistance(), state.angle.getRadians());
+        m_turningPIDController.calculate(
+            m_turningEncoder.getDistance(), desiredState.angle.getRadians());
 
     // Calculate the turning motor output from the turning PID controller.
     m_driveMotor.set(driveOutput);

wpilibjExamples/src/main/java/edu/wpi/first/wpilibj/examples/swervedriveposeestimator/SwerveModule.java

@@ -115,23 +115,26 @@ public class SwerveModule {
     var encoderRotation = new Rotation2d(m_turningEncoder.getDistance());
 
     // Optimize the reference state to avoid spinning further than 90 degrees
-    SwerveModuleState state = SwerveModuleState.optimize(desiredState, encoderRotation);
+    desiredState.optimize(encoderRotation);
 
     // Scale speed by cosine of angle error. This scales down movement perpendicular to the desired
     // direction of travel that can occur when modules change directions. This results in smoother
     // driving.
-    state.speedMetersPerSecond *= state.angle.minus(encoderRotation).getCos();
+    desiredState.cosineScale(encoderRotation);
 
     // Calculate the drive output from the drive PID controller.
     final double driveOutput =
-        m_drivePIDController.calculate(m_driveEncoder.getRate(), state.speedMetersPerSecond);
+        m_drivePIDController.calculate(m_driveEncoder.getRate(), desiredState.speedMetersPerSecond);
 
     final double driveFeedforward =
-        m_driveFeedforward.calculate(MetersPerSecond.of(state.speedMetersPerSecond)).in(Volts);
+        m_driveFeedforward
+            .calculate(MetersPerSecond.of(desiredState.speedMetersPerSecond))
+            .in(Volts);
 
     // Calculate the turning motor output from the turning PID controller.
     final double turnOutput =
-        m_turningPIDController.calculate(m_turningEncoder.getDistance(), state.angle.getRadians());
+        m_turningPIDController.calculate(
+            m_turningEncoder.getDistance(), desiredState.angle.getRadians());
 
     final double turnFeedforward =
         m_turnFeedforward

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

@@ -83,6 +83,21 @@ public class SwerveModuleState
         "SwerveModuleState(Speed: %.2f m/s, Angle: %s)", speedMetersPerSecond, angle);
   }
 
+  /**
+   * Minimize the change in heading this swerve module state would require by potentially reversing
+   * the direction the wheel spins. If this is used with the PIDController class's continuous input
+   * functionality, the furthest a wheel will ever rotate is 90 degrees.
+   *
+   * @param currentAngle The current module angle.
+   */
+  public void optimize(Rotation2d currentAngle) {
+    var delta = angle.minus(currentAngle);
+    if (Math.abs(delta.getDegrees()) > 90.0) {
+      speedMetersPerSecond *= -1;
+      angle = angle.rotateBy(Rotation2d.kPi);
+    }
+  }
+
   /**
    * Minimize the change in heading the desired swerve module state would require by potentially
    * reversing the direction the wheel spins. If this is used with the PIDController class's
@@ -91,7 +106,9 @@ public class SwerveModuleState
    * @param desiredState The desired state.
    * @param currentAngle The current module angle.
    * @return Optimized swerve module state.
+   * @deprecated Use the instance method instead.
    */
+  @Deprecated
   public static SwerveModuleState optimize(
       SwerveModuleState desiredState, Rotation2d currentAngle) {
     var delta = desiredState.angle.minus(currentAngle);
@@ -102,4 +119,15 @@ public class SwerveModuleState
       return new SwerveModuleState(desiredState.speedMetersPerSecond, desiredState.angle);
     }
   }
+
+  /**
+   * Scales speed by cosine of angle error. This scales down movement perpendicular to the desired
+   * direction of travel that can occur when modules change directions. This results in smoother
+   * driving.
+   *
+   * @param currentAngle The current module angle.
+   */
+  public void cosineScale(Rotation2d currentAngle) {
+    speedMetersPerSecond *= angle.minus(currentAngle).getCos();
+  }
 }

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

@@ -34,6 +34,22 @@ struct WPILIB_DLLEXPORT SwerveModuleState {
    */
   bool operator==(const SwerveModuleState& other) const;
 
+  /**
+   * Minimize the change in the heading this swerve module state would
+   * require by potentially reversing the direction the wheel spins. If this is
+   * used with the PIDController class's continuous input functionality, the
+   * furthest a wheel will ever rotate is 90 degrees.
+   *
+   * @param currentAngle The current module angle.
+   */
+  void Optimize(const Rotation2d& currentAngle) {
+    auto delta = angle - currentAngle;
+    if (units::math::abs(delta.Degrees()) > 90_deg) {
+      speed *= -1;
+      angle = angle + Rotation2d{180_deg};
+    }
+  }
+
   /**
    * Minimize the change in heading the desired swerve module state would
    * require by potentially reversing the direction the wheel spins. If this is
@@ -43,8 +59,20 @@ struct WPILIB_DLLEXPORT SwerveModuleState {
    * @param desiredState The desired state.
    * @param currentAngle The current module angle.
    */
+  [[deprecated("Use instance method instead.")]]
   static SwerveModuleState Optimize(const SwerveModuleState& desiredState,
                                     const Rotation2d& currentAngle);
+
+  /**
+   * Scales speed by cosine of angle error. This scales down movement
+   * perpendicular to the desired direction of travel that can occur when
+   * modules change directions. This results in smoother driving.
+   *
+   * @param currentAngle The current module angle.
+   */
+  void CosineScale(const Rotation2d& currentAngle) {
+    speed *= (angle - currentAngle).Cos();
+  }
 };
 }  // namespace frc
 

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

@@ -17,37 +17,136 @@ class SwerveModuleStateTest {
   void testOptimize() {
     var angleA = Rotation2d.fromDegrees(45);
     var refA = new SwerveModuleState(-2.0, Rotation2d.kPi);
-    var optimizedA = SwerveModuleState.optimize(refA, angleA);
+    refA.optimize(angleA);
 
     assertAll(
-        () -> assertEquals(2.0, optimizedA.speedMetersPerSecond, kEpsilon),
-        () -> assertEquals(0.0, optimizedA.angle.getDegrees(), kEpsilon));
+        () -> assertEquals(2.0, refA.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(0.0, refA.angle.getDegrees(), kEpsilon));
 
     var angleB = Rotation2d.fromDegrees(-50);
     var refB = new SwerveModuleState(4.7, Rotation2d.fromDegrees(41));
-    var optimizedB = SwerveModuleState.optimize(refB, angleB);
+    refB.optimize(angleB);
 
     assertAll(
-        () -> assertEquals(-4.7, optimizedB.speedMetersPerSecond, kEpsilon),
-        () -> assertEquals(-139.0, optimizedB.angle.getDegrees(), kEpsilon));
+        () -> assertEquals(-4.7, refB.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(-139.0, refB.angle.getDegrees(), kEpsilon));
   }
 
   @Test
   void testNoOptimize() {
     var angleA = Rotation2d.kZero;
     var refA = new SwerveModuleState(2.0, Rotation2d.fromDegrees(89));
-    var optimizedA = SwerveModuleState.optimize(refA, angleA);
+    refA.optimize(angleA);
 
     assertAll(
-        () -> assertEquals(2.0, optimizedA.speedMetersPerSecond, kEpsilon),
-        () -> assertEquals(89.0, optimizedA.angle.getDegrees(), kEpsilon));
+        () -> assertEquals(2.0, refA.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(89.0, refA.angle.getDegrees(), kEpsilon));
 
     var angleB = Rotation2d.kZero;
     var refB = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(-2));
-    var optimizedB = SwerveModuleState.optimize(refB, angleB);
+    refB.optimize(angleB);
 
     assertAll(
-        () -> assertEquals(-2.0, optimizedB.speedMetersPerSecond, kEpsilon),
-        () -> assertEquals(-2.0, optimizedB.angle.getDegrees(), kEpsilon));
+        () -> assertEquals(-2.0, refB.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(-2.0, refB.angle.getDegrees(), kEpsilon));
+  }
+
+  @Test
+  void testCosineScale() {
+    var angleA = Rotation2d.fromDegrees(0.0);
+    var refA = new SwerveModuleState(2.0, Rotation2d.fromDegrees(45.0));
+    refA.cosineScale(angleA);
+
+    assertAll(
+        () -> assertEquals(Math.sqrt(2.0), refA.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refA.angle.getDegrees(), kEpsilon));
+
+    var angleB = Rotation2d.fromDegrees(45.0);
+    var refB = new SwerveModuleState(2.0, Rotation2d.fromDegrees(45.0));
+    refB.cosineScale(angleB);
+
+    assertAll(
+        () -> assertEquals(2.0, refB.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refB.angle.getDegrees(), kEpsilon));
+
+    var angleC = Rotation2d.fromDegrees(-45.0);
+    var refC = new SwerveModuleState(2.0, Rotation2d.fromDegrees(45.0));
+    refC.cosineScale(angleC);
+
+    assertAll(
+        () -> assertEquals(0.0, refC.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refC.angle.getDegrees(), kEpsilon));
+
+    var angleD = Rotation2d.fromDegrees(135.0);
+    var refD = new SwerveModuleState(2.0, Rotation2d.fromDegrees(45.0));
+    refD.cosineScale(angleD);
+
+    assertAll(
+        () -> assertEquals(0.0, refD.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refD.angle.getDegrees(), kEpsilon));
+
+    var angleE = Rotation2d.fromDegrees(-135.0);
+    var refE = new SwerveModuleState(2.0, Rotation2d.fromDegrees(45.0));
+    refE.cosineScale(angleE);
+
+    assertAll(
+        () -> assertEquals(-2.0, refE.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refE.angle.getDegrees(), kEpsilon));
+
+    var angleF = Rotation2d.fromDegrees(180.0);
+    var refF = new SwerveModuleState(2.0, Rotation2d.fromDegrees(45.0));
+    refF.cosineScale(angleF);
+
+    assertAll(
+        () -> assertEquals(-Math.sqrt(2.0), refF.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refF.angle.getDegrees(), kEpsilon));
+
+    var angleG = Rotation2d.fromDegrees(0.0);
+    var refG = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(45.0));
+    refG.cosineScale(angleG);
+
+    assertAll(
+        () -> assertEquals(-Math.sqrt(2.0), refG.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refG.angle.getDegrees(), kEpsilon));
+
+    var angleH = Rotation2d.fromDegrees(45.0);
+    var refH = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(45.0));
+    refH.cosineScale(angleH);
+
+    assertAll(
+        () -> assertEquals(-2.0, refH.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refH.angle.getDegrees(), kEpsilon));
+
+    var angleI = Rotation2d.fromDegrees(-45.0);
+    var refI = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(45.0));
+    refI.cosineScale(angleI);
+
+    assertAll(
+        () -> assertEquals(0.0, refI.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refI.angle.getDegrees(), kEpsilon));
+
+    var angleJ = Rotation2d.fromDegrees(135.0);
+    var refJ = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(45.0));
+    refJ.cosineScale(angleJ);
+
+    assertAll(
+        () -> assertEquals(0.0, refJ.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refJ.angle.getDegrees(), kEpsilon));
+
+    var angleK = Rotation2d.fromDegrees(-135.0);
+    var refK = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(45.0));
+    refK.cosineScale(angleK);
+
+    assertAll(
+        () -> assertEquals(2.0, refK.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refK.angle.getDegrees(), kEpsilon));
+
+    var angleL = Rotation2d.fromDegrees(180.0);
+    var refL = new SwerveModuleState(-2.0, Rotation2d.fromDegrees(45.0));
+    refL.cosineScale(angleL);
+
+    assertAll(
+        () -> assertEquals(Math.sqrt(2.0), refL.speedMetersPerSecond, kEpsilon),
+        () -> assertEquals(45.0, refL.angle.getDegrees(), kEpsilon));
   }
 }

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

@@ -12,33 +12,119 @@ static constexpr double kEpsilon = 1E-9;
 TEST(SwerveModuleStateTest, Optimize) {
   frc::Rotation2d angleA{45_deg};
   frc::SwerveModuleState refA{-2_mps, 180_deg};
-  auto optimizedA = frc::SwerveModuleState::Optimize(refA, angleA);
+  refA.Optimize(angleA);
 
-  EXPECT_NEAR(optimizedA.speed.value(), 2.0, kEpsilon);
-  EXPECT_NEAR(optimizedA.angle.Degrees().value(), 0.0, kEpsilon);
+  EXPECT_NEAR(refA.speed.value(), 2.0, kEpsilon);
+  EXPECT_NEAR(refA.angle.Degrees().value(), 0.0, kEpsilon);
 
   frc::Rotation2d angleB{-50_deg};
   frc::SwerveModuleState refB{4.7_mps, 41_deg};
-  auto optimizedB = frc::SwerveModuleState::Optimize(refB, angleB);
+  refB.Optimize(angleB);
 
-  EXPECT_NEAR(optimizedB.speed.value(), -4.7, kEpsilon);
-  EXPECT_NEAR(optimizedB.angle.Degrees().value(), -139.0, kEpsilon);
+  EXPECT_NEAR(refB.speed.value(), -4.7, kEpsilon);
+  EXPECT_NEAR(refB.angle.Degrees().value(), -139.0, kEpsilon);
 }
 
 TEST(SwerveModuleStateTest, NoOptimize) {
   frc::Rotation2d angleA{0_deg};
   frc::SwerveModuleState refA{2_mps, 89_deg};
-  auto optimizedA = frc::SwerveModuleState::Optimize(refA, angleA);
+  refA.Optimize(angleA);
 
-  EXPECT_NEAR(optimizedA.speed.value(), 2.0, kEpsilon);
-  EXPECT_NEAR(optimizedA.angle.Degrees().value(), 89.0, kEpsilon);
+  EXPECT_NEAR(refA.speed.value(), 2.0, kEpsilon);
+  EXPECT_NEAR(refA.angle.Degrees().value(), 89.0, kEpsilon);
 
   frc::Rotation2d angleB{0_deg};
   frc::SwerveModuleState refB{-2_mps, -2_deg};
-  auto optimizedB = frc::SwerveModuleState::Optimize(refB, angleB);
+  refB.Optimize(angleB);
 
-  EXPECT_NEAR(optimizedB.speed.value(), -2.0, kEpsilon);
-  EXPECT_NEAR(optimizedB.angle.Degrees().value(), -2.0, kEpsilon);
+  EXPECT_NEAR(refB.speed.value(), -2.0, kEpsilon);
+  EXPECT_NEAR(refB.angle.Degrees().value(), -2.0, kEpsilon);
+}
+
+TEST(SwerveModuleStateTest, CosineScaling) {
+  frc::Rotation2d angleA{0_deg};
+  frc::SwerveModuleState refA{2_mps, 45_deg};
+  refA.CosineScale(angleA);
+
+  EXPECT_NEAR(refA.speed.value(), std::sqrt(2.0), kEpsilon);
+  EXPECT_NEAR(refA.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleB{45_deg};
+  frc::SwerveModuleState refB{2_mps, 45_deg};
+  refB.CosineScale(angleB);
+
+  EXPECT_NEAR(refB.speed.value(), 2.0, kEpsilon);
+  EXPECT_NEAR(refB.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleC{-45_deg};
+  frc::SwerveModuleState refC{2_mps, 45_deg};
+  refC.CosineScale(angleC);
+
+  EXPECT_NEAR(refC.speed.value(), 0.0, kEpsilon);
+  EXPECT_NEAR(refC.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleD{135_deg};
+  frc::SwerveModuleState refD{2_mps, 45_deg};
+  refD.CosineScale(angleD);
+
+  EXPECT_NEAR(refD.speed.value(), 0.0, kEpsilon);
+  EXPECT_NEAR(refD.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleE{-135_deg};
+  frc::SwerveModuleState refE{2_mps, 45_deg};
+  refE.CosineScale(angleE);
+
+  EXPECT_NEAR(refE.speed.value(), -2.0, kEpsilon);
+  EXPECT_NEAR(refE.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleF{180_deg};
+  frc::SwerveModuleState refF{2_mps, 45_deg};
+  refF.CosineScale(angleF);
+
+  EXPECT_NEAR(refF.speed.value(), -std::sqrt(2.0), kEpsilon);
+  EXPECT_NEAR(refF.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleG{0_deg};
+  frc::SwerveModuleState refG{-2_mps, 45_deg};
+  refG.CosineScale(angleG);
+
+  EXPECT_NEAR(refG.speed.value(), -std::sqrt(2.0), kEpsilon);
+  EXPECT_NEAR(refG.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleH{45_deg};
+  frc::SwerveModuleState refH{-2_mps, 45_deg};
+  refH.CosineScale(angleH);
+
+  EXPECT_NEAR(refH.speed.value(), -2.0, kEpsilon);
+  EXPECT_NEAR(refH.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleI{-45_deg};
+  frc::SwerveModuleState refI{-2_mps, 45_deg};
+  refI.CosineScale(angleI);
+
+  EXPECT_NEAR(refI.speed.value(), 0.0, kEpsilon);
+  EXPECT_NEAR(refI.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleJ{135_deg};
+  frc::SwerveModuleState refJ{-2_mps, 45_deg};
+  refJ.CosineScale(angleJ);
+
+  EXPECT_NEAR(refJ.speed.value(), 0.0, kEpsilon);
+  EXPECT_NEAR(refJ.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleK{-135_deg};
+  frc::SwerveModuleState refK{-2_mps, 45_deg};
+  refK.CosineScale(angleK);
+
+  EXPECT_NEAR(refK.speed.value(), 2.0, kEpsilon);
+  EXPECT_NEAR(refK.angle.Degrees().value(), 45.0, kEpsilon);
+
+  frc::Rotation2d angleL{180_deg};
+  frc::SwerveModuleState refL{-2_mps, 45_deg};
+  refL.CosineScale(angleL);
+
+  EXPECT_NEAR(refL.speed.value(), std::sqrt(2.0), kEpsilon);
+  EXPECT_NEAR(refL.angle.Degrees().value(), 45.0, kEpsilon);
 }
 
 TEST(SwerveModuleStateTest, Equality) {