yoinked armvisualizer from mech advantage

TitaniumTitans · Feb 28, 2024 · f15de83 · f15de83
1 parent a9fc13e
commit f15de83
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 43 deletions.
diff --git a/src/main/java/frc/robot/subsystems/arm/ArmSubsystem.java b/src/main/java/frc/robot/subsystems/arm/ArmSubsystem.java
@@ -2,20 +2,15 @@
 
 import edu.wpi.first.math.MathUtil;
 import edu.wpi.first.math.geometry.Pose2d;
-import edu.wpi.first.math.geometry.Rotation2d;
-import edu.wpi.first.math.geometry.Translation2d;
-import edu.wpi.first.math.geometry.*;
-import edu.wpi.first.math.util.Units;
+import edu.wpi.first.math.geometry.Pose3d;
 import edu.wpi.first.wpilibj2.command.Command;
-
 import edu.wpi.first.wpilibj2.command.SubsystemBase;
-import lib.utils.AllianceFlipUtil;
-import lib.utils.FieldConstants;
-import lib.utils.GeomUtils;
-import frc.robot.Constants.ArmSetpoints;
 import frc.robot.Constants.ArmConstants;
+import frc.robot.Constants.ArmSetpoints;
+import lib.utils.AimbotUtils;
 import org.littletonrobotics.junction.Logger;
-import frc.robot.Constants;
+
+import java.util.function.Supplier;
 
 public class ArmSubsystem extends SubsystemBase {
 
@@ -39,14 +34,22 @@ public enum ArmState {
   private ArmState m_desiredState = ArmState.DISABLED;
   private ArmState m_currentState = ArmState.DISABLED;
 
+  private final Supplier<Pose2d> m_poseSupplier;
+
   public ArmSubsystem(ArmIO io) {
+    this(io, Pose2d::new);
+  }
+
+  public ArmSubsystem(ArmIO io, Supplier<Pose2d> supplier) {
     m_io = io;
     m_inputs = new ArmIOInputsAutoLogged();
 
     m_desiredWristPoseDegs = Double.NEGATIVE_INFINITY;
     m_desiredArmPoseDegs = Double.NEGATIVE_INFINITY;
 
     m_io.resetPosition();
+
+    m_poseSupplier = supplier;
   }
 
   //TODO: Finite state machine logic
@@ -98,6 +101,13 @@ public void handleState() {
         m_desiredArmPoseDegs = ArmSetpoints.STOW_SETPOINT.armPoseDegs();
         m_desiredWristPoseDegs = ArmSetpoints.STOW_SETPOINT.wristPoseDegs();
       }
+      case AUTO_AIM -> {
+        ArmSetpoints.ArmSetpoint aimSetpoint = AimbotUtils.aimbotCalculate(
+                new Pose3d(m_poseSupplier.get()), m_inputs.armPositionDegs);
+
+        m_desiredArmPoseDegs = aimSetpoint.armPoseDegs();
+        m_desiredWristPoseDegs = aimSetpoint.wristPoseDegs();
+      }
       case INTAKE -> {
         m_desiredArmPoseDegs = ArmSetpoints.INTAKE_SETPOINT.armPoseDegs();
         m_desiredWristPoseDegs = ArmSetpoints.INTAKE_SETPOINT.wristPoseDegs();
@@ -109,38 +119,7 @@ public void handleState() {
     }
   }
 
-  /** Gets the top point of the shooter for checking limits*/
-  public Translation2d calculateArmPosition(double armAngle, double wristAngle) {
-    return ArmConstants.PIVOT_JOINT_TRANSLATION
-        // translate the length + direction of the arm
-        .plus(new Translation2d(ArmConstants.ARM_LENGTH_METERS,
-            Rotation2d.fromDegrees(armAngle)))
-        // translate the length + direction of the wrist
-        .plus(new Translation2d(ArmConstants.WRIST_LENGTH_METERS,
-            Rotation2d.fromDegrees(360)
-                .minus(Rotation2d.fromDegrees(wristAngle))));
-  }
 
-  /** Gets the transformation of the shooter relative to the drive base */
-  public Transform3d getShooterTransformation() {
-    return Constants.ArmConstants.PIVOT_TRANSLATION_METERS.plus(
-            // Add the movement of the arm
-            GeomUtils.translationToTransform(new Translation3d(
-                    ArmConstants.ARM_LENGTH_METERS,
-                    new Rotation3d(0.0, Units.degreesToRadians(m_inputs.armPositionDegs), 0.0)
-            ))
-    );
-  }
-
-  public double aimbotCalculate(Pose3d robotPose) {
-    Pose3d speakerPose = new Pose3d(AllianceFlipUtil.apply(FieldConstants.CENTER_SPEAKER), new Rotation3d());
-    Pose3d shooterPivotPose = robotPose.plus(getShooterTransformation());
-    Transform3d robotToSpeaker = new Transform3d(shooterPivotPose.plus(getShooterTransformation()), speakerPose);
-
-    double groundDistance =
-            Math.sqrt(Math.pow(robotToSpeaker.getX(), 2) + Math.pow(robotToSpeaker.getY(), 2));
-    return Math.atan2(groundDistance, robotToSpeaker.getZ());
-  }
 
   /* Command Factories */
 

diff --git a/src/main/java/frc/robot/subsystems/arm/ArmVisualizer.java b/src/main/java/frc/robot/subsystems/arm/ArmVisualizer.java
@@ -0,0 +1,39 @@
+package frc.robot.subsystems.arm;
+
+import edu.wpi.first.math.geometry.Pose3d;
+import edu.wpi.first.math.geometry.Rotation2d;
+import edu.wpi.first.math.geometry.Rotation3d;
+import edu.wpi.first.wpilibj.smartdashboard.Mechanism2d;
+import edu.wpi.first.wpilibj.smartdashboard.MechanismLigament2d;
+import edu.wpi.first.wpilibj.smartdashboard.MechanismRoot2d;
+import edu.wpi.first.wpilibj.util.Color;
+import edu.wpi.first.wpilibj.util.Color8Bit;
+import frc.robot.Constants.ArmConstants;
+import org.littletonrobotics.junction.Logger;
+
+public class ArmVisualizer {
+  private final Mechanism2d mechanism;
+  private final MechanismLigament2d arm;
+  private final String key;
+
+  public ArmVisualizer(String key, Color color) {
+    this.key = key;
+    mechanism = new Mechanism2d(3.0, 3.0, new Color8Bit(Color.kWhite));
+    MechanismRoot2d root = mechanism.getRoot("pivot", 1.0, 0.4);
+    arm = new MechanismLigament2d("arm", ArmConstants.ARM_LENGTH_METERS, 20.0, 6, new Color8Bit(color));
+    root.append(arm);
+  }
+
+  /** Update arm visualizer with current arm angle */
+  public void update(double angleDegs) {
+    // Log Mechanism2d
+    arm.setAngle(Rotation2d.fromDegrees(angleDegs));
+    Logger.recordOutput("Arm/Mechanism2d/" + key, mechanism);
+
+    // Log 3D poses
+    Pose3d pivotArm =
+            new Pose3d(ArmConstants.PIVOT_JOINT_TRANSLATION.getX(), 0.0,
+                    ArmConstants.PIVOT_JOINT_TRANSLATION.getY(), new Rotation3d(0.0, -angleDegs, 0.0));
+    Logger.recordOutput("Arm/Mechanism3d/" + key, pivotArm);
+  }
+}
diff --git a/src/main/java/lib/utils/AimbotUtils.java b/src/main/java/lib/utils/AimbotUtils.java
@@ -0,0 +1,45 @@
+package lib.utils;
+
+import edu.wpi.first.math.geometry.*;
+import edu.wpi.first.math.util.Units;
+import frc.robot.Constants.ArmConstants;
+import frc.robot.Constants.ArmSetpoints;
+
+public class AimbotUtils {
+  /** Gets the top point of the shooter for checking limits*/
+  public static Translation2d calculateArmPosition(double armAngle, double wristAngle) {
+    return ArmConstants.PIVOT_JOINT_TRANSLATION
+            // translate the length + direction of the arm
+            .plus(new Translation2d(ArmConstants.ARM_LENGTH_METERS,
+                    Rotation2d.fromDegrees(armAngle)))
+            // translate the length + direction of the wrist
+            .plus(new Translation2d(ArmConstants.WRIST_LENGTH_METERS,
+                    Rotation2d.fromDegrees(360)
+                            .minus(Rotation2d.fromDegrees(wristAngle))));
+  }
+
+  /** Gets the transformation of the shooter relative to the drive base */
+  public static Transform3d getShooterTransformation(double armAngle) {
+    return ArmConstants.PIVOT_TRANSLATION_METERS.plus(
+            // Add the movement of the arm
+            GeomUtils.translationToTransform(new Translation3d(
+                    ArmConstants.ARM_LENGTH_METERS,
+                    new Rotation3d(0.0, Units.degreesToRadians(armAngle), 0.0)
+            ))
+    );
+  }
+
+  public static ArmSetpoints.ArmSetpoint aimbotCalculate(Pose3d robotPose, double armAngle) {
+    Pose3d speakerPose = new Pose3d(AllianceFlipUtil.apply(FieldConstants.CENTER_SPEAKER), new Rotation3d());
+    Pose3d shooterPivotPose = robotPose.plus(getShooterTransformation(armAngle));
+    Transform3d robotToSpeaker =
+            new Transform3d(shooterPivotPose.plus(getShooterTransformation(armAngle)), speakerPose);
+
+    double groundDistance =
+            Math.sqrt(Math.pow(robotToSpeaker.getX(), 2) + Math.pow(robotToSpeaker.getY(), 2));
+
+    double desiredWristAngle = Math.atan2(groundDistance, robotToSpeaker.getZ());
+    double safeArmAngle = desiredWristAngle - ArmConstants.WRIST_ARM_GAP.getValue();
+    return new ArmSetpoints.ArmSetpoint(safeArmAngle, desiredWristAngle);
+  }
+}
diff --git a/src/main/java/lib/GeomUtil.java → src/main/java/lib/utils/GeomUtil.java b/src/main/java/lib/GeomUtil.java → src/main/java/lib/utils/GeomUtil.java
@@ -5,7 +5,7 @@
 // license that can be found in the LICENSE file at
 // the root directory of this project.
 
-package lib;
+package lib.utils;
 
 import edu.wpi.first.math.geometry.Pose2d;
 import edu.wpi.first.math.geometry.Pose3d;

diff --git a/src/main/java/lib/utils/PoseEstimator.java b/src/main/java/lib/utils/PoseEstimator.java
@@ -20,7 +20,6 @@
 import java.util.List;
 import java.util.NavigableMap;
 import java.util.TreeMap;
-import lib.GeomUtil;
 
 public class PoseEstimator {
   private static final double HISTORY_LENGTH_SECS = 0.3;