Files
allwpilib/robotpyExamples/SwerveBot/swervemodule.py
Gold856 056d7bcbbe [wpimath] Make swerve and differential kinematics functions immutable (#8274)
Originally started with just swerve, but expanded to diff and mecanum
(docs only) for parity across the drivetrains. Return value checks are
applied when possible to make migration easier and to error loudly if
people forget.

---------

Co-authored-by: Joseph Eng <91924258+KangarooKoala@users.noreply.github.com>
2026-04-16 21:23:32 -07:00

133 lines
5.2 KiB
Python

#
# Copyright (c) FIRST and other WPILib contributors.
# Open Source Software; you can modify and/or share it under the terms of
# the WPILib BSD license file in the root directory of this project.
#
import math
import wpilib
import wpimath
kWheelRadius = 0.0508
kEncoderResolution = 4096
kModuleMaxAngularVelocity = math.pi
kModuleMaxAngularAcceleration = math.tau
class SwerveModule:
def __init__(
self,
driveMotorChannel: int,
turningMotorChannel: int,
driveEncoderChannelA: int,
driveEncoderChannelB: int,
turningEncoderChannelA: int,
turningEncoderChannelB: int,
) -> None:
"""Constructs a SwerveModule with a drive motor, turning motor, drive encoder and turning encoder.
:param driveMotorChannel: PWM output for the drive motor.
:param turningMotorChannel: PWM output for the turning motor.
:param driveEncoderChannelA: DIO input for the drive encoder channel A
:param driveEncoderChannelB: DIO input for the drive encoder channel B
:param turningEncoderChannelA: DIO input for the turning encoder channel A
:param turningEncoderChannelB: DIO input for the turning encoder channel B
"""
self.driveMotor = wpilib.PWMSparkMax(driveMotorChannel)
self.turningMotor = wpilib.PWMSparkMax(turningMotorChannel)
self.driveEncoder = wpilib.Encoder(driveEncoderChannelA, driveEncoderChannelB)
self.turningEncoder = wpilib.Encoder(
turningEncoderChannelA, turningEncoderChannelB
)
# Gains are for example purposes only - must be determined for your own robot!
self.drivePIDController = wpimath.PIDController(1, 0, 0)
# Gains are for example purposes only - must be determined for your own robot!
self.turningPIDController = wpimath.ProfiledPIDController(
1,
0,
0,
wpimath.TrapezoidProfile.Constraints(
kModuleMaxAngularVelocity,
kModuleMaxAngularAcceleration,
),
)
# Gains are for example purposes only - must be determined for your own robot!
self.driveFeedforward = wpimath.SimpleMotorFeedforwardMeters(1, 3)
self.turnFeedforward = wpimath.SimpleMotorFeedforwardMeters(1, 0.5)
# Set the distance per pulse for the drive encoder. We can simply use the
# distance traveled for one rotation of the wheel divided by the encoder
# resolution.
self.driveEncoder.setDistancePerPulse(
math.tau * kWheelRadius / kEncoderResolution
)
# Set the distance (in this case, angle) in radians per pulse for the turning encoder.
# This is the the angle through an entire rotation (2 * pi) divided by the
# encoder resolution.
self.turningEncoder.setDistancePerPulse(math.tau / kEncoderResolution)
# Limit the PID Controller's input range between -pi and pi and set the input
# to be continuous.
self.turningPIDController.enableContinuousInput(-math.pi, math.pi)
def getVelocity(self) -> wpimath.SwerveModuleVelocity:
"""Returns the current velocity of the module.
:returns: The current velocity of the module.
"""
return wpimath.SwerveModuleVelocity(
self.driveEncoder.getRate(),
wpimath.Rotation2d(self.turningEncoder.getDistance()),
)
def getPosition(self) -> wpimath.SwerveModulePosition:
"""Returns the current position of the module.
:returns: The current position of the module.
"""
return wpimath.SwerveModulePosition(
self.driveEncoder.getDistance(),
wpimath.Rotation2d(self.turningEncoder.getDistance()),
)
def setDesiredVelocity(self, desiredVelocity: wpimath.SwerveModuleVelocity) -> None:
"""Sets the desired velocity for the module.
:param desiredVelocity: Desired state with velocity and angle.
"""
encoderRotation = wpimath.Rotation2d(self.turningEncoder.getDistance())
# Optimize the desired velocity to avoid spinning further than 90 degrees, then scale
# velocity by cosine of angle error. This scales down movement perpendicular to the desired
# direction of travel that can occur when modules change directions. This results in
# smoother driving.
velocity = desiredVelocity.optimize(encoderRotation).cosineScale(
encoderRotation
)
# Calculate the drive output from the drive PID controller.
driveOutput = self.drivePIDController.calculate(
self.driveEncoder.getRate(), velocity.velocity
)
driveFeedforward = self.driveFeedforward.calculate(velocity.velocity)
# Calculate the turning motor output from the turning PID controller.
turnOutput = self.turningPIDController.calculate(
self.turningEncoder.getDistance(), velocity.angle.radians()
)
turnFeedforward = self.turnFeedforward.calculate(
self.turningPIDController.getSetpoint().velocity
)
self.driveMotor.setVoltage(driveOutput + driveFeedforward)
self.turningMotor.setVoltage(turnOutput + turnFeedforward)