001/*
002 * Copyright (C) Cross The Road Electronics.  All rights reserved.
003 * License information can be found in CTRE_LICENSE.txt
004 * For support and suggestions contact support@ctr-electronics.com or file
005 * an issue tracker at https://github.com/CrossTheRoadElec/Phoenix-Releases
006 */
007package com.ctre.phoenix6.swerve;
008
009import com.ctre.phoenix6.sim.CANcoderSimState;
010import com.ctre.phoenix6.sim.ChassisReference;
011import com.ctre.phoenix6.sim.Pigeon2SimState;
012import com.ctre.phoenix6.sim.TalonFXSimState;
013
014import edu.wpi.first.math.geometry.Rotation2d;
015import edu.wpi.first.math.geometry.Translation2d;
016import edu.wpi.first.math.kinematics.SwerveDriveKinematics;
017import edu.wpi.first.math.kinematics.SwerveModuleState;
018import edu.wpi.first.math.system.plant.DCMotor;
019import edu.wpi.first.math.system.plant.LinearSystemId;
020import edu.wpi.first.math.util.Units;
021import edu.wpi.first.wpilibj.simulation.DCMotorSim;
022
023/**
024 * Simplified swerve drive simulation class.
025 * <p>
026 * This class assumes that the swerve drive is perfect, meaning
027 * that there is no scrub and the wheels do not slip.
028 * <p>
029 * In addition, it assumes the inertia of the robot is governed only
030 * by the inertia of the steer module and the individual drive wheels.
031 * Robot-wide inertia is not accounted for, and neither is translational
032 * vs rotational inertia of the robot.
033 * <p>
034 * These assumptions provide a simplified example that can demonstrate the
035 * behavior of a swerve drive in simulation. Users are encouraged to
036 * expand this model for their own use.
037 */
038public class SimSwerveDrivetrain {
039    protected static class SimSwerveModule {
040        /** Reference to motor simulation for drive motor */
041        public final DCMotorSim DriveMotor;
042        /** Reference to motor simulation for the steer motor */
043        public final DCMotorSim SteerMotor;
044        /** Reference to steer gearing for updating CANcoder */
045        public final double DriveGearing;
046        /** Reference to steer gearing for updating CANcoder */
047        public final double SteerGearing;
048        /** Voltage necessary for the drive motor to overcome friction */
049        public final double DriveFrictionVoltage;
050        /** Voltage necessary for the steer motor to overcome friction */
051        public final double SteerFrictionVoltage;
052        /** Whether the drive motor is inverted */
053        public final boolean DriveMotorInverted;
054        /** Whether the steer motor is inverted */
055        public final boolean SteerMotorInverted;
056        /** Whether the CANcoder is inverted */
057        public final boolean CANcoderInverted;
058
059
060        public SimSwerveModule(double driveGearing, double driveInertia,
061                            double driveFrictionVoltage, boolean driveMotorInverted,
062                            double steerGearing, double steerInertia,
063                            double steerFrictionVoltage, boolean steerMotorInverted,
064                            boolean cancoderInverted)
065        {
066            var gearbox = DCMotor.getKrakenX60Foc(1);
067
068            DriveMotor = new DCMotorSim(LinearSystemId.createDCMotorSystem(gearbox, driveInertia, driveGearing), gearbox);
069            SteerMotor = new DCMotorSim(LinearSystemId.createDCMotorSystem(gearbox, steerInertia, steerGearing), gearbox);
070            DriveGearing = driveGearing;
071            SteerGearing = steerGearing;
072            DriveFrictionVoltage = driveFrictionVoltage;
073            SteerFrictionVoltage = steerFrictionVoltage;
074            DriveMotorInverted = driveMotorInverted;
075            SteerMotorInverted = steerMotorInverted;
076            CANcoderInverted = cancoderInverted;
077        }
078    }
079
080    protected final Pigeon2SimState m_pigeonSim;
081    protected final SimSwerveModule[] m_modules;
082
083    protected final SwerveDriveKinematics m_kinem;
084    protected Rotation2d m_lastAngle = new Rotation2d();
085
086    public SimSwerveDrivetrain(Translation2d[] wheelLocations,
087                               Pigeon2SimState pigeonSim,
088                               SwerveModuleConstants... moduleConstants) {
089        m_pigeonSim = pigeonSim;
090        m_modules = new SimSwerveModule[moduleConstants.length];
091        for (int i = 0; i < m_modules.length; ++i) {
092            m_modules[i] = new SimSwerveModule(
093                moduleConstants[i].DriveMotorGearRatio, moduleConstants[i].DriveInertia,
094                moduleConstants[i].DriveFrictionVoltage, moduleConstants[i].DriveMotorInverted,
095                moduleConstants[i].SteerMotorGearRatio, moduleConstants[i].SteerInertia,
096                moduleConstants[i].SteerFrictionVoltage, moduleConstants[i].SteerMotorInverted,
097                moduleConstants[i].CANcoderInverted
098            );
099        }
100
101        m_kinem = new SwerveDriveKinematics(wheelLocations);
102    }
103
104    /**
105     * Update this simulation for the time duration.
106     * <p>
107     * This performs a simulation update on all the simulated devices
108     *
109     * @param dtSeconds The time delta between this update and the previous update
110     * @param supplyVoltage The voltage as seen at the motor controllers
111     * @param modulesToApply What modules to apply the update to
112     */
113    public void update(double dtSeconds, double supplyVoltage, SwerveModule ... modulesToApply) {
114        if (modulesToApply.length != m_modules.length) return;
115
116        SwerveModuleState[] states = new SwerveModuleState[m_modules.length];
117        /* Update our sim devices */
118        for (int i = 0; i < m_modules.length; ++i) {
119            TalonFXSimState driveMotor = modulesToApply[i].getDriveMotor().getSimState();
120            TalonFXSimState steerMotor = modulesToApply[i].getSteerMotor().getSimState();
121            CANcoderSimState cancoder = modulesToApply[i].getCANcoder().getSimState();
122
123            driveMotor.Orientation = m_modules[i].DriveMotorInverted ? ChassisReference.Clockwise_Positive : ChassisReference.CounterClockwise_Positive;
124            steerMotor.Orientation = m_modules[i].SteerMotorInverted ? ChassisReference.Clockwise_Positive : ChassisReference.CounterClockwise_Positive;
125            cancoder.Orientation = m_modules[i].CANcoderInverted ? ChassisReference.Clockwise_Positive : ChassisReference.CounterClockwise_Positive;
126
127            driveMotor.setSupplyVoltage(supplyVoltage);
128            steerMotor.setSupplyVoltage(supplyVoltage);
129            cancoder.setSupplyVoltage(supplyVoltage);
130
131            m_modules[i].DriveMotor.setInputVoltage(addFriction(driveMotor.getMotorVoltage(), m_modules[i].DriveFrictionVoltage));
132            m_modules[i].SteerMotor.setInputVoltage(addFriction(steerMotor.getMotorVoltage(), m_modules[i].SteerFrictionVoltage));
133
134            m_modules[i].DriveMotor.update(dtSeconds);
135            m_modules[i].SteerMotor.update(dtSeconds);
136
137            driveMotor.setRawRotorPosition(m_modules[i].DriveMotor.getAngularPositionRotations() * m_modules[i].DriveGearing);
138            driveMotor.setRotorVelocity(m_modules[i].DriveMotor.getAngularVelocityRPM() / 60.0 * m_modules[i].DriveGearing);
139
140            steerMotor.setRawRotorPosition(m_modules[i].SteerMotor.getAngularPositionRotations() * m_modules[i].SteerGearing);
141            steerMotor.setRotorVelocity(m_modules[i].SteerMotor.getAngularVelocityRPM() / 60.0 * m_modules[i].SteerGearing);
142
143            /* CANcoders see the mechanism, so don't account for the steer gearing */
144            cancoder.setRawPosition(m_modules[i].SteerMotor.getAngularPositionRotations());
145            cancoder.setVelocity(m_modules[i].SteerMotor.getAngularVelocityRPM() / 60.0);
146
147            states[i] = modulesToApply[i].getCurrentState();
148        }
149
150        double angularVelRadPerSec = m_kinem.toChassisSpeeds(states).omegaRadiansPerSecond;
151        m_lastAngle = m_lastAngle.plus(Rotation2d.fromRadians(angularVelRadPerSec * dtSeconds));
152        m_pigeonSim.setRawYaw(m_lastAngle.getDegrees());
153        m_pigeonSim.setAngularVelocityZ(Units.radiansToDegrees(angularVelRadPerSec));
154    }
155
156    /**
157     * Applies the effects of friction to dampen the motor voltage.
158     *
159     * @param motorVoltage Voltage output by the motor
160     * @param frictionVoltage Voltage required to overcome friction
161     * @return Friction-dampened motor voltage
162     */
163    protected double addFriction(double motorVoltage, double frictionVoltage) {
164        if (Math.abs(motorVoltage) < frictionVoltage) {
165            motorVoltage = 0.0;
166        } else if (motorVoltage > 0.0) {
167            motorVoltage -= frictionVoltage;
168        } else {
169            motorVoltage += frictionVoltage;
170        }
171        return motorVoltage;
172    }
173}