[py] Add wpimath/kinematics tests (#8461)

Used Gemini to convert the C++ tests in wpimath/.../kinematics to pytest. For ease of use required a better constructor for `MecanumDriveWheelPositions` I did this conversion months ago (before robotpy landed and before the reorg), so new tests might be missing. At least its better than nothing 🤷 Broke the huge pr that does almost everything into parts so its easier to review. --------- Co-authored-by: David Vo <auscompgeek@users.noreply.github.com>
2026-07-04 03:11:43 +00:00 · 2026-01-12 22:14:14 -05:00
parent c2b339cfc9
commit eccc2301ac
15 changed files with 2061 additions and 1 deletions
--- a/wpimath/src/test/python/kinematics/test_mecanum_drive_odometry3d.py
+++ b/wpimath/src/test/python/kinematics/test_mecanum_drive_odometry3d.py
@@ -0,0 +1,264 @@
+import pytest
+import math
+import random
+
+from wpimath import (
+    ChassisSpeeds,
+    MecanumDriveKinematics,
+    MecanumDriveOdometry3d,
+    MecanumDriveWheelPositions,
+    Pose2d,
+    Pose3d,
+    Rotation2d,
+    Rotation3d,
+    TrajectoryGenerator,
+    TrajectoryConfig,
+    Translation2d,
+)
+
+
+@pytest.fixture
+def odometry3d_test():
+    class MecanumDriveOdometry3dTest:
+        def __init__(self):
+            self.m_fl = Translation2d(x=12, y=12)
+            self.m_fr = Translation2d(x=12, y=-12)
+            self.m_bl = Translation2d(x=-12, y=12)
+            self.m_br = Translation2d(x=-12, y=-12)
+
+            self.zero = MecanumDriveWheelPositions()
+
+            self.kinematics = MecanumDriveKinematics(
+                self.m_fl, self.m_fr, self.m_bl, self.m_br
+            )
+            self.odometry = MecanumDriveOdometry3d(
+                self.kinematics, Rotation3d(), self.zero
+            )
+
+    return MecanumDriveOdometry3dTest()
+
+
+def test_initialize(odometry3d_test):
+    odometry = MecanumDriveOdometry3d(
+        odometry3d_test.kinematics,
+        gyroAngle=Rotation3d(),
+        wheelPositions=odometry3d_test.zero,
+        initialPose=Pose3d(x=1, y=2, z=0, rotation=Rotation3d.fromDegrees(0, 0, 45)),
+    )
+
+    pose = odometry.getPose()
+
+    assert pose.x == pytest.approx(1, abs=1e-9)
+    assert pose.y == pytest.approx(2, abs=1e-9)
+    assert pose.z == pytest.approx(0, abs=1e-9)
+    assert pose.rotation().toRotation2d().degrees() == pytest.approx(45, abs=1e-9)
+
+
+def test_multiple_consecutive_updates(odometry3d_test):
+    wheel_deltas = MecanumDriveWheelPositions(
+        frontLeft=3.536, frontRight=3.536, rearLeft=3.536, rearRight=3.536
+    )
+
+    odometry3d_test.odometry.resetPosition(Rotation3d(), wheel_deltas, Pose3d())
+    odometry3d_test.odometry.update(Rotation3d(), wheel_deltas)
+    second_pose = odometry3d_test.odometry.update(Rotation3d(), wheel_deltas)
+
+    assert second_pose.x == pytest.approx(0.0, abs=0.01)
+    assert second_pose.y == pytest.approx(0.0, abs=0.01)
+    assert second_pose.z == pytest.approx(0.0, abs=0.01)
+    assert second_pose.rotation().toRotation2d().radians() == pytest.approx(
+        0.0, abs=0.01
+    )
+
+
+def test_two_iterations(odometry3d_test):
+    odometry3d_test.odometry.resetPosition(Rotation3d(), odometry3d_test.zero, Pose3d())
+    wheel_deltas = MecanumDriveWheelPositions(
+        frontLeft=0.3536, frontRight=0.3536, rearLeft=0.3536, rearRight=0.3536
+    )
+
+    odometry3d_test.odometry.update(Rotation3d(), MecanumDriveWheelPositions())
+    pose = odometry3d_test.odometry.update(Rotation3d(), wheel_deltas)
+
+    assert pose.x == pytest.approx(0.3536, abs=0.01)
+    assert pose.y == pytest.approx(0.0, abs=0.01)
+    assert pose.z == pytest.approx(0.0, abs=0.01)
+    assert pose.rotation().toRotation2d().radians() == pytest.approx(0.0, abs=0.01)
+
+
+def test_90_degree_turn(odometry3d_test):
+    odometry3d_test.odometry.resetPosition(Rotation3d(), odometry3d_test.zero, Pose3d())
+    wheel_deltas = MecanumDriveWheelPositions(
+        frontLeft=-13.328, frontRight=39.986, rearLeft=-13.329, rearRight=39.986
+    )
+    odometry3d_test.odometry.update(Rotation3d(), MecanumDriveWheelPositions())
+    pose = odometry3d_test.odometry.update(
+        Rotation3d.fromDegrees(0, 0, 90), wheel_deltas
+    )
+
+    assert pose.x == pytest.approx(8.4855, abs=0.01)
+    assert pose.y == pytest.approx(8.4855, abs=0.01)
+    assert pose.z == pytest.approx(0, abs=0.01)
+    assert pose.rotation().toRotation2d().degrees() == pytest.approx(90.0, abs=0.01)
+
+
+def test_gyro_angle_reset(odometry3d_test):
+    odometry3d_test.odometry.resetPosition(
+        Rotation3d.fromDegrees(0, 0, 90), odometry3d_test.zero, Pose3d()
+    )
+
+    wheel_deltas = MecanumDriveWheelPositions(
+        frontLeft=0.3536, frontRight=0.3536, rearLeft=0.3536, rearRight=0.3536
+    )
+
+    pose = odometry3d_test.odometry.update(
+        Rotation3d.fromDegrees(0, 0, 90), wheel_deltas
+    )
+
+    assert pose.x == pytest.approx(0.3536, abs=0.01)
+    assert pose.y == pytest.approx(0.0, abs=0.01)
+    assert pose.z == pytest.approx(0.0, abs=0.01)
+    assert pose.rotation().toRotation2d().radians() == pytest.approx(0.0, abs=0.01)
+
+
+def test_accuracy_facing_trajectory():
+    kinematics = MecanumDriveKinematics(
+        frontLeftWheel=Translation2d(x=1, y=1),
+        frontRightWheel=Translation2d(x=1, y=-1),
+        rearLeftWheel=Translation2d(x=-1, y=-1),
+        rearRightWheel=Translation2d(x=-1, y=1),
+    )
+
+    wheel_positions = MecanumDriveWheelPositions()
+
+    odometry = MecanumDriveOdometry3d(kinematics, Rotation3d(), wheel_positions)
+
+    trajectory = TrajectoryGenerator.generateTrajectory(
+        [
+            Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(45)),
+            Pose2d(x=3, y=0, rotation=Rotation2d.fromDegrees(-90)),
+            Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(135)),
+            Pose2d(x=-3, y=0, rotation=Rotation2d.fromDegrees(-90)),
+            Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(45)),
+        ],
+        TrajectoryConfig(maxVelocity=5.0, maxAcceleration=2.0),
+    )
+
+    random.seed(4915)
+
+    dt = 0.02
+    t = 0
+
+    max_error = -1e10
+    error_sum = 0
+
+    while t < trajectory.totalTime():
+        ground_truth_state = trajectory.sample(t)
+
+        wheel_speeds = kinematics.toWheelSpeeds(
+            ChassisSpeeds(
+                ground_truth_state.velocity,
+                0,
+                ground_truth_state.velocity * ground_truth_state.curvature,
+            )
+        )
+
+        wheel_speeds.frontLeft += random.gauss(0.0, 1.0) * 0.1
+        wheel_speeds.frontRight += random.gauss(0.0, 1.0) * 0.1
+        wheel_speeds.rearLeft += random.gauss(0.0, 1.0) * 0.1
+        wheel_speeds.rearRight += random.gauss(0.0, 1.0) * 0.1
+
+        wheel_positions.frontLeft += wheel_speeds.frontLeft * dt
+        wheel_positions.frontRight += wheel_speeds.frontRight * dt
+        wheel_positions.rearLeft += wheel_speeds.rearLeft * dt
+        wheel_positions.rearRight += wheel_speeds.rearRight * dt
+
+        xhat = odometry.update(
+            Rotation3d(
+                ground_truth_state.pose.rotation()
+                + Rotation2d(random.gauss(0.0, 1.0) * 0.05)
+            ),
+            wheel_positions,
+        )
+
+        error = ground_truth_state.pose.translation().distance(
+            xhat.translation().toTranslation2d()
+        )
+        if error > max_error:
+            max_error = error
+        error_sum += error
+
+        t += dt
+
+    assert error_sum / (trajectory.totalTime() / dt) < 0.07
+    assert max_error < 0.125
+
+
+def test_accuracy_facing_x_axis():
+    kinematics = MecanumDriveKinematics(
+        frontLeftWheel=Translation2d(x=1, y=1),
+        frontRightWheel=Translation2d(x=1, y=-1),
+        rearLeftWheel=Translation2d(x=-1, y=-1),
+        rearRightWheel=Translation2d(x=-1, y=1),
+    )
+
+    wheel_positions = MecanumDriveWheelPositions()
+
+    odometry = MecanumDriveOdometry3d(kinematics, Rotation3d(), wheel_positions)
+
+    trajectory = TrajectoryGenerator.generateTrajectory(
+        [
+            Pose2d(x=0, y=0, rotation=Rotation2d(45)),
+            Pose2d(x=3, y=0, rotation=Rotation2d(-90)),
+            Pose2d(x=0, y=0, rotation=Rotation2d(135)),
+            Pose2d(x=-3, y=0, rotation=Rotation2d(-90)),
+            Pose2d(x=0, y=0, rotation=Rotation2d(45)),
+        ],
+        TrajectoryConfig(maxVelocity=5.0, maxAcceleration=2.0),
+    )
+
+    random.seed(4915)
+
+    dt = 0.02
+    t = 0
+
+    max_error = -1e10
+    error_sum = 0
+
+    while t < trajectory.totalTime():
+        ground_truth_state = trajectory.sample(t)
+
+        wheel_speeds = kinematics.toWheelSpeeds(
+            ChassisSpeeds(
+                ground_truth_state.velocity * ground_truth_state.pose.rotation().cos(),
+                ground_truth_state.velocity * ground_truth_state.pose.rotation().sin(),
+                0,
+            )
+        )
+
+        wheel_speeds.frontLeft += random.gauss(0.0, 1.0) * 0.1
+        wheel_speeds.frontRight += random.gauss(0.0, 1.0) * 0.1
+        wheel_speeds.rearLeft += random.gauss(0.0, 1.0) * 0.1
+        wheel_speeds.rearRight += random.gauss(0.0, 1.0) * 0.1
+
+        wheel_positions.frontLeft += wheel_speeds.frontLeft * dt
+        wheel_positions.frontRight += wheel_speeds.frontRight * dt
+        wheel_positions.rearLeft += wheel_speeds.rearLeft * dt
+        wheel_positions.rearRight += wheel_speeds.rearRight * dt
+
+        xhat = odometry.update(
+            Rotation3d(0, 0, random.gauss(0.0, 1.0) * 0.05),
+            wheel_positions,
+        )
+
+        error = ground_truth_state.pose.translation().distance(
+            xhat.translation().toTranslation2d()
+        )
+        if error > max_error:
+            max_error = error
+        error_sum += error
+
+        t += dt
+
+    assert error_sum / (trajectory.totalTime() / dt) < 0.06
+    assert max_error < 0.125