Files
allwpilib/robotpyExamples/ArmSimulation/subsystems/arm.py
2026-03-20 13:24:22 -06:00

115 lines
4.4 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 wpilib
import wpilib.simulation
import wpimath
import wpimath.units
import wpiutil
from constants import Constants
class Arm:
def __init__(self) -> None:
# The P gain for the PID controller that drives this arm.
self.armKp = Constants.kDefaultArmKp
self.armSetpointDegrees = Constants.kDefaultArmSetpointDegrees
# The arm gearbox represents a gearbox containing two Vex 775pro motors.
self.armGearbox = wpimath.DCMotor.vex775Pro(2)
# Standard classes for controlling our arm
self.controller = wpimath.PIDController(self.armKp, 0, 0)
self.encoder = wpilib.Encoder(
Constants.kEncoderAChannel, Constants.kEncoderBChannel
)
self.motor = wpilib.PWMSparkMax(Constants.kMotorPort)
# Simulation classes help us simulate what's going on, including gravity.
# This arm sim represents an arm that can travel from -75 degrees (rotated down front)
# to 255 degrees (rotated down in the back).
self.armSim = wpilib.simulation.SingleJointedArmSim(
self.armGearbox,
Constants.kArmReduction,
wpilib.simulation.SingleJointedArmSim.estimateMOI(
Constants.kArmLength, Constants.kArmMass
),
Constants.kArmLength,
Constants.kMinAngleRads,
Constants.kMaxAngleRads,
True,
0,
[Constants.kArmEncoderDistPerPulse, 0.0],
)
self.encoderSim = wpilib.simulation.EncoderSim(self.encoder)
# Create a Mechanism2d display of an Arm with a fixed ArmTower and moving Arm.
self.mech2d = wpilib.Mechanism2d(60, 60)
self.armPivot = self.mech2d.getRoot("ArmPivot", 30, 30)
self.armTower = self.armPivot.appendLigament("ArmTower", 30, -90)
self.armLigament = self.armPivot.appendLigament(
"Arm",
30,
wpimath.units.radiansToDegrees(self.armSim.getAngle()),
6,
wpiutil.Color8Bit(wpiutil.Color.YELLOW),
)
# Subsystem constructor.
self.encoder.setDistancePerPulse(Constants.kArmEncoderDistPerPulse)
# Put Mechanism 2d to SmartDashboard
wpilib.SmartDashboard.putData("Arm Sim", self.mech2d)
self.armTower.setColor(wpiutil.Color8Bit(wpiutil.Color.BLUE))
# Set the Arm position setpoint and P constant to Preferences if the keys don't already exist
wpilib.Preferences.initDouble(
Constants.kArmPositionKey, self.armSetpointDegrees
)
wpilib.Preferences.initDouble(Constants.kArmPKey, self.armKp)
def simulationPeriodic(self) -> None:
# In this method, we update our simulation of what our arm is doing
# First, we set our "inputs" (voltages)
self.armSim.setInput(
[self.motor.getDutyCycle() * wpilib.RobotController.getBatteryVoltage()]
)
# Next, we update it. The standard loop time is 20ms.
self.armSim.update(0.020)
# Finally, we set our simulated encoder's readings and simulated battery voltage
self.encoderSim.setDistance(self.armSim.getAngle())
# SimBattery estimates loaded battery voltages
wpilib.simulation.RoboRioSim.setVInVoltage(
wpilib.simulation.BatterySim.calculate([self.armSim.getCurrentDraw()])
)
# Update the Mechanism Arm angle based on the simulated arm angle
self.armLigament.setAngle(
wpimath.units.radiansToDegrees(self.armSim.getAngle())
)
def loadPreferences(self) -> None:
# Read Preferences for Arm setpoint and kP on entering Teleop
self.armSetpointDegrees = wpilib.Preferences.getDouble(
Constants.kArmPositionKey, self.armSetpointDegrees
)
if self.armKp != wpilib.Preferences.getDouble(Constants.kArmPKey, self.armKp):
self.armKp = wpilib.Preferences.getDouble(Constants.kArmPKey, self.armKp)
self.controller.setP(self.armKp)
def reachSetpoint(self) -> None:
pidOutput = self.controller.calculate(
self.encoder.getDistance(),
wpimath.units.degreesToRadians(self.armSetpointDegrees),
)
self.motor.setVoltage(pidOutput)
def stop(self) -> None:
self.motor.setDutyCycle(0.0)