SimSwerveDrivetrain.hpp

Go to the documentation of this file.

/*
 * Copyright (C) Cross The Road Electronics.  All rights reserved.
 * License information can be found in CTRE_LICENSE.txt
 * For support and suggestions contact support@ctr-electronics.com or file
 * an issue tracker at https://github.com/CrossTheRoadElec/Phoenix-Releases
 */
#pragma once
 
#include "ctre/phoenix6/sim/Pigeon2SimState.hpp"
#include "ctre/phoenix6/swerve/impl/SwerveDriveKinematics.hpp"
#include "ctre/phoenix6/swerve/SwerveModule.hpp"
#include <wpi/simulation/DCMotorSim.hpp>
#include <wpi/math/system/Models.hpp>
 
namespace ctre {
namespace phoenix6 {
namespace swerve {
 
/**
 * \brief Simplified swerve drive simulation class.
 *
 * This class assumes that the swerve drive is perfect, meaning
 * that there is no scrub and the wheels do not slip.
 *
 * In addition, it assumes the inertia of the robot is governed only
 * by the inertia of the steer module and the individual drive wheels.
 * Robot-wide inertia is not accounted for, and neither is translational
 * vs rotational inertia of the robot.
 *
 * These assumptions provide a simplified example that can demonstrate the
 * behavior of a swerve drive in simulation. Users are encouraged to
 * expand this model for their own use.
 */
template <
    std::derived_from<hardware::traits::CommonTalon> DriveMotorT,
    std::derived_from<hardware::traits::CommonTalon> SteerMotorT,
    typename EncoderT
>
    requires std::same_as<EncoderT, hardware::CANcoder> ||
        std::same_as<EncoderT, hardware::CANdi> ||
        std::same_as<EncoderT, hardware::TalonFXS>
class SimSwerveDrivetrain {
protected:
    class SimSwerveModule {
    public:
        /** \brief Reference to motor simulation for drive motor */
        wpi::sim::DCMotorSim DriveMotor;
        /** \brief Reference to motor simulation for the steer motor */
        wpi::sim::DCMotorSim SteerMotor;
        /** \brief Reference to steer gearing for updating encoder */
        wpi::units::scalar_t DriveGearing;
        /** \brief Reference to steer gearing for updating encoder */
        wpi::units::scalar_t SteerGearing;
        /** \brief Voltage necessary for the drive motor to overcome friction */
        wpi::units::volt_t DriveFrictionVoltage;
        /** \brief Voltage necessary for the steer motor to overcome friction */
        wpi::units::volt_t SteerFrictionVoltage;
        /** \brief Whether the drive motor is inverted */
        bool DriveMotorInverted;
        /** \brief Whether the steer motor is inverted */
        bool SteerMotorInverted;
        /** \brief Whether the azimuth encoder is inverted */
        bool EncoderInverted;
        /** \brief Offset of the azimuth encoder */
        wpi::units::turn_t EncoderOffset;
        /** \brief The type of encoder to use for the azimuth */
        SteerFeedbackType EncoderType;
 
        SimSwerveModule(
            wpi::units::scalar_t driveGearing, wpi::units::kilogram_square_meter_t driveInertia,
            wpi::units::volt_t driveFrictionVoltage, bool driveMotorInverted,
            DriveMotorArrangement driveMotorType,
            wpi::units::scalar_t steerGearing, wpi::units::kilogram_square_meter_t steerInertia,
            wpi::units::volt_t steerFrictionVoltage, bool steerMotorInverted,
            SteerMotorArrangement steerMotorType,
            bool encoderInverted,
            wpi::units::turn_t encoderOffset,
            SteerFeedbackType encoderType
        ) :
            DriveMotor{
                wpi::math::Models::SingleJointedArmFromPhysicalConstants(GetDriveMotorGearbox(driveMotorType), driveInertia, driveGearing),
                GetDriveMotorGearbox(driveMotorType)
            },
            SteerMotor{
                wpi::math::Models::SingleJointedArmFromPhysicalConstants(GetSteerMotorGearbox(steerMotorType), steerInertia, steerGearing),
                GetSteerMotorGearbox(steerMotorType)
            },
            DriveGearing{driveGearing},
            SteerGearing{steerGearing},
            DriveFrictionVoltage{driveFrictionVoltage},
            SteerFrictionVoltage{steerFrictionVoltage},
            DriveMotorInverted{driveMotorInverted},
            SteerMotorInverted{steerMotorInverted},
            EncoderInverted{encoderInverted},
            EncoderOffset{encoderOffset},
            EncoderType{encoderType}
        {}
 
        static wpi::math::DCMotor GetDriveMotorGearbox(DriveMotorArrangement driveMotorType)
        {
            switch (driveMotorType) {
                case DriveMotorArrangement::TalonFX_Integrated:
                default:
                    return wpi::math::DCMotor::KrakenX60FOC(1);
                case DriveMotorArrangement::TalonFXS_NEO_JST:
                    return wpi::math::DCMotor::NEO(1);
                case DriveMotorArrangement::TalonFXS_VORTEX_JST:
                    return wpi::math::DCMotor::NeoVortex(1);
            }
        }
 
        static wpi::math::DCMotor GetSteerMotorGearbox(SteerMotorArrangement steerMotorType)
        {
            switch (steerMotorType) {
                case SteerMotorArrangement::TalonFX_Integrated:
                default:
                    return wpi::math::DCMotor::KrakenX60FOC(1);
                case SteerMotorArrangement::TalonFXS_Minion_JST:
                    return wpi::math::DCMotor{12_V, 3.1_Nm, 202_A, 4_A, 774_rad_per_s, 1};
                case SteerMotorArrangement::TalonFXS_NEO_JST:
                    return wpi::math::DCMotor::NEO(1);
                case SteerMotorArrangement::TalonFXS_VORTEX_JST:
                    return wpi::math::DCMotor::NeoVortex(1);
                case SteerMotorArrangement::TalonFXS_NEO550_JST:
                    return wpi::math::DCMotor::NEO550(1);
                case SteerMotorArrangement::TalonFXS_Brushed_AB:
                case SteerMotorArrangement::TalonFXS_Brushed_AC:
                case SteerMotorArrangement::TalonFXS_Brushed_BC:
                    return wpi::math::DCMotor::CIM(1);
            }
        }
    };
 
    sim::Pigeon2SimState &_pigeonSim;
    std::vector<SimSwerveModule> _modules;
 
    impl::SwerveDriveKinematics _kinem;
    Rotation2d _lastAngle{};
 
public:
    template <
        std::same_as<SwerveModuleConstants<
            typename DriveMotorT::Configuration,
            typename SteerMotorT::Configuration,
            typename EncoderT::Configuration
        >>... ModuleConstants
    >
    SimSwerveDrivetrain(
        std::vector<Translation2d> wheelLocations,
        sim::Pigeon2SimState &pigeonSim,
        ModuleConstants const &... moduleConstants
    ) :
        _pigeonSim{pigeonSim},
        _modules{
            SimSwerveModule{
                moduleConstants.DriveMotorGearRatio,
                moduleConstants.DriveInertia,
                moduleConstants.DriveFrictionVoltage,
                moduleConstants.DriveMotorInverted,
                moduleConstants.DriveMotorType,
                moduleConstants.SteerMotorGearRatio,
                moduleConstants.SteerInertia,
                moduleConstants.SteerFrictionVoltage,
                moduleConstants.SteerMotorInverted,
                moduleConstants.SteerMotorType,
                moduleConstants.EncoderInverted,
                moduleConstants.EncoderOffset,
                moduleConstants.FeedbackSource
            }...
        },
        _kinem{std::move(wheelLocations)}
    {}
 
    void Update(
        wpi::units::second_t dt,
        wpi::units::volt_t supplyVoltage,
        std::span<std::unique_ptr<SwerveModule<DriveMotorT, SteerMotorT, EncoderT>> const> modulesToApply
    ) {
        if (modulesToApply.size() != _modules.size()) return;
 
        std::vector<SwerveModuleVelocity> velocities(modulesToApply.size());
        /* update our sim devices */
        for (size_t i = 0; i < modulesToApply.size(); ++i) {
            auto &driveMotor = modulesToApply[i]->GetDriveMotor().GetSimState();
            auto &steerMotor = modulesToApply[i]->GetSteerMotor().GetSimState();
            auto &encoder = modulesToApply[i]->GetEncoder().GetSimState();
 
            if constexpr (std::same_as<hardware::TalonFXS, DriveMotorT>) {
                driveMotor.MotorOrientation = _modules[i].DriveMotorInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
            } else {
                driveMotor.Orientation = _modules[i].DriveMotorInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
            }
 
            if constexpr (std::same_as<hardware::TalonFXS, SteerMotorT>) {
                steerMotor.MotorOrientation = _modules[i].SteerMotorInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
                steerMotor.ExtSensorOrientation = _modules[i].SteerMotorInverted != _modules[i].EncoderInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
                steerMotor.PulseWidthSensorOffset = _modules[i].EncoderOffset;
            } else {
                steerMotor.Orientation = _modules[i].SteerMotorInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
            }
 
            if constexpr (std::same_as<hardware::CANcoder, EncoderT>) {
                encoder.Orientation = _modules[i].EncoderInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
                encoder.SensorOffset = _modules[i].EncoderOffset;
            }
 
            driveMotor.SetSupplyVoltage(supplyVoltage);
            steerMotor.SetSupplyVoltage(supplyVoltage);
            encoder.SetSupplyVoltage(supplyVoltage);
 
            _modules[i].DriveMotor.SetInputVoltage(AddFriction(driveMotor.GetMotorVoltage(), _modules[i].DriveFrictionVoltage));
            _modules[i].SteerMotor.SetInputVoltage(AddFriction(steerMotor.GetMotorVoltage(), _modules[i].SteerFrictionVoltage));
 
            _modules[i].DriveMotor.Update(dt);
            _modules[i].SteerMotor.Update(dt);
 
            driveMotor.SetRawRotorPosition(_modules[i].DriveMotor.GetAngularPosition() * _modules[i].DriveGearing);
            driveMotor.SetRotorVelocity(_modules[i].DriveMotor.GetAngularVelocity() * _modules[i].DriveGearing);
            driveMotor.SetRotorAcceleration(_modules[i].DriveMotor.GetAngularAcceleration() * _modules[i].DriveGearing);
 
            steerMotor.SetRawRotorPosition(_modules[i].SteerMotor.GetAngularPosition() * _modules[i].SteerGearing);
            steerMotor.SetRotorVelocity(_modules[i].SteerMotor.GetAngularVelocity() * _modules[i].SteerGearing);
            steerMotor.SetRotorAcceleration(_modules[i].SteerMotor.GetAngularAcceleration() * _modules[i].SteerGearing);
            if constexpr (std::same_as<hardware::TalonFXS, SteerMotorT>) {
                /* azimuth encoders see the mechanism, so don't account for the steer gearing */
                steerMotor.SetPulseWidthPosition(_modules[i].SteerMotor.GetAngularPosition());
                steerMotor.SetPulseWidthVelocity(_modules[i].SteerMotor.GetAngularVelocity());
            }
 
            if constexpr (std::same_as<hardware::CANcoder, EncoderT>) {
                /* azimuth encoders see the mechanism, so don't account for the steer gearing */
                encoder.SetRawPosition(_modules[i].SteerMotor.GetAngularPosition());
                encoder.SetVelocity(_modules[i].SteerMotor.GetAngularVelocity());
            } else if constexpr (std::same_as<hardware::CANdi, EncoderT>) {
                switch (_modules[i].EncoderType) {
                    case SteerFeedbackType::FusedCANdiPWM1:
                    case SteerFeedbackType::SyncCANdiPWM1:
                    case SteerFeedbackType::RemoteCANdiPWM1:
                        encoder.Pwm1Orientation = _modules[i].EncoderInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
                        encoder.Pwm1SensorOffset = _modules[i].EncoderOffset;
 
                        /* azimuth encoders see the mechanism, so don't account for the steer gearing */
                        encoder.SetPwm1Connected(true);
                        encoder.SetPwm1Position(_modules[i].SteerMotor.GetAngularPosition());
                        encoder.SetPwm1Velocity(_modules[i].SteerMotor.GetAngularVelocity());
                        break;
 
                    case SteerFeedbackType::FusedCANdiPWM2:
                    case SteerFeedbackType::SyncCANdiPWM2:
                    case SteerFeedbackType::RemoteCANdiPWM2:
                        encoder.Pwm2Orientation = _modules[i].EncoderInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
                        encoder.Pwm1SensorOffset = _modules[i].EncoderOffset;
 
                        /* azimuth encoders see the mechanism, so don't account for the steer gearing */
                        encoder.SetPwm2Connected(true);
                        encoder.SetPwm2Position(_modules[i].SteerMotor.GetAngularPosition());
                        encoder.SetPwm2Velocity(_modules[i].SteerMotor.GetAngularVelocity());
                        break;
 
                    default:
                        break;
                }
            }
 
            velocities[i] = modulesToApply[i]->GetCurrentVelocity();
        }
 
        auto const angularVel = _kinem.ToChassisVelocities(velocities).omega;
        _lastAngle = _lastAngle + Rotation2d{angularVel * dt};
        _pigeonSim.SetRawYaw(_lastAngle.Degrees());
        _pigeonSim.SetAngularVelocityZ(angularVel);
    }
 
protected:
    /**
     * \brief Applies the effects of friction to dampen the motor voltage.
     *
     * \param motorVoltage Voltage output by the motor
     * \param frictionVoltage Voltage required to overcome friction
     * \returns Friction-dampened motor voltage
     */
    static wpi::units::volt_t AddFriction(wpi::units::volt_t motorVoltage, wpi::units::volt_t frictionVoltage)
    {
        if (wpi::units::math::abs(motorVoltage) < frictionVoltage) {
            motorVoltage = 0_V;
        } else if (motorVoltage > 0_V) {
            motorVoltage -= frictionVoltage;
        } else {
            motorVoltage += frictionVoltage;
        }
        return motorVoltage;
    }
};
 
}
}
}