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/core/CorePigeon2.hpp"
#include "ctre/phoenix6/swerve/impl/SwerveDriveKinematics.hpp"
#include "ctre/phoenix6/swerve/SwerveModule.hpp"
#include <frc/simulation/DCMotorSim.h>
#include <frc/system/plant/LinearSystemId.h>
 
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.
 */
class SimSwerveDrivetrain {
protected:
    class SimSwerveModule {
    public:
        /** \brief Reference to motor simulation for drive motor */
        frc::sim::DCMotorSim DriveMotor;
        /** \brief Reference to motor simulation for the steer motor */
        frc::sim::DCMotorSim SteerMotor;
        /** \brief Reference to steer gearing for updating encoder */
        units::scalar_t DriveGearing;
        /** \brief Reference to steer gearing for updating encoder */
        units::scalar_t SteerGearing;
        /** \brief Voltage necessary for the drive motor to overcome friction */
        units::volt_t DriveFrictionVoltage;
        /** \brief Voltage necessary for the steer motor to overcome friction */
        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;
 
        SimSwerveModule(
            units::scalar_t driveGearing, units::kilogram_square_meter_t driveInertia,
            units::volt_t driveFrictionVoltage, bool driveMotorInverted,
            units::scalar_t steerGearing, units::kilogram_square_meter_t steerInertia,
            units::volt_t steerFrictionVoltage, bool steerMotorInverted,
            bool encoderInverted
        ) :
                DriveMotor{frc::LinearSystemId::DCMotorSystem(frc::DCMotor::KrakenX60FOC(1), driveInertia, driveGearing), frc::DCMotor::KrakenX60FOC(1)},
                SteerMotor{frc::LinearSystemId::DCMotorSystem(frc::DCMotor::KrakenX60FOC(1), steerInertia, steerGearing), frc::DCMotor::KrakenX60FOC(1)},
                DriveGearing{driveGearing},
                SteerGearing{steerGearing},
                DriveFrictionVoltage{driveFrictionVoltage},
                SteerFrictionVoltage{steerFrictionVoltage},
                DriveMotorInverted{driveMotorInverted},
                SteerMotorInverted{steerMotorInverted},
                EncoderInverted{encoderInverted}
        {}
    };
 
    sim::Pigeon2SimState &_pigeonSim;
    std::vector<SimSwerveModule> _modules;
 
    impl::SwerveDriveKinematics _kinem;
    Rotation2d _lastAngle{};
 
public:
    template <
        typename... ModuleConstants,
        typename = std::enable_if_t<std::conjunction_v<
            std::is_same< ModuleConstants, SwerveModuleConstants<configs::TalonFXConfiguration, configs::TalonFXConfiguration, configs::CANcoderConfiguration> >...
        >>
    >
    SimSwerveDrivetrain(
        std::vector<Translation2d> wheelLocations,
        sim::Pigeon2SimState &pigeonSim,
        ModuleConstants const &... moduleConstants
    ) :
        _pigeonSim{pigeonSim},
        _modules{
            SimSwerveModule{
                moduleConstants.DriveMotorGearRatio,
                moduleConstants.DriveInertia,
                moduleConstants.DriveFrictionVoltage,
                moduleConstants.DriveMotorInverted,
                moduleConstants.SteerMotorGearRatio,
                moduleConstants.SteerInertia,
                moduleConstants.SteerFrictionVoltage,
                moduleConstants.SteerMotorInverted,
                moduleConstants.EncoderInverted
            }...
        },
        _kinem{std::move(wheelLocations)}
    {}
 
    template <
        typename DriveMotorT,
        typename SteerMotorT,
        typename EncoderT,
        typename = std::enable_if_t<std::is_base_of_v<hardware::traits::CommonTalon, DriveMotorT>>,
        typename = std::enable_if_t<std::is_base_of_v<hardware::traits::CommonTalon, SteerMotorT>>,
        typename = std::enable_if_t<std::is_base_of_v<hardware::ParentDevice, EncoderT>>
    >
    void Update(
        units::second_t dt,
        units::volt_t supplyVoltage,
        std::vector<std::unique_ptr<SwerveModule<DriveMotorT, SteerMotorT, EncoderT>>> const &modulesToApply
    ) {
        if (modulesToApply.size() != _modules.size()) return;
 
        if constexpr (!(std::is_same_v<hardware::TalonFX, DriveMotorT> &&
            std::is_same_v<hardware::TalonFX, SteerMotorT> &&
            std::is_same_v<hardware::CANcoder, EncoderT>))
        {
            return;
        }
 
        std::vector<SwerveModuleState> states(modulesToApply.size());
        /* update our sim devices */
        for (size_t i = 0; i < modulesToApply.size(); ++i) {
            sim::TalonFXSimState &driveMotor = modulesToApply[i]->GetDriveMotor().GetSimState();
            sim::TalonFXSimState &steerMotor = modulesToApply[i]->GetSteerMotor().GetSimState();
            sim::CANcoderSimState &encoder = modulesToApply[i]->GetEncoder().GetSimState();
 
            driveMotor.Orientation = _modules[i].DriveMotorInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
            steerMotor.Orientation = _modules[i].SteerMotorInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
            encoder.Orientation = _modules[i].EncoderInverted ? sim::ChassisReference::Clockwise_Positive : sim::ChassisReference::CounterClockwise_Positive;
 
            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);
 
            steerMotor.SetRawRotorPosition(_modules[i].SteerMotor.GetAngularPosition() * _modules[i].SteerGearing);
            steerMotor.SetRotorVelocity(_modules[i].SteerMotor.GetAngularVelocity() * _modules[i].SteerGearing);
 
            /* 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());
 
            states[i] = modulesToApply[i]->GetCurrentState();
        }
 
        auto const angularVel = _kinem.ToChassisSpeeds(states).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 units::volt_t AddFriction(units::volt_t motorVoltage, units::volt_t frictionVoltage)
    {
        if (units::math::abs(motorVoltage) < frictionVoltage) {
            motorVoltage = 0_V;
        } else if (motorVoltage > 0_V) {
            motorVoltage -= frictionVoltage;
        } else {
            motorVoltage += frictionVoltage;
        }
        return motorVoltage;
    }
};
 
}
}
}