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.Rotation2d;
011import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
012import edu.wpi.first.math.kinematics.SwerveDriveOdometry;
013import edu.wpi.first.math.kinematics.SwerveDriveWheelPositions;
014import edu.wpi.first.math.kinematics.SwerveModulePosition;
015import edu.wpi.first.math.numbers.N1;
016import edu.wpi.first.math.numbers.N3;
017
018/**
019 * This class wraps {@link SwerveDriveOdometry Swerve Drive Odometry} to fuse latency-compensated
020 * vision measurements with swerve drive encoder distance measurements. It is intended to be a
021 * drop-in replacement for {@link edu.wpi.first.math.kinematics.SwerveDriveOdometry}.
022 *
023 * <p>{@link SwerveDrivePoseEstimator#update} should be called every robot loop.
024 *
025 * <p>{@link SwerveDrivePoseEstimator#addVisionMeasurement} can be called as infrequently as you
026 * want; if you never call it, then this class will behave as regular encoder odometry.
027 */
028public class SwerveDrivePoseEstimator extends PoseEstimator<SwerveDriveWheelPositions> {
029  private final int m_numModules;
030
031  /**
032   * Constructs a SwerveDrivePoseEstimator with default standard deviations for the model and vision
033   * measurements.
034   *
035   * <p>The default standard deviations of the model states are 0.1 meters for x, 0.1 meters for y,
036   * and 0.1 radians for heading. The default standard deviations of the vision measurements are 0.9
037   * meters for x, 0.9 meters for y, and 0.9 radians for heading.
038   *
039   * @param kinematics A correctly-configured kinematics object for your drivetrain.
040   * @param gyroAngle The current gyro angle.
041   * @param modulePositions The current distance measurements and rotations of the swerve modules.
042   * @param initialPoseMeters The starting pose estimate.
043   */
044  public SwerveDrivePoseEstimator(
045      SwerveDriveKinematics kinematics,
046      Rotation2d gyroAngle,
047      SwerveModulePosition[] modulePositions,
048      Pose2d initialPoseMeters) {
049    this(
050        kinematics,
051        gyroAngle,
052        modulePositions,
053        initialPoseMeters,
054        VecBuilder.fill(0.1, 0.1, 0.1),
055        VecBuilder.fill(0.9, 0.9, 0.9));
056  }
057
058  /**
059   * Constructs a SwerveDrivePoseEstimator.
060   *
061   * @param kinematics A correctly-configured kinematics object for your drivetrain.
062   * @param gyroAngle The current gyro angle.
063   * @param modulePositions The current distance and rotation measurements of the swerve modules.
064   * @param initialPoseMeters The starting pose estimate.
065   * @param stateStdDevs Standard deviations of the pose estimate (x position in meters, y position
066   *     in meters, and heading in radians). Increase these numbers to trust your state estimate
067   *     less.
068   * @param visionMeasurementStdDevs Standard deviations of the vision pose measurement (x position
069   *     in meters, y position in meters, and heading in radians). Increase these numbers to trust
070   *     the vision pose measurement less.
071   */
072  public SwerveDrivePoseEstimator(
073      SwerveDriveKinematics kinematics,
074      Rotation2d gyroAngle,
075      SwerveModulePosition[] modulePositions,
076      Pose2d initialPoseMeters,
077      Matrix<N3, N1> stateStdDevs,
078      Matrix<N3, N1> visionMeasurementStdDevs) {
079    super(
080        kinematics,
081        new SwerveDriveOdometry(kinematics, gyroAngle, modulePositions, initialPoseMeters),
082        stateStdDevs,
083        visionMeasurementStdDevs);
084
085    m_numModules = modulePositions.length;
086  }
087
088  /**
089   * Resets the robot's position on the field.
090   *
091   * <p>The gyroscope angle does not need to be reset in the user's robot code. The library
092   * automatically takes care of offsetting the gyro angle.
093   *
094   * @param gyroAngle The angle reported by the gyroscope.
095   * @param modulePositions The current distance measurements and rotations of the swerve modules.
096   * @param poseMeters The position on the field that your robot is at.
097   */
098  public void resetPosition(
099      Rotation2d gyroAngle, SwerveModulePosition[] modulePositions, Pose2d poseMeters) {
100    resetPosition(gyroAngle, new SwerveDriveWheelPositions(modulePositions), poseMeters);
101  }
102
103  /**
104   * Updates the pose estimator with wheel encoder and gyro information. This should be called every
105   * loop.
106   *
107   * @param gyroAngle The current gyro angle.
108   * @param modulePositions The current distance measurements and rotations of the swerve modules.
109   * @return The estimated pose of the robot in meters.
110   */
111  public Pose2d update(Rotation2d gyroAngle, SwerveModulePosition[] modulePositions) {
112    return update(gyroAngle, new SwerveDriveWheelPositions(modulePositions));
113  }
114
115  /**
116   * Updates the pose estimator with wheel encoder and gyro information. This should be called every
117   * loop.
118   *
119   * @param currentTimeSeconds Time at which this method was called, in seconds.
120   * @param gyroAngle The current gyroscope angle.
121   * @param modulePositions The current distance measurements and rotations of the swerve modules.
122   * @return The estimated pose of the robot in meters.
123   */
124  public Pose2d updateWithTime(
125      double currentTimeSeconds, Rotation2d gyroAngle, SwerveModulePosition[] modulePositions) {
126    return updateWithTime(
127        currentTimeSeconds, gyroAngle, new SwerveDriveWheelPositions(modulePositions));
128  }
129
130  @Override
131  public Pose2d updateWithTime(
132      double currentTimeSeconds, Rotation2d gyroAngle, SwerveDriveWheelPositions wheelPositions) {
133    if (wheelPositions.positions.length != m_numModules) {
134      throw new IllegalArgumentException(
135          "Number of modules is not consistent with number of wheel locations provided in "
136              + "constructor");
137    }
138
139    return super.updateWithTime(currentTimeSeconds, gyroAngle, wheelPositions);
140  }
141}