-
Notifications
You must be signed in to change notification settings - Fork 0
/
swervemodule.py
130 lines (104 loc) · 4.92 KB
/
swervemodule.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#
# 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.kinematics
import wpimath.geometry
import wpimath.controller
import wpimath.trajectory
import wpimath.units
import rev
kWheelRadius = wpimath.units.inchesToMeters(4)
kEncoderResolution = 4096
kModuleMaxAngularVelocity = math.pi
kModuleMaxAngularAcceleration = math.tau
class SwerveModule:
def __init__(
self,
driveMotorId: int,
steerMotorId: int,
) -> None:
"""Constructs a SwerveModule with a drive motor, steer motor, drive encoder and steer encoder.
:param driveMotorId: CAN id for the drive motor.
:param steerMotorId: CAN id for the steer motor.
"""
self.driveMotor = rev.CANSparkMax(driveMotorId, rev.CANSparkLowLevel.MotorType.kBrushless)
self.steerMotor = rev.CANSparkMax(steerMotorId, rev.CANSparkLowLevel.MotorType.kBrushless)
self.driveEncoder = self.driveMotor.getEncoder()
self.steerEncoder = self.steerMotor.getEncoder()
# Gains are for example purposes only - must be determined for your own robot!
self.drivePIDController = wpimath.controller.PIDController(1, 0, 0)
# Gains are for example purposes only - must be determined for your own robot!
self.steerPIDController = wpimath.controller.ProfiledPIDController(
1,
0,
0,
wpimath.trajectory.TrapezoidProfile.Constraints(
kModuleMaxAngularVelocity,
kModuleMaxAngularAcceleration,
),
)
# Gains are for example purposes only - must be determined for your own robot!
self.driveFeedforward = wpimath.controller.SimpleMotorFeedforwardMeters(1, 3)
self.turnFeedforward = wpimath.controller.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.setPositionConversionFactor(
math.tau * kWheelRadius / 7.13
)
# Set the distance (in this case, angle) in radians per pulse for the steer encoder.
# This is the the angle through an entire rotation (2 * pi) divided by the
# encoder resolution.
self.steerEncoder.setPositionConversionFactor(math.tau / 11.3143)
# Limit the PID Controller's input range between -pi and pi and set the input
# to be continuous.
self.steerPIDController.enableContinuousInput(-math.pi, math.pi)
def getState(self) -> wpimath.kinematics.SwerveModuleState:
"""Returns the current state of the module.
:returns: The current state of the module.
"""
return wpimath.kinematics.SwerveModuleState(
self.driveEncoder.getVelocity(),
wpimath.geometry.Rotation2d(self.steerEncoder.getPosition()),
)
def getPosition(self) -> wpimath.kinematics.SwerveModulePosition:
"""Returns the current position of the module.
:returns: The current position of the module.
"""
return wpimath.kinematics.SwerveModulePosition(
self.driveEncoder.getVelocity(),
wpimath.geometry.Rotation2d(self.steerEncoder.getPosition()),
)
def setDesiredState(
self, desiredState: wpimath.kinematics.SwerveModuleState
) -> None:
"""Sets the desired state for the module.
:param desiredState: Desired state with speed and angle.
"""
encoderRotation = wpimath.geometry.Rotation2d(self.steerEncoder.getVelocity())
# Optimize the reference state to avoid spinning further than 90 degrees
state = wpimath.kinematics.SwerveModuleState.optimize(
desiredState, encoderRotation
)
# Scale speed 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.
state.speed *= (state.angle - encoderRotation).cos()
# Calculate the drive output from the drive PID controller.
driveOutput = self.drivePIDController.calculate(
self.driveEncoder.getVelocity(), state.speed
)
driveFeedforward = self.driveFeedforward.calculate(state.speed)
# Calculate the steer motor output from the steer PID controller.
turnOutput = self.steerPIDController.calculate(
self.steerEncoder.getPosition(), state.angle.radians()
)
turnFeedforward = self.turnFeedforward.calculate(
self.steerPIDController.getSetpoint().velocity
)
self.driveMotor.setVoltage(driveOutput + driveFeedforward)
self.steerMotor.setVoltage(turnOutput + turnFeedforward)