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.estimator;
006
007import edu.wpi.first.math.Matrix;
008import edu.wpi.first.math.VecBuilder;
009import edu.wpi.first.math.geometry.Pose2d;
010import edu.wpi.first.math.geometry.Pose3d;
011import edu.wpi.first.math.geometry.Rotation2d;
012import edu.wpi.first.math.geometry.Rotation3d;
013import edu.wpi.first.math.geometry.Translation2d;
014import edu.wpi.first.math.geometry.Translation3d;
015import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
016import edu.wpi.first.math.kinematics.SwerveDriveOdometry3d;
017import edu.wpi.first.math.kinematics.SwerveModulePosition;
018import edu.wpi.first.math.numbers.N1;
019import edu.wpi.first.math.numbers.N4;
020
021/**
022 * This class wraps {@link SwerveDriveOdometry3d Swerve Drive Odometry} to fuse latency-compensated
023 * vision measurements with swerve drive encoder distance measurements. It is intended to be a
024 * drop-in replacement for {@link SwerveDriveOdometry3d}. It is also intended to be an easy
025 * replacement for {@link SwerveDrivePoseEstimator}, only requiring the addition of a standard
026 * deviation for Z and appropriate conversions between 2D and 3D versions of geometry classes. (See
027 * {@link Pose3d#Pose3d(Pose2d)}, {@link Rotation3d#Rotation3d(Rotation2d)}, {@link
028 * Translation3d#Translation3d(Translation2d)}, and {@link Pose3d#toPose2d()}.)
029 *
030 * <p>{@link SwerveDrivePoseEstimator3d#update} should be called every robot loop.
031 *
032 * <p>{@link SwerveDrivePoseEstimator3d#addVisionMeasurement} can be called as infrequently as you
033 * want; if you never call it, then this class will behave as regular encoder odometry.
034 */
035public class SwerveDrivePoseEstimator3d extends PoseEstimator3d<SwerveModulePosition[]> {
036  private final int m_numModules;
037
038  /**
039   * Constructs a SwerveDrivePoseEstimator3d with default standard deviations for the model and
040   * vision measurements.
041   *
042   * <p>The default standard deviations of the model states are 0.1 meters for x, 0.1 meters for y,
043   * 0.1 meters for z, and 0.1 radians for angle. The default standard deviations of the vision
044   * measurements are 0.9 meters for x, 0.9 meters for y, 0.9 meters for z, and 0.9 radians for
045   * angle.
046   *
047   * @param kinematics A correctly-configured kinematics object for your drivetrain.
048   * @param gyroAngle The current gyro angle.
049   * @param modulePositions The current distance measurements and rotations of the swerve modules.
050   * @param initialPoseMeters The starting pose estimate.
051   */
052  public SwerveDrivePoseEstimator3d(
053      SwerveDriveKinematics kinematics,
054      Rotation3d gyroAngle,
055      SwerveModulePosition[] modulePositions,
056      Pose3d initialPoseMeters) {
057    this(
058        kinematics,
059        gyroAngle,
060        modulePositions,
061        initialPoseMeters,
062        VecBuilder.fill(0.1, 0.1, 0.1, 0.1),
063        VecBuilder.fill(0.9, 0.9, 0.9, 0.9));
064  }
065
066  /**
067   * Constructs a SwerveDrivePoseEstimator3d.
068   *
069   * @param kinematics A correctly-configured kinematics object for your drivetrain.
070   * @param gyroAngle The current gyro angle.
071   * @param modulePositions The current distance and rotation measurements of the swerve modules.
072   * @param initialPoseMeters The starting pose estimate.
073   * @param stateStdDevs Standard deviations of the pose estimate (x position in meters, y position
074   *     in meters, and heading in radians). Increase these numbers to trust your state estimate
075   *     less.
076   * @param visionMeasurementStdDevs Standard deviations of the vision pose measurement (x position
077   *     in meters, y position in meters, and heading in radians). Increase these numbers to trust
078   *     the vision pose measurement less.
079   */
080  public SwerveDrivePoseEstimator3d(
081      SwerveDriveKinematics kinematics,
082      Rotation3d gyroAngle,
083      SwerveModulePosition[] modulePositions,
084      Pose3d initialPoseMeters,
085      Matrix<N4, N1> stateStdDevs,
086      Matrix<N4, N1> visionMeasurementStdDevs) {
087    super(
088        kinematics,
089        new SwerveDriveOdometry3d(kinematics, gyroAngle, modulePositions, initialPoseMeters),
090        stateStdDevs,
091        visionMeasurementStdDevs);
092
093    m_numModules = modulePositions.length;
094  }
095
096  @Override
097  public Pose3d updateWithTime(
098      double currentTimeSeconds, Rotation3d gyroAngle, SwerveModulePosition[] wheelPositions) {
099    if (wheelPositions.length != m_numModules) {
100      throw new IllegalArgumentException(
101          "Number of modules is not consistent with number of wheel locations provided in "
102              + "constructor");
103    }
104
105    return super.updateWithTime(currentTimeSeconds, gyroAngle, wheelPositions);
106  }
107}