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