001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.math.kinematics;
006
007import static edu.wpi.first.units.Units.MetersPerSecond;
008import static edu.wpi.first.units.Units.RadiansPerSecond;
009
010import edu.wpi.first.math.MathSharedStore;
011import edu.wpi.first.math.geometry.Rotation2d;
012import edu.wpi.first.math.geometry.Translation2d;
013import edu.wpi.first.math.geometry.Twist2d;
014import edu.wpi.first.math.kinematics.proto.SwerveDriveKinematicsProto;
015import edu.wpi.first.math.kinematics.struct.SwerveDriveKinematicsStruct;
016import edu.wpi.first.units.measure.AngularVelocity;
017import edu.wpi.first.units.measure.LinearVelocity;
018import edu.wpi.first.util.protobuf.ProtobufSerializable;
019import edu.wpi.first.util.struct.Struct;
020import edu.wpi.first.util.struct.StructSerializable;
021import java.util.Arrays;
022import org.ejml.simple.SimpleMatrix;
023
024/**
025 * Helper class that converts a chassis velocity (dx, dy, and dtheta components) into individual
026 * module states (speed and angle).
027 *
028 * <p>The inverse kinematics (converting from a desired chassis velocity to individual module
029 * states) uses the relative locations of the modules with respect to the center of rotation. The
030 * center of rotation for inverse kinematics is also variable. This means that you can set your
031 * center of rotation in a corner of the robot to perform special evasion maneuvers.
032 *
033 * <p>Forward kinematics (converting an array of module states into the overall chassis motion) is
034 * performs the exact opposite of what inverse kinematics does. Since this is an overdetermined
035 * system (more equations than variables), we use a least-squares approximation.
036 *
037 * <p>The inverse kinematics: [moduleStates] = [moduleLocations] * [chassisSpeeds] We take the
038 * Moore-Penrose pseudoinverse of [moduleLocations] and then multiply by [moduleStates] to get our
039 * chassis speeds.
040 *
041 * <p>Forward kinematics is also used for odometry -- determining the position of the robot on the
042 * field using encoders and a gyro.
043 */
044@SuppressWarnings("overrides")
045public class SwerveDriveKinematics
046    implements Kinematics<SwerveModuleState[], SwerveModulePosition[]>,
047        ProtobufSerializable,
048        StructSerializable {
049  private final SimpleMatrix m_inverseKinematics;
050  private final SimpleMatrix m_forwardKinematics;
051
052  private final int m_numModules;
053  private final Translation2d[] m_modules;
054  private Rotation2d[] m_moduleHeadings;
055  private Translation2d m_prevCoR = Translation2d.kZero;
056
057  /**
058   * Constructs a swerve drive kinematics object. This takes in a variable number of module
059   * locations as Translation2d objects. The order in which you pass in the module locations is the
060   * same order that you will receive the module states when performing inverse kinematics. It is
061   * also expected that you pass in the module states in the same order when calling the forward
062   * kinematics methods.
063   *
064   * @param moduleTranslations The locations of the modules relative to the physical center of the
065   *     robot.
066   */
067  public SwerveDriveKinematics(Translation2d... moduleTranslations) {
068    if (moduleTranslations.length < 2) {
069      throw new IllegalArgumentException("A swerve drive requires at least two modules");
070    }
071    m_numModules = moduleTranslations.length;
072    m_modules = Arrays.copyOf(moduleTranslations, m_numModules);
073    m_moduleHeadings = new Rotation2d[m_numModules];
074    Arrays.fill(m_moduleHeadings, Rotation2d.kZero);
075    m_inverseKinematics = new SimpleMatrix(m_numModules * 2, 3);
076
077    for (int i = 0; i < m_numModules; i++) {
078      m_inverseKinematics.setRow(i * 2 + 0, 0, /* Start Data */ 1, 0, -m_modules[i].getY());
079      m_inverseKinematics.setRow(i * 2 + 1, 0, /* Start Data */ 0, 1, +m_modules[i].getX());
080    }
081    m_forwardKinematics = m_inverseKinematics.pseudoInverse();
082
083    MathSharedStore.reportUsage("SwerveDriveKinematics", "");
084  }
085
086  /**
087   * Reset the internal swerve module headings.
088   *
089   * @param moduleHeadings The swerve module headings. The order of the module headings should be
090   *     same as passed into the constructor of this class.
091   */
092  public void resetHeadings(Rotation2d... moduleHeadings) {
093    if (moduleHeadings.length != m_numModules) {
094      throw new IllegalArgumentException(
095          "Number of headings is not consistent with number of module locations provided in "
096              + "constructor");
097    }
098    m_moduleHeadings = Arrays.copyOf(moduleHeadings, m_numModules);
099  }
100
101  /**
102   * Performs inverse kinematics to return the module states from a desired chassis velocity. This
103   * method is often used to convert joystick values into module speeds and angles.
104   *
105   * <p>This function also supports variable centers of rotation. During normal operations, the
106   * center of rotation is usually the same as the physical center of the robot; therefore, the
107   * argument is defaulted to that use case. However, if you wish to change the center of rotation
108   * for evasive maneuvers, vision alignment, or for any other use case, you can do so.
109   *
110   * <p>In the case that the desired chassis speeds are zero (i.e. the robot will be stationary),
111   * the previously calculated module angle will be maintained.
112   *
113   * @param chassisSpeeds The desired chassis speed.
114   * @param centerOfRotation The center of rotation. For example, if you set the center of rotation
115   *     at one corner of the robot and provide a chassis speed that only has a dtheta component,
116   *     the robot will rotate around that corner.
117   * @return An array containing the module states. Use caution because these module states are not
118   *     normalized. Sometimes, a user input may cause one of the module speeds to go above the
119   *     attainable max velocity. Use the {@link #desaturateWheelSpeeds(SwerveModuleState[], double)
120   *     DesaturateWheelSpeeds} function to rectify this issue.
121   */
122  public SwerveModuleState[] toSwerveModuleStates(
123      ChassisSpeeds chassisSpeeds, Translation2d centerOfRotation) {
124    var moduleStates = new SwerveModuleState[m_numModules];
125
126    if (chassisSpeeds.vx == 0.0 && chassisSpeeds.vy == 0.0 && chassisSpeeds.omega == 0.0) {
127      for (int i = 0; i < m_numModules; i++) {
128        moduleStates[i] = new SwerveModuleState(0.0, m_moduleHeadings[i]);
129      }
130
131      return moduleStates;
132    }
133
134    if (!centerOfRotation.equals(m_prevCoR)) {
135      for (int i = 0; i < m_numModules; i++) {
136        m_inverseKinematics.setRow(
137            i * 2 + 0, 0, /* Start Data */ 1, 0, -m_modules[i].getY() + centerOfRotation.getY());
138        m_inverseKinematics.setRow(
139            i * 2 + 1, 0, /* Start Data */ 0, 1, +m_modules[i].getX() - centerOfRotation.getX());
140      }
141      m_prevCoR = centerOfRotation;
142    }
143
144    var chassisSpeedsVector = new SimpleMatrix(3, 1);
145    chassisSpeedsVector.setColumn(0, 0, chassisSpeeds.vx, chassisSpeeds.vy, chassisSpeeds.omega);
146
147    var moduleStatesMatrix = m_inverseKinematics.mult(chassisSpeedsVector);
148
149    for (int i = 0; i < m_numModules; i++) {
150      double x = moduleStatesMatrix.get(i * 2, 0);
151      double y = moduleStatesMatrix.get(i * 2 + 1, 0);
152
153      double speed = Math.hypot(x, y);
154      Rotation2d angle = speed > 1e-6 ? new Rotation2d(x, y) : m_moduleHeadings[i];
155
156      moduleStates[i] = new SwerveModuleState(speed, angle);
157      m_moduleHeadings[i] = angle;
158    }
159
160    return moduleStates;
161  }
162
163  /**
164   * Performs inverse kinematics. See {@link #toSwerveModuleStates(ChassisSpeeds, Translation2d)}
165   * toSwerveModuleStates for more information.
166   *
167   * @param chassisSpeeds The desired chassis speed.
168   * @return An array containing the module states.
169   */
170  public SwerveModuleState[] toSwerveModuleStates(ChassisSpeeds chassisSpeeds) {
171    return toSwerveModuleStates(chassisSpeeds, Translation2d.kZero);
172  }
173
174  @Override
175  public SwerveModuleState[] toWheelSpeeds(ChassisSpeeds chassisSpeeds) {
176    return toSwerveModuleStates(chassisSpeeds);
177  }
178
179  /**
180   * Performs forward kinematics to return the resulting chassis state from the given module states.
181   * This method is often used for odometry -- determining the robot's position on the field using
182   * data from the real-world speed and angle of each module on the robot.
183   *
184   * @param moduleStates The state of the modules (as a SwerveModuleState type) as measured from
185   *     respective encoders and gyros. The order of the swerve module states should be same as
186   *     passed into the constructor of this class.
187   * @return The resulting chassis speed.
188   */
189  @Override
190  public ChassisSpeeds toChassisSpeeds(SwerveModuleState... moduleStates) {
191    if (moduleStates.length != m_numModules) {
192      throw new IllegalArgumentException(
193          "Number of modules is not consistent with number of module locations provided in "
194              + "constructor");
195    }
196    var moduleStatesMatrix = new SimpleMatrix(m_numModules * 2, 1);
197
198    for (int i = 0; i < m_numModules; i++) {
199      var module = moduleStates[i];
200      moduleStatesMatrix.set(i * 2, 0, module.speed * module.angle.getCos());
201      moduleStatesMatrix.set(i * 2 + 1, module.speed * module.angle.getSin());
202    }
203
204    var chassisSpeedsVector = m_forwardKinematics.mult(moduleStatesMatrix);
205    return new ChassisSpeeds(
206        chassisSpeedsVector.get(0, 0),
207        chassisSpeedsVector.get(1, 0),
208        chassisSpeedsVector.get(2, 0));
209  }
210
211  /**
212   * Performs forward kinematics to return the resulting chassis state from the given module states.
213   * This method is often used for odometry -- determining the robot's position on the field using
214   * data from the real-world speed and angle of each module on the robot.
215   *
216   * @param moduleDeltas The latest change in position of the modules (as a SwerveModulePosition
217   *     type) as measured from respective encoders and gyros. The order of the swerve module states
218   *     should be same as passed into the constructor of this class.
219   * @return The resulting Twist2d.
220   */
221  public Twist2d toTwist2d(SwerveModulePosition... moduleDeltas) {
222    if (moduleDeltas.length != m_numModules) {
223      throw new IllegalArgumentException(
224          "Number of modules is not consistent with number of module locations provided in "
225              + "constructor");
226    }
227    var moduleDeltaMatrix = new SimpleMatrix(m_numModules * 2, 1);
228
229    for (int i = 0; i < m_numModules; i++) {
230      var module = moduleDeltas[i];
231      moduleDeltaMatrix.set(i * 2, 0, module.distance * module.angle.getCos());
232      moduleDeltaMatrix.set(i * 2 + 1, module.distance * module.angle.getSin());
233    }
234
235    var chassisDeltaVector = m_forwardKinematics.mult(moduleDeltaMatrix);
236    return new Twist2d(
237        chassisDeltaVector.get(0, 0), chassisDeltaVector.get(1, 0), chassisDeltaVector.get(2, 0));
238  }
239
240  @Override
241  public Twist2d toTwist2d(SwerveModulePosition[] start, SwerveModulePosition[] end) {
242    if (start.length != end.length) {
243      throw new IllegalArgumentException("Inconsistent number of modules!");
244    }
245    var newPositions = new SwerveModulePosition[start.length];
246    for (int i = 0; i < start.length; i++) {
247      newPositions[i] = new SwerveModulePosition(end[i].distance - start[i].distance, end[i].angle);
248    }
249    return toTwist2d(newPositions);
250  }
251
252  /**
253   * Renormalizes the wheel speeds if any individual speed is above the specified maximum.
254   *
255   * <p>Sometimes, after inverse kinematics, the requested speed from one or more modules may be
256   * above the max attainable speed for the driving motor on that module. To fix this issue, one can
257   * reduce all the wheel speeds to make sure that all requested module speeds are at-or-below the
258   * absolute threshold, while maintaining the ratio of speeds between modules.
259   *
260   * <p>Scaling down the module speeds rotates the direction of net motion in the opposite direction
261   * of rotational velocity, which makes discretizing the chassis speeds inaccurate because the
262   * discretization did not account for this translational skew.
263   *
264   * @param moduleStates Reference to array of module states. The array will be mutated with the
265   *     normalized speeds!
266   * @param attainableMaxSpeed The absolute max speed in meters per second that a module can reach.
267   */
268  public static void desaturateWheelSpeeds(
269      SwerveModuleState[] moduleStates, double attainableMaxSpeed) {
270    double realMaxSpeed = 0;
271    for (SwerveModuleState moduleState : moduleStates) {
272      realMaxSpeed = Math.max(realMaxSpeed, Math.abs(moduleState.speed));
273    }
274    if (realMaxSpeed > attainableMaxSpeed) {
275      for (SwerveModuleState moduleState : moduleStates) {
276        moduleState.speed = moduleState.speed / realMaxSpeed * attainableMaxSpeed;
277      }
278    }
279  }
280
281  /**
282   * Renormalizes the wheel speeds if any individual speed is above the specified maximum.
283   *
284   * <p>Sometimes, after inverse kinematics, the requested speed from one or more modules may be
285   * above the max attainable speed for the driving motor on that module. To fix this issue, one can
286   * reduce all the wheel speeds to make sure that all requested module speeds are at-or-below the
287   * absolute threshold, while maintaining the ratio of speeds between modules.
288   *
289   * <p>Scaling down the module speeds rotates the direction of net motion in the opposite direction
290   * of rotational velocity, which makes discretizing the chassis speeds inaccurate because the
291   * discretization did not account for this translational skew.
292   *
293   * @param moduleStates Reference to array of module states. The array will be mutated with the
294   *     normalized speeds!
295   * @param attainableMaxSpeed The absolute max speed in meters per second that a module can reach.
296   */
297  public static void desaturateWheelSpeeds(
298      SwerveModuleState[] moduleStates, LinearVelocity attainableMaxSpeed) {
299    desaturateWheelSpeeds(moduleStates, attainableMaxSpeed.in(MetersPerSecond));
300  }
301
302  /**
303   * Renormalizes the wheel speeds if any individual speed is above the specified maximum, as well
304   * as getting rid of joystick saturation at edges of joystick.
305   *
306   * <p>Sometimes, after inverse kinematics, the requested speed from one or more modules may be
307   * above the max attainable speed for the driving motor on that module. To fix this issue, one can
308   * reduce all the wheel speeds to make sure that all requested module speeds are at-or-below the
309   * absolute threshold, while maintaining the ratio of speeds between modules.
310   *
311   * <p>Scaling down the module speeds rotates the direction of net motion in the opposite direction
312   * of rotational velocity, which makes discretizing the chassis speeds inaccurate because the
313   * discretization did not account for this translational skew.
314   *
315   * @param moduleStates Reference to array of module states. The array will be mutated with the
316   *     normalized speeds!
317   * @param desiredChassisSpeed The desired speed of the robot
318   * @param attainableMaxModuleSpeed The absolute max speed in meters per second that a module can
319   *     reach
320   * @param attainableMaxTranslationalSpeed The absolute max speed in meters per second that your
321   *     robot can reach while translating
322   * @param attainableMaxRotationalVelocity The absolute max speed in radians per second the robot
323   *     can reach while rotating
324   */
325  public static void desaturateWheelSpeeds(
326      SwerveModuleState[] moduleStates,
327      ChassisSpeeds desiredChassisSpeed,
328      double attainableMaxModuleSpeed,
329      double attainableMaxTranslationalSpeed,
330      double attainableMaxRotationalVelocity) {
331    double realMaxSpeed = 0;
332    for (SwerveModuleState moduleState : moduleStates) {
333      realMaxSpeed = Math.max(realMaxSpeed, Math.abs(moduleState.speed));
334    }
335
336    if (attainableMaxTranslationalSpeed == 0
337        || attainableMaxRotationalVelocity == 0
338        || realMaxSpeed == 0) {
339      return;
340    }
341    double translationalK =
342        Math.hypot(desiredChassisSpeed.vx, desiredChassisSpeed.vy)
343            / attainableMaxTranslationalSpeed;
344    double rotationalK = Math.abs(desiredChassisSpeed.omega) / attainableMaxRotationalVelocity;
345    double k = Math.max(translationalK, rotationalK);
346    double scale = Math.min(k * attainableMaxModuleSpeed / realMaxSpeed, 1);
347    for (SwerveModuleState moduleState : moduleStates) {
348      moduleState.speed *= scale;
349    }
350  }
351
352  /**
353   * Renormalizes the wheel speeds if any individual speed is above the specified maximum, as well
354   * as getting rid of joystick saturation at edges of joystick.
355   *
356   * <p>Sometimes, after inverse kinematics, the requested speed from one or more modules may be
357   * above the max attainable speed for the driving motor on that module. To fix this issue, one can
358   * reduce all the wheel speeds to make sure that all requested module speeds are at-or-below the
359   * absolute threshold, while maintaining the ratio of speeds between modules.
360   *
361   * <p>Scaling down the module speeds rotates the direction of net motion in the opposite direction
362   * of rotational velocity, which makes discretizing the chassis speeds inaccurate because the
363   * discretization did not account for this translational skew.
364   *
365   * @param moduleStates Reference to array of module states. The array will be mutated with the
366   *     normalized speeds!
367   * @param desiredChassisSpeed The desired speed of the robot
368   * @param attainableMaxModuleSpeed The absolute max speed that a module can reach
369   * @param attainableMaxTranslationalSpeed The absolute max speed that your robot can reach while
370   *     translating
371   * @param attainableMaxRotationalVelocity The absolute max speed the robot can reach while
372   *     rotating
373   */
374  public static void desaturateWheelSpeeds(
375      SwerveModuleState[] moduleStates,
376      ChassisSpeeds desiredChassisSpeed,
377      LinearVelocity attainableMaxModuleSpeed,
378      LinearVelocity attainableMaxTranslationalSpeed,
379      AngularVelocity attainableMaxRotationalVelocity) {
380    desaturateWheelSpeeds(
381        moduleStates,
382        desiredChassisSpeed,
383        attainableMaxModuleSpeed.in(MetersPerSecond),
384        attainableMaxTranslationalSpeed.in(MetersPerSecond),
385        attainableMaxRotationalVelocity.in(RadiansPerSecond));
386  }
387
388  @Override
389  public SwerveModulePosition[] copy(SwerveModulePosition[] positions) {
390    var newPositions = new SwerveModulePosition[positions.length];
391    for (int i = 0; i < positions.length; ++i) {
392      newPositions[i] = positions[i].copy();
393    }
394    return newPositions;
395  }
396
397  @Override
398  public void copyInto(SwerveModulePosition[] positions, SwerveModulePosition[] output) {
399    if (positions.length != output.length) {
400      throw new IllegalArgumentException("Inconsistent number of modules!");
401    }
402    for (int i = 0; i < positions.length; ++i) {
403      output[i].distance = positions[i].distance;
404      output[i].angle = positions[i].angle;
405    }
406  }
407
408  @Override
409  public SwerveModulePosition[] interpolate(
410      SwerveModulePosition[] startValue, SwerveModulePosition[] endValue, double t) {
411    if (endValue.length != startValue.length) {
412      throw new IllegalArgumentException("Inconsistent number of modules!");
413    }
414    var newPositions = new SwerveModulePosition[startValue.length];
415    for (int i = 0; i < startValue.length; ++i) {
416      newPositions[i] = startValue[i].interpolate(endValue[i], t);
417    }
418    return newPositions;
419  }
420
421  /**
422   * Gets the locations of the modules relative to the center of rotation.
423   *
424   * @return The locations of the modules relative to the center of rotation. This array should not
425   *     be modified.
426   */
427  @SuppressWarnings("PMD.MethodReturnsInternalArray")
428  public Translation2d[] getModules() {
429    return m_modules;
430  }
431
432  /** SwerveDriveKinematics protobuf for serialization. */
433  public static final SwerveDriveKinematicsProto proto = new SwerveDriveKinematicsProto();
434
435  /**
436   * Creates an implementation of the {@link Struct} interface for swerve drive kinematics objects.
437   *
438   * @param numModules The number of modules of the kinematics objects this serializer processes.
439   * @return The struct implementation.
440   */
441  public static final SwerveDriveKinematicsStruct getStruct(int numModules) {
442    return new SwerveDriveKinematicsStruct(numModules);
443  }
444}