From 16d2952bfed3bfa3b60c7045e1c937d171e82d0c Mon Sep 17 00:00:00 2001
From: Vignesh Balasubramaniam <vigneshb.hp@gmail.com>
Date: Sun, 22 Sep 2024 18:53:07 -0500
Subject: [PATCH] Fix vector math

---
 .../lasarobotics/drive/MAXSwerveModule.java   |  6 +++--
 .../drive/MAXSwerveModuleTest.java            | 23 +++++++++++++++++++
 .../drive/TractionControlTest.java            |  6 ++---
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/src/main/java/org/lasarobotics/drive/MAXSwerveModule.java b/src/main/java/org/lasarobotics/drive/MAXSwerveModule.java
index b822003b..aa26a1e8 100644
--- a/src/main/java/org/lasarobotics/drive/MAXSwerveModule.java
+++ b/src/main/java/org/lasarobotics/drive/MAXSwerveModule.java
@@ -366,7 +366,7 @@ private SwerveModuleState getDesiredState(SwerveModuleState requestedState) {
 
   /**
    * Get velocity of module parallel to the desired orientation
-   * @param state Desired swerve module state representing desired speed
+   * @param state Desired swerve module state representing desired speed and orientation
    * @param realSpeeds Real speeds of robot from IMU
    * @return Velocity of module parallel to desired module orientation
    */
@@ -386,7 +386,9 @@ private Measure<Velocity<Distance>> getParallelVelocity(SwerveModuleState state,
 
     // Calculate portion of inertial velocity that is in direction of desired
     var parallelModuleVelocity = moduleDesiredVelocityVector.scalarMultiply(
-      moduleDesiredVelocityVector.getNorm() != 0.0 ? moduleInertialVelocityVector.dotProduct(moduleDesiredVelocityVector.normalize()) : 1.0
+      moduleDesiredVelocityVector.getNorm() != 0.0
+        ? moduleInertialVelocityVector.dotProduct(moduleDesiredVelocityVector) / moduleDesiredVelocityVector.getNormSq()
+        : 1.0
     );
 
     // Return velocity
diff --git a/src/test/java/org/lasarobotics/drive/MAXSwerveModuleTest.java b/src/test/java/org/lasarobotics/drive/MAXSwerveModuleTest.java
index 37d6dd03..d82866bf 100644
--- a/src/test/java/org/lasarobotics/drive/MAXSwerveModuleTest.java
+++ b/src/test/java/org/lasarobotics/drive/MAXSwerveModuleTest.java
@@ -29,6 +29,7 @@
 
 import edu.wpi.first.math.geometry.Rotation2d;
 import edu.wpi.first.math.geometry.Translation2d;
+import edu.wpi.first.math.kinematics.ChassisSpeeds;
 import edu.wpi.first.math.kinematics.SwerveModulePosition;
 import edu.wpi.first.math.kinematics.SwerveModuleState;
 import edu.wpi.first.units.Current;
@@ -296,4 +297,26 @@ public void maxLinearSpeed() {
     assertTrue(m_lRearModule.getMaxLinearSpeed().isNear(VORTEX_MAX_LINEAR_SPEED, DELTA));
     assertTrue(m_rRearModule.getMaxLinearSpeed().isNear(VORTEX_MAX_LINEAR_SPEED, DELTA));
   }
+
+  @Test
+  @Order(6)
+  @DisplayName("Test if module limits speed")
+  public void tractionControl() {
+    // Hardcode sensor values
+    SparkInputsAutoLogged sparkInputs = new SparkInputsAutoLogged();
+    when(m_lFrontRotateMotor.getInputs()).thenReturn(sparkInputs);
+    sparkInputs.encoderVelocity = 3.0;
+    when(m_lFrontDriveMotor.getInputs()).thenReturn(sparkInputs);
+
+    // Hardcode other inputs
+    var moduleState = new SwerveModuleState(NEO_MAX_LINEAR_SPEED, Rotation2d.fromDegrees(0.0));
+    var realSpeeds = new ChassisSpeeds(Units.MetersPerSecond.of(3.0), Units.MetersPerSecond.zero(), Units.DegreesPerSecond.zero());
+
+    // Attempt to drive module
+    m_lFrontModule.set(moduleState, realSpeeds);
+
+    // Verify that motors are being driven with expected values
+    verify(m_lFrontDriveMotor).set(AdditionalMatchers.leq(NEO_MAX_LINEAR_SPEED.in(Units.MetersPerSecond)), ArgumentMatchers.eq(ControlType.kVelocity));
+    verify(m_lFrontRotateMotor).set(AdditionalMatchers.eq(-Math.PI / 2, DELTA), ArgumentMatchers.eq(ControlType.kPosition));
+  }
 }
diff --git a/src/test/java/org/lasarobotics/drive/TractionControlTest.java b/src/test/java/org/lasarobotics/drive/TractionControlTest.java
index 4d1a0d48..e4f21fe0 100644
--- a/src/test/java/org/lasarobotics/drive/TractionControlTest.java
+++ b/src/test/java/org/lasarobotics/drive/TractionControlTest.java
@@ -26,10 +26,10 @@
 @TestMethodOrder(MethodOrderer.OrderAnnotation.class)
 public class TractionControlTest {
   private final Measure<Dimensionless> SLIP_RATIO = Units.Percent.of(8.0);
-  private final Measure<Dimensionless> STATIC_FRICTION_COEFFICIENT = Units.Value.of(1.1);
+  private final Measure<Dimensionless> STATIC_FRICTION_COEFFICIENT = Units.Value.of(0.9);
   private final Measure<Dimensionless> DYNAMIC_FRICTION_COEFFICIENT = Units.Value.of(0.8);
-  private final Measure<Mass> MASS = Units.Pounds.of(110.0);
-  private final Measure<Velocity<Distance>> MAX_LINEAR_SPEED = Units.MetersPerSecond.of(4.30);
+  private final Measure<Mass> MASS = Units.Pounds.of(135.0);
+  private final Measure<Velocity<Distance>> MAX_LINEAR_SPEED = Units.MetersPerSecond.of(5.2);
   private final double THRESHOLD = 0.05;
 
   private TractionControlController m_tractionControlController;