Team2890/allwpilib
4
0
Fork 0
mirror of https://github.com/wpilibsuite/allwpilib synced 2026-06-21 01:01:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

[py] Add wpimath/geometry tests (#8460)

Browse Source
This commit is contained in:
PJ Reiniger
2026-01-02 12:10:02 -05:00
committed by GitHub
parent 32fa124166
commit 8518b7014c
16 changed files with 2004 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

220
wpimath/src/test/python/geometry/test_coordinate_system.py Normal file
View File

@@ -0,0 +1,220 @@
import pytest
import math
from wpimath import (
CoordinateSystem,
Pose3d,
Transform3d,
Translation3d,
Rotation3d,
)
def check_pose3d_convert(pose_from, pose_to, coord_from, coord_to):
# "from" to "to"
assert pose_to.translation() == CoordinateSystem.convert(
pose_from.translation(), coord_from, coord_to
)
assert pose_to.rotation() == CoordinateSystem.convert(
pose_from.rotation(), coord_from, coord_to
)
assert pose_to == CoordinateSystem.convert(pose_from, coord_from, coord_to)
# "to" to "from"
assert pose_from.translation() == CoordinateSystem.convert(
pose_to.translation(), coord_to, coord_from
)
assert pose_from.rotation() == CoordinateSystem.convert(
pose_to.rotation(), coord_to, coord_from
)
assert pose_from == CoordinateSystem.convert(pose_to, coord_to, coord_from)
def check_transform3d_convert(transform_from, transform_to, coord_from, coord_to):
# "from" to "to"
assert transform_to.translation() == CoordinateSystem.convert(
transform_from.translation(), coord_from, coord_to
)
assert transform_to == CoordinateSystem.convert(
transform_from, coord_from, coord_to
)
# "to" to "from"
assert transform_from.translation() == CoordinateSystem.convert(
transform_to.translation(), coord_to, coord_from
)
assert transform_from == CoordinateSystem.convert(
transform_to, coord_to, coord_from
)
def test_pose3d_edn_to_nwu():
# No rotation from EDN to NWU
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d()),
Pose3d(x=3, y=-1, z=-2, rotation=Rotation3d.fromDegrees(-90, 0, -90)),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
# 45° roll from EDN to NWU
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(45, 0, 0)),
Pose3d(x=3, y=-1, z=-2, rotation=Rotation3d.fromDegrees(-45, 0, -90)),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
# 45° pitch from EDN to NWU
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(0, 45, 0)),
Pose3d(x=3, y=-1, z=-2, rotation=Rotation3d.fromDegrees(-90, 0, -135)),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
# 45° yaw from EDN to NWU
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(0, 0, 45)),
Pose3d(x=3, y=-1, z=-2, rotation=Rotation3d.fromDegrees(-90, 45, -90)),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
def test_pose3d_edn_to_ned():
# No rotation from EDN to NED
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d()),
Pose3d(x=3, y=1, z=2, rotation=Rotation3d.fromDegrees(90, 0, 90)),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
# 45° roll from EDN to NED
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(45, 0, 0)),
Pose3d(x=3, y=1, z=2, rotation=Rotation3d.fromDegrees(135, 0, 90)),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
# 45° pitch from EDN to NED
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(0, 45, 0)),
Pose3d(x=3, y=1, z=2, rotation=Rotation3d.fromDegrees(90, 0, 135)),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
# 45° yaw from EDN to NED
check_pose3d_convert(
Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(0, 0, 45)),
Pose3d(x=3, y=1, z=2, rotation=Rotation3d.fromDegrees(90, -45, 90)),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
def test_transform3d_edn_to_nwu():
# No rotation from EDN to NWU
check_transform3d_convert(
Transform3d(Translation3d(x=1, y=2, z=3), rotation=Rotation3d()),
Transform3d(Translation3d(x=3, y=-1, z=-2), rotation=Rotation3d()),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
# 45° roll from EDN to NWU
check_transform3d_convert(
Transform3d(
Translation3d(x=1, y=2, z=3),
Rotation3d.fromDegrees(45, 0, 0),
),
Transform3d(
Translation3d(x=3, y=-1, z=-2),
Rotation3d.fromDegrees(0, -45, 0),
),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
# 45° pitch from EDN to NWU
check_transform3d_convert(
Transform3d(
Translation3d(x=1, y=2, z=3),
Rotation3d.fromDegrees(0, 45, 0),
),
Transform3d(
Translation3d(x=3, y=-1, z=-2),
Rotation3d.fromDegrees(0, 0, -45),
),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
# 45° yaw from EDN to NWU
check_transform3d_convert(
Transform3d(
Translation3d(x=1, y=2, z=3),
Rotation3d.fromDegrees(0, 0, 45),
),
Transform3d(
Translation3d(x=3, y=-1, z=-2),
Rotation3d.fromDegrees(45, 0, 0),
),
CoordinateSystem.EDN(),
CoordinateSystem.NWU(),
)
def test_transform3d_edn_to_ned():
# No rotation from EDN to NED
check_transform3d_convert(
Transform3d(Translation3d(x=1, y=2, z=3), rotation=Rotation3d()),
Transform3d(Translation3d(x=3, y=1, z=2), rotation=Rotation3d()),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
# 45° roll from EDN to NED
check_transform3d_convert(
Transform3d(
Translation3d(x=1, y=2, z=3),
Rotation3d.fromDegrees(45, 0, 0),
),
Transform3d(
Translation3d(x=3, y=1, z=2),
Rotation3d.fromDegrees(0, 45, 0),
),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
# 45° pitch from EDN to NED
check_transform3d_convert(
Transform3d(
Translation3d(x=1, y=2, z=3),
Rotation3d.fromDegrees(0, 45, 0),
),
Transform3d(
Translation3d(x=3, y=1, z=2),
Rotation3d.fromDegrees(0, 0, 45),
),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)
# 45° yaw from EDN to NED
check_transform3d_convert(
Transform3d(
Translation3d(x=1, y=2, z=3),
Rotation3d.fromDegrees(0, 0, 45),
),
Transform3d(
Translation3d(x=3, y=1, z=2),
Rotation3d.fromDegrees(45, 0, 0),
),
CoordinateSystem.EDN(),
CoordinateSystem.NED(),
)

97
wpimath/src/test/python/geometry/test_ellipse2d.py Normal file
View File

@@ -0,0 +1,97 @@
import pytest
import math
from wpimath import Pose2d, Translation2d, Ellipse2d, Rotation2d
def test_focal_points():
center = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(0))
ellipse = Ellipse2d(center, xSemiAxis=5, ySemiAxis=4)
a, b = ellipse.focalPoints()
assert a == Translation2d(x=-2, y=2)
assert b == Translation2d(x=4, y=2)
def test_intersects():
center = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(0))
ellipse = Ellipse2d(center, xSemiAxis=2, ySemiAxis=1)
point_a = Translation2d(x=1, y=3)
point_b = Translation2d(x=0, y=3)
assert ellipse.intersects(point_a)
assert not ellipse.intersects(point_b)
def test_contains():
center = Pose2d(x=-1, y=-2, rotation=Rotation2d.fromDegrees(45))
ellipse = Ellipse2d(center, xSemiAxis=2, ySemiAxis=1)
point_a = Translation2d(x=0, y=-1)
point_b = Translation2d(x=0.5, y=-2)
assert ellipse.contains(point_a)
assert not ellipse.contains(point_b)
def test_distance():
k_epsilon = 1e-9
center = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(270))
ellipse = Ellipse2d(center, xSemiAxis=1, ySemiAxis=2)
point1 = Translation2d(x=2.5, y=2)
assert ellipse.distance(point1) == pytest.approx(0, abs=k_epsilon)
point2 = Translation2d(x=1, y=2)
assert ellipse.distance(point2) == pytest.approx(0, abs=k_epsilon)
point3 = Translation2d(x=1, y=1)
assert ellipse.distance(point3) == pytest.approx(0, abs=k_epsilon)
point4 = Translation2d(x=-1, y=2.5)
assert ellipse.distance(point4) == pytest.approx(0.19210128384806818, abs=k_epsilon)
def test_nearest():
k_epsilon = 1e-9
center = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(270))
ellipse = Ellipse2d(center, xSemiAxis=1, ySemiAxis=2)
point1 = Translation2d(x=2.5, y=2)
nearest_point1 = ellipse.nearest(point1)
assert nearest_point1.x == pytest.approx(2.5, abs=k_epsilon)
assert nearest_point1.y == pytest.approx(2.0, abs=k_epsilon)
point2 = Translation2d(x=1, y=2)
nearest_point2 = ellipse.nearest(point2)
assert nearest_point2.x == pytest.approx(1.0, abs=k_epsilon)
assert nearest_point2.y == pytest.approx(2.0, abs=k_epsilon)
point3 = Translation2d(x=1, y=1)
nearest_point3 = ellipse.nearest(point3)
assert nearest_point3.x == pytest.approx(1.0, abs=k_epsilon)
assert nearest_point3.y == pytest.approx(1.0, abs=k_epsilon)
point4 = Translation2d(x=-1, y=2.5)
nearest_point4 = ellipse.nearest(point4)
assert nearest_point4.x == pytest.approx(-0.8512799937611617, abs=k_epsilon)
assert nearest_point4.y == pytest.approx(2.378405333174535, abs=k_epsilon)
def test_equals():
center1 = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(90))
ellipse1 = Ellipse2d(center1, xSemiAxis=2, ySemiAxis=3)
center2 = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(90))
ellipse2 = Ellipse2d(center2, xSemiAxis=2, ySemiAxis=3)
center3 = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(90))
ellipse3 = Ellipse2d(center3, xSemiAxis=3, ySemiAxis=2)
assert ellipse1 == ellipse2
assert ellipse1 != ellipse3
assert ellipse3 != ellipse2

149
wpimath/src/test/python/geometry/test_pose2d.py Normal file
View File

@@ -0,0 +1,149 @@
import pytest
import math
from wpimath import Pose2d, Rotation2d, Translation2d, Transform2d
from wpimath.units import feetToMeters
def test_rotate_by():
x = 1
y = 2
initial = Pose2d(x, y, math.radians(45))
rotation = Rotation2d(math.radians(5))
rotated = initial.rotateBy(rotation)
# Translation is rotated by CCW rotation matrix
c = rotation.cos()
s = rotation.sin()
assert rotated.x == pytest.approx(c * x - s * y)
assert rotated.y == pytest.approx(s * x + c * y)
assert rotated.rotation().degrees() == pytest.approx(
initial.rotation().degrees() + rotation.degrees()
)
def test_transform_by():
initial = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(45))
transform = Transform2d(Translation2d(x=5, y=0), Rotation2d.fromDegrees(5))
transformed = initial + transform
assert transformed.x == pytest.approx(1.0 + 5.0 / math.sqrt(2.0))
assert transformed.y == pytest.approx(2.0 + 5.0 / math.sqrt(2.0))
assert transformed.rotation().degrees() == pytest.approx(50.0)
def test_relative_to():
initial = Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(45))
final = Pose2d(x=5, y=5, rotation=Rotation2d.fromDegrees(45.0))
final_relative_to_initial = final.relativeTo(initial)
assert final_relative_to_initial.x == pytest.approx(5.0 * math.sqrt(2.0), abs=1e-9)
assert final_relative_to_initial.y == pytest.approx(0.0, abs=1e-9)
assert final_relative_to_initial.rotation().degrees() == pytest.approx(
0.0, abs=1e-9
)
def test_rotate_around():
initial = Pose2d(x=5, y=0, rotation=Rotation2d.fromDegrees(0))
point = Translation2d(x=0, y=0)
rotated = initial.rotateAround(point, Rotation2d.fromDegrees(180))
assert rotated.x == pytest.approx(-5.0, abs=1e-9)
assert rotated.y == pytest.approx(0.0, abs=1e-9)
assert rotated.rotation().degrees() == pytest.approx(180.0, abs=1e-9)
def test_equality():
a = Pose2d(x=0, y=5, rotation=Rotation2d.fromDegrees(43))
b = Pose2d(x=0, y=5, rotation=Rotation2d.fromDegrees(43))
assert a == b
def test_inequality():
a = Pose2d(x=0, y=5, rotation=Rotation2d.fromDegrees(43))
b = Pose2d(x=0, y=feetToMeters(5), rotation=Rotation2d.fromDegrees(43))
assert a != b
def test_minus():
initial = Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(45))
final = Pose2d(x=5, y=5, rotation=Rotation2d.fromDegrees(45))
transform = final - initial
assert transform.x == pytest.approx(5.0 * math.sqrt(2.0), abs=1e-9)
assert transform.y == pytest.approx(0.0, abs=1e-9)
assert transform.rotation().degrees() == pytest.approx(0.0, abs=1e-9)
def test_nearest():
origin = Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(0))
pose1 = Pose2d(
Translation2d(distance=1.0, angle=Rotation2d.fromDegrees((45))),
Rotation2d.fromDegrees(0),
)
pose2 = Pose2d(
Translation2d(distance=2.0, angle=Rotation2d.fromDegrees((90))),
Rotation2d.fromDegrees(0),
)
pose3 = Pose2d(
Translation2d(distance=3.0, angle=Rotation2d.fromDegrees((135))),
Rotation2d.fromDegrees(0),
)
pose4 = Pose2d(
Translation2d(distance=4.0, angle=Rotation2d.fromDegrees((180))),
Rotation2d.fromDegrees(0),
)
pose5 = Pose2d(
Translation2d(distance=5.0, angle=Rotation2d.fromDegrees((270))),
Rotation2d.fromDegrees(0),
)
assert origin.nearest([pose5, pose3, pose4]).x == pytest.approx(pose3.x)
assert origin.nearest([pose5, pose3, pose4]).y == pytest.approx(pose3.y)
assert origin.nearest([pose1, pose2, pose3]).x == pytest.approx(pose1.x)
assert origin.nearest([pose1, pose2, pose3]).y == pytest.approx(pose1.y)
assert origin.nearest([pose4, pose2, pose3]).x == pytest.approx(pose2.x)
assert origin.nearest([pose4, pose2, pose3]).y == pytest.approx(pose2.y)
# Rotation component sort (when distance is the same)
# Use the same translation because using different angles at the same
# distance can cause rounding error.
translation = Translation2d(distance=1.0, angle=Rotation2d.fromDegrees(0))
pose_a = Pose2d(translation, Rotation2d.fromDegrees(0))
pose_b = Pose2d(translation, Rotation2d.fromDegrees(30))
pose_c = Pose2d(translation, Rotation2d.fromDegrees(120))
pose_d = Pose2d(translation, Rotation2d.fromDegrees(90))
pose_e = Pose2d(translation, Rotation2d.fromDegrees(-180))
assert Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(360)).nearest(
[pose_a, pose_b, pose_d]
).rotation().degrees() == pytest.approx(pose_a.rotation().degrees())
assert Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(-335)).nearest(
[pose_b, pose_c, pose_d]
).rotation().degrees() == pytest.approx(pose_b.rotation().degrees())
assert Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(-120)).nearest(
[pose_b, pose_c, pose_d]
).rotation().degrees() == pytest.approx(pose_c.rotation().degrees())
assert Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(85)).nearest(
[pose_a, pose_c, pose_d]
).rotation().degrees() == pytest.approx(pose_d.rotation().degrees())
assert Pose2d(x=0, y=0, rotation=Rotation2d.fromDegrees(170)).nearest(
[pose_a, pose_d, pose_e]
).rotation().degrees() == pytest.approx(pose_e.rotation().degrees())
def test_to_matrix():
before = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(20))
after = Pose2d.fromMatrix(before.toMatrix())
assert before == after

372
wpimath/src/test/python/geometry/test_pose3d.py Normal file
View File

@@ -0,0 +1,372 @@
import pytest
import math
import numpy as np
from wpimath import (
Pose2d,
Pose3d,
Rotation2d,
Rotation3d,
Transform2d,
Transform3d,
Translation2d,
Translation3d,
Quaternion,
)
from wpimath.units import feetToMeters
def test_rotate_by():
x = 1
y = 2
initial = Pose3d(x, y, 0, Rotation3d.fromDegrees(0, 0, 45))
yaw = math.radians(5)
rotation = Rotation3d.fromDegrees(0, 0, math.degrees(yaw))
rotated = initial.rotateBy(rotation)
# Translation is rotated by CCW rotation matrix
c = math.cos(yaw)
s = math.sin(yaw)
assert rotated.x == pytest.approx(c * x - s * y)
assert rotated.y == pytest.approx(s * x + c * y)
assert rotated.z == pytest.approx(0.0)
assert rotated.rotation().x == pytest.approx(0.0)
assert rotated.rotation().y == pytest.approx(0.0)
assert rotated.rotation().z == pytest.approx(initial.rotation().z + rotation.z)
def test_test_transform_by_rotations():
k_epsilon = 1e-9
initial_pose = Pose3d(x=0, y=0, z=0, rotation=Rotation3d())
transform1 = Transform3d(Translation3d(), Rotation3d.fromDegrees(90, 45, 0))
transform2 = Transform3d(Translation3d(), Rotation3d.fromDegrees(-90, 0, 0))
transform3 = Transform3d(Translation3d(), Rotation3d.fromDegrees(0, -45, 0))
final_pose = (
initial_pose.transformBy(transform1)
.transformBy(transform2)
.transformBy(transform3)
)
assert final_pose.rotation().x == pytest.approx(
initial_pose.rotation().x, abs=k_epsilon
)
assert final_pose.rotation().y == pytest.approx(
initial_pose.rotation().y, abs=k_epsilon
)
assert final_pose.rotation().z == pytest.approx(
initial_pose.rotation().z, abs=k_epsilon
)
def test_transform_by():
z_axis = np.array([0.0, 0.0, 1.0])
initial = Pose3d(x=1, y=2, z=0, rotation=Rotation3d(z_axis, math.radians(45.0)))
transform = Transform3d(
Translation3d(x=5, y=0, z=0), rotation=Rotation3d(z_axis, math.radians(5))
)
transformed = initial + transform
assert transformed.x == pytest.approx(1.0 + 5.0 / math.sqrt(2.0))
assert transformed.y == pytest.approx(2.0 + 5.0 / math.sqrt(2.0))
assert transformed.rotation().z == pytest.approx(math.radians(50))
def test_relative_to():
z_axis = np.array([0.0, 0.0, 1.0])
initial = Pose3d(x=0, y=0, z=0, rotation=Rotation3d(z_axis, math.radians(45)))
final = Pose3d(x=5, y=5, z=0, rotation=Rotation3d(z_axis, math.radians(45)))
final_relative_to_initial = final.relativeTo(initial)
assert final_relative_to_initial.x == pytest.approx(5.0 * math.sqrt(2.0))
assert final_relative_to_initial.y == pytest.approx(0.0)
assert final_relative_to_initial.rotation().z == pytest.approx(0.0, abs=1e-9)
def test_rotate_around():
initial = Pose3d(x=5, y=0, z=0, rotation=Rotation3d())
point = Translation3d(x=0, y=0, z=0)
rotated = initial.rotateAround(point, Rotation3d.fromDegrees(0, 0, 180))
assert rotated.x == pytest.approx(-5.0, abs=1e-9)
assert rotated.y == pytest.approx(0.0, abs=1e-9)
assert rotated.rotation().z == pytest.approx(math.radians(180), abs=1e-9)
def test_equality():
z_axis = np.array([0.0, 0.0, 1.0])
a = Pose3d(x=0, y=5, z=0, rotation=Rotation3d(z_axis, math.radians(43)))
b = Pose3d(x=0, y=5, z=0, rotation=Rotation3d(z_axis, math.radians(43)))
assert a == b
def test_inequality():
z_axis = np.array([0.0, 0.0, 1.0])
a = Pose3d(x=0, y=5, z=0, rotation=Rotation3d(z_axis, math.radians(43)))
b = Pose3d(
x=0, y=feetToMeters(5), z=0, rotation=Rotation3d(z_axis, math.radians(43))
)
assert a != b
def test_minus():
z_axis = np.array([0.0, 0.0, 1.0])
initial = Pose3d(x=0, y=0, z=0, rotation=Rotation3d(z_axis, math.radians(45)))
final = Pose3d(x=5, y=5, z=0, rotation=Rotation3d(z_axis, math.radians(45)))
transform = final - initial
assert transform.x == pytest.approx(5.0 * math.sqrt(2.0))
assert transform.y == pytest.approx(0.0)
assert transform.rotation().z == pytest.approx(0.0, abs=1e-9)
def test_to_matrix():
before = Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(10, 20, 30))
after = Pose3d.fromMatrix(before.toMatrix())
assert before == after
def test_to_pose2d():
pose = Pose3d(x=1, y=2, z=3, rotation=Rotation3d.fromDegrees(20, 30, 40))
expected = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(40))
assert expected == pose.toPose2d()
def test_complex_twists():
initial_poses = [
Pose3d(
x=0.698303,
y=-0.959096,
z=0.271076,
rotation=Rotation3d(Quaternion(0.86403, -0.076866, 0.147234, 0.475254)),
),
Pose3d(
x=0.634892,
y=-0.765209,
z=0.117543,
rotation=Rotation3d(Quaternion(0.84987, -0.070829, 0.162097, 0.496415)),
),
Pose3d(
x=0.584827,
y=-0.590303,
z=-0.02557,
rotation=Rotation3d(Quaternion(0.832743, -0.041991, 0.202188, 0.513708)),
),
Pose3d(
x=0.505038,
y=-0.451479,
z=-0.112835,
rotation=Rotation3d(Quaternion(0.816515, -0.002673, 0.226182, 0.531166)),
),
Pose3d(
x=0.428178,
y=-0.329692,
z=-0.189707,
rotation=Rotation3d(Quaternion(0.807886, 0.029298, 0.257788, 0.529157)),
),
]
final_poses = [
Pose3d(
x=-0.230448,
y=-0.511957,
z=0.198406,
rotation=Rotation3d(Quaternion(0.753984, 0.347016, 0.409105, 0.379106)),
),
Pose3d(
x=-0.088932,
y=-0.343253,
z=0.095018,
rotation=Rotation3d(Quaternion(0.638738, 0.413016, 0.536281, 0.365833)),
),
Pose3d(
x=-0.107908,
y=-0.317552,
z=0.133946,
rotation=Rotation3d(Quaternion(0.653444, 0.417069, 0.465505, 0.427046)),
),
Pose3d(
x=-0.123383,
y=-0.156411,
z=-0.047435,
rotation=Rotation3d(Quaternion(0.652983, 0.40644, 0.431566, 0.47135)),
),
Pose3d(
x=-0.084654,
y=-0.019305,
z=-0.030022,
rotation=Rotation3d(Quaternion(0.620243, 0.429104, 0.479384, 0.44873)),
),
]
for i in range(len(initial_poses)):
start = initial_poses[i]
end = final_poses[i]
twist = (end - start).log()
start_exp = start + twist.exp()
eps = 1e-5
assert start_exp.x == pytest.approx(end.x, abs=eps)
assert start_exp.y == pytest.approx(end.y, abs=eps)
assert start_exp.z == pytest.approx(end.z, abs=eps)
assert start_exp.rotation().getQuaternion().w == pytest.approx(
end.rotation().getQuaternion().w, abs=eps
)
assert start_exp.rotation().getQuaternion().x == pytest.approx(
end.rotation().getQuaternion().x, abs=eps
)
assert start_exp.rotation().getQuaternion().y == pytest.approx(
end.rotation().getQuaternion().y, abs=eps
)
assert start_exp.rotation().getQuaternion().z == pytest.approx(
end.rotation().getQuaternion().z, abs=eps
)
def test_twist_nan():
initial_poses = [
Pose3d(
x=6.32,
y=4.12,
z=0.00,
rotation=Rotation3d(
Quaternion(-0.9999999999999999, 0.0, 0.0, 1.9208309264993548e-08)
),
),
Pose3d(
x=3.75,
y=2.95,
z=0.00,
rotation=Rotation3d(
Quaternion(0.9999999999999793, 0.0, 0.0, 2.0352360299846772e-07)
),
),
]
final_poses = [
Pose3d(
x=6.33,
y=4.15,
z=0.00,
rotation=Rotation3d(
Quaternion(-0.9999999999999999, 0.0, 0.0, 2.416890209039172e-08)
),
),
Pose3d(
x=3.66,
y=2.93,
z=0.00,
rotation=Rotation3d(
Quaternion(0.9999999999999782, 0.0, 0.0, 2.0859477994905617e-07)
),
),
]
for i in range(len(initial_poses)):
start = initial_poses[i]
end = final_poses[i]
twist = (end - start).log()
assert not math.isnan(twist.dx)
assert not math.isnan(twist.dy)
assert not math.isnan(twist.dz)
assert not math.isnan(twist.rx)
assert not math.isnan(twist.ry)
assert not math.isnan(twist.rz)
def test_nearest():
origin = Pose3d(x=0, y=0, z=0, rotation=Rotation3d())
# Distance sort
# poses are in order of closest to farthest away from the origin at
# various positions in 3D space.
pose1 = Pose3d(x=1, y=0, z=0, rotation=Rotation3d())
pose2 = Pose3d(x=0, y=2, z=0, rotation=Rotation3d())
pose3 = Pose3d(x=0, y=0, z=3, rotation=Rotation3d())
pose4 = Pose3d(x=2, y=2, z=2, rotation=Rotation3d())
pose5 = Pose3d(x=3, y=3, z=3, rotation=Rotation3d())
assert origin.nearest([pose5, pose3, pose4]).x == pytest.approx(pose3.x)
assert origin.nearest([pose5, pose3, pose4]).y == pytest.approx(pose3.y)
assert origin.nearest([pose5, pose3, pose4]).z == pytest.approx(pose3.z)
assert origin.nearest([pose1, pose2, pose3]).x == pytest.approx(pose1.x)
assert origin.nearest([pose1, pose2, pose3]).y == pytest.approx(pose1.y)
assert origin.nearest([pose1, pose2, pose3]).z == pytest.approx(pose1.z)
assert origin.nearest([pose4, pose2, pose3]).x == pytest.approx(pose2.x)
assert origin.nearest([pose4, pose2, pose3]).y == pytest.approx(pose2.y)
assert origin.nearest([pose4, pose2, pose3]).z == pytest.approx(pose2.z)
# Rotation component sort (when distance is the same)
# Use the same translation to avoid distance differences
translation = Translation3d(x=1, y=0, z=0)
pose_a = Pose3d(translation, Rotation3d()) # No rotation
pose_b = Pose3d(translation, Rotation3d.fromDegrees(30, 0, 0))
pose_c = Pose3d(translation, Rotation3d.fromDegrees(0, 45, 0))
pose_d = Pose3d(translation, Rotation3d.fromDegrees(0, 0, 90))
pose_e = Pose3d(translation, Rotation3d.fromDegrees(180, 0, 0))
result1 = Pose3d(x=0, y=0, z=0, rotation=Rotation3d()).nearest(
[pose_a, pose_b, pose_d]
)
assert result1.rotation().x == pytest.approx(pose_a.rotation().x)
assert result1.rotation().y == pytest.approx(pose_a.rotation().y)
assert result1.rotation().z == pytest.approx(pose_a.rotation().z)
result2 = Pose3d(x=0, y=0, z=0, rotation=Rotation3d.fromDegrees(25, 0, 0)).nearest(
[pose_b, pose_c, pose_d]
)
assert result2.rotation().x == pytest.approx(pose_b.rotation().x)
assert result2.rotation().y == pytest.approx(pose_b.rotation().y)
assert result2.rotation().z == pytest.approx(pose_b.rotation().z)
result3 = Pose3d(x=0, y=0, z=0, rotation=Rotation3d.fromDegrees(0, 50, 0)).nearest(
[pose_b, pose_c, pose_d]
)
assert result3.rotation().x == pytest.approx(pose_c.rotation().x)
assert result3.rotation().y == pytest.approx(pose_c.rotation().y)
assert result3.rotation().z == pytest.approx(pose_c.rotation().z)
result4 = Pose3d(x=0, y=0, z=0, rotation=Rotation3d.fromDegrees(0, 0, 85)).nearest(
[pose_a, pose_c, pose_d]
)
assert result4.rotation().x == pytest.approx(pose_d.rotation().x)
assert result4.rotation().y == pytest.approx(pose_d.rotation().y)
assert result4.rotation().z == pytest.approx(pose_d.rotation().z)
result5 = Pose3d(x=0, y=0, z=0, rotation=Rotation3d.fromDegrees(170, 0, 0)).nearest(
[pose_a, pose_d, pose_e]
)
assert result5.rotation().x == pytest.approx(pose_e.rotation().x)
assert result5.rotation().y == pytest.approx(pose_e.rotation().y)
assert result5.rotation().z == pytest.approx(pose_e.rotation().z)
# Test with complex 3D rotations (combining roll, pitch, yaw)
complex_pose1 = Pose3d(translation, Rotation3d.fromDegrees(45, 30, 60))
complex_pose2 = Pose3d(translation, Rotation3d.fromDegrees(90, 45, 90))
complex_pose3 = Pose3d(translation, Rotation3d.fromDegrees(10, 15, 20))
complex_result = Pose3d(
x=0, y=0, z=0, rotation=Rotation3d.fromDegrees(5, 10, 15)
).nearest([complex_pose1, complex_pose2, complex_pose3])
assert complex_result.rotation().x == pytest.approx(complex_pose3.rotation().x)
assert complex_result.rotation().y == pytest.approx(complex_pose3.rotation().y)
assert complex_result.rotation().z == pytest.approx(complex_pose3.rotation().z)

227
wpimath/src/test/python/geometry/test_quaternion.py Normal file
View File

@@ -0,0 +1,227 @@
import pytest
import math
from wpimath import Quaternion, Rotation3d
def test_init():
# Identity
q1 = Quaternion()
assert q1.w == pytest.approx(1.0)
assert q1.x == pytest.approx(0.0)
assert q1.y == pytest.approx(0.0)
assert q1.z == pytest.approx(0.0)
# Normalized
q2 = Quaternion(0.5, 0.5, 0.5, 0.5)
assert q2.w == pytest.approx(0.5)
assert q2.x == pytest.approx(0.5)
assert q2.y == pytest.approx(0.5)
assert q2.z == pytest.approx(0.5)
# Unnormalized
q3 = Quaternion(0.75, 0.3, 0.4, 0.5)
assert q3.w == pytest.approx(0.75)
assert q3.x == pytest.approx(0.3)
assert q3.y == pytest.approx(0.4)
assert q3.z == pytest.approx(0.5)
q3 = q3.normalize()
norm = math.sqrt(0.75 * 0.75 + 0.3 * 0.3 + 0.4 * 0.4 + 0.5 * 0.5)
assert q3.w == pytest.approx(0.75 / norm)
assert q3.x == pytest.approx(0.3 / norm)
assert q3.y == pytest.approx(0.4 / norm)
assert q3.z == pytest.approx(0.5 / norm)
assert q3.w * q3.w + q3.x * q3.x + q3.y * q3.y + q3.z * q3.z == pytest.approx(1.0)
def test_addition():
q = Quaternion(0.1, 0.2, 0.3, 0.4)
p = Quaternion(0.5, 0.6, 0.7, 0.8)
sum = q + p
assert sum.w == pytest.approx(q.w + p.w)
assert sum.x == pytest.approx(q.x + p.x)
assert sum.y == pytest.approx(q.y + p.y)
assert sum.z == pytest.approx(q.z + p.z)
def test_subtraction():
q = Quaternion(0.1, 0.2, 0.3, 0.4)
p = Quaternion(0.5, 0.6, 0.7, 0.8)
difference = q - p
assert difference.w == pytest.approx(q.w - p.w)
assert difference.x == pytest.approx(q.x - p.x)
assert difference.y == pytest.approx(q.y - p.y)
assert difference.z == pytest.approx(q.z - p.z)
def test_scalar_multiplication():
q = Quaternion(0.1, 0.2, 0.3, 0.4)
scalar = 2
product = q * scalar
assert product.w == pytest.approx(q.w * scalar)
assert product.x == pytest.approx(q.x * scalar)
assert product.y == pytest.approx(q.y * scalar)
assert product.z == pytest.approx(q.z * scalar)
def test_scalar_division():
q = Quaternion(0.1, 0.2, 0.3, 0.4)
scalar = 2
product = q / scalar
assert product.w == pytest.approx(q.w / scalar)
assert product.x == pytest.approx(q.x / scalar)
assert product.y == pytest.approx(q.y / scalar)
assert product.z == pytest.approx(q.z / scalar)
def test_multiply():
# 90° CCW rotations around each axis
c = math.cos(math.radians(90) / 2.0)
s = math.sin(math.radians(90) / 2.0)
x_rot = Quaternion(c, s, 0.0, 0.0)
y_rot = Quaternion(c, 0.0, s, 0.0)
z_rot = Quaternion(c, 0.0, 0.0, s)
# 90° CCW X rotation, 90° CCW Y rotation, and 90° CCW Z rotation should
# produce a 90° CCW Y rotation
expected = y_rot
actual = z_rot * y_rot * x_rot
assert actual.w == pytest.approx(expected.w, abs=1e-9)
assert actual.x == pytest.approx(expected.x, abs=1e-9)
assert actual.y == pytest.approx(expected.y, abs=1e-9)
assert actual.z == pytest.approx(expected.z, abs=1e-9)
# Identity
q = Quaternion(
0.72760687510899891,
0.29104275004359953,
0.38805700005813276,
0.48507125007266594,
)
actual = q * q.inverse()
assert actual.w == pytest.approx(1.0, abs=1e-9)
assert actual.x == pytest.approx(0.0, abs=1e-9)
assert actual.y == pytest.approx(0.0, abs=1e-9)
assert actual.z == pytest.approx(0.0, abs=1e-9)
def test_conjugate():
q = Quaternion(
0.72760687510899891,
0.29104275004359953,
0.38805700005813276,
0.48507125007266594,
)
conj = q.conjugate()
assert conj.w == pytest.approx(q.w)
assert conj.x == pytest.approx(-q.x)
assert conj.y == pytest.approx(-q.y)
assert conj.z == pytest.approx(-q.z)
def test_inverse():
q = Quaternion(
0.72760687510899891,
0.29104275004359953,
0.38805700005813276,
0.48507125007266594,
)
norm = q.norm()
inv = q.inverse()
assert inv.w == pytest.approx(q.w / (norm * norm))
assert inv.x == pytest.approx(-q.x / (norm * norm))
assert inv.y == pytest.approx(-q.y / (norm * norm))
assert inv.z == pytest.approx(-q.z / (norm * norm))
def test_norm():
q = Quaternion(3, 4, 12, 84)
norm = q.norm()
assert norm == pytest.approx(85, abs=1e-9)
def test_exponential():
q = Quaternion(1.1, 2.2, 3.3, 4.4)
expect = Quaternion(
2.81211398529184, -0.392521193481878, -0.588781790222817, -0.785042386963756
)
q_exp = q.exp()
assert q_exp == expect
def test_logarithm():
q = Quaternion(1.1, 2.2, 3.3, 4.4)
expect = Quaternion(
1.7959088706354, 0.515190292664085, 0.772785438996128, 1.03038058532817
)
q_log = q.log()
assert q_log == expect
zero = Quaternion(0, 0, 0, 0)
one = Quaternion(1, 0, 0, 0)
i = Quaternion(0, 1, 0, 0)
j = Quaternion(0, 0, 1, 0)
k = Quaternion(0, 0, 0, 1)
ln_half = Quaternion(math.log(0.5), -math.pi, 0, 0)
assert one.log() == zero
assert i.log() == Quaternion(0, math.pi / 2, 0, 0)
assert j.log() == Quaternion(0, 0, math.pi / 2, 0)
assert k.log() == Quaternion(0, 0, 0, math.pi / 2)
assert (one * -1).log() == Quaternion(0, -math.pi, 0, 0)
assert (one * -0.5).log() == ln_half
def test_logarithm_and_exponential_inverse():
q = Quaternion(1.1, 2.2, 3.3, 4.4)
# These operations are order-dependent: ln(exp(q)) is congruent but not
# necessarily equal to exp(ln(q)) due to the multi-valued nature of the
# complex logarithm.
q_log_exp = q.log().exp()
assert q_log_exp == q
start = Quaternion(1, 2, 3, 4)
expect = Quaternion(5, 6, 7, 8)
twist = start.log(expect)
actual = start.exp(twist)
assert actual == expect
def test_dot_product():
q = Quaternion(1.1, 2.2, 3.3, 4.4)
p = Quaternion(5.5, 6.6, 7.7, 8.8)
assert q.dot(p) == pytest.approx(
q.w * p.w + q.x * p.x + q.y * p.y + q.z * p.z, abs=1e-9
)
def test_dot_product_as_equality():
q = Quaternion(1.1, 2.2, 3.3, 4.4)
q_conj = q.conjugate()
assert q.dot(q) == pytest.approx(q.norm() * q.norm(), abs=1e-9)
assert abs(q.dot(q_conj) - q.norm() * q_conj.norm()) > 1e-9

85
wpimath/src/test/python/geometry/test_rectangle2d.py Normal file
View File

@@ -0,0 +1,85 @@
import pytest
import math
from wpimath import Pose2d, Translation2d, Rotation2d, Rectangle2d
def test_new_with_corners():
corner_a = Translation2d(x=1, y=2)
corner_b = Translation2d(x=4, y=6)
rect = Rectangle2d(corner_a, corner_b)
assert rect.xwidth == pytest.approx(3.0)
assert rect.ywidth == pytest.approx(4.0)
assert rect.center().x == pytest.approx(2.5)
assert rect.center().y == pytest.approx(4.0)
def test_intersects():
center = Pose2d(x=4, y=3, rotation=Rotation2d.fromDegrees(90))
rect = Rectangle2d(center, xWidth=2, yWidth=3.0)
assert rect.intersects(Translation2d(x=5.5, y=4))
assert rect.intersects(Translation2d(x=3, y=2))
assert not rect.intersects(Translation2d(x=4, y=1.5))
assert not rect.intersects(Translation2d(x=4, y=3.5))
def test_contains():
center = Pose2d(x=2.0, y=3.0, rotation=Rotation2d.fromDegrees(45))
rect = Rectangle2d(center, xWidth=3, yWidth=1)
assert rect.contains(Translation2d(x=2, y=3))
assert rect.contains(Translation2d(x=3, y=4))
assert not rect.contains(Translation2d(x=3, y=3))
def test_distance():
k_epsilon = 1e-9
center = Pose2d(x=1.0, y=2.0, rotation=Rotation2d.fromDegrees(270))
rect = Rectangle2d(center, xWidth=1, yWidth=2.0)
point1 = Translation2d(x=2.5, y=2)
assert rect.distance(point1) == pytest.approx(0.5, abs=k_epsilon)
point2 = Translation2d(x=1, y=2)
assert rect.distance(point2) == pytest.approx(0, abs=k_epsilon)
point3 = Translation2d(x=1, y=1)
assert rect.distance(point3) == pytest.approx(0.5, abs=k_epsilon)
point4 = Translation2d(x=-1, y=2.5)
assert rect.distance(point4) == pytest.approx(1, abs=k_epsilon)
def test_nearest():
k_epsilon = 1e-9
center = Pose2d(x=1, y=1, rotation=Rotation2d.fromDegrees(90))
rect = Rectangle2d(center, xWidth=3, yWidth=4)
point1 = Translation2d(x=1, y=3)
nearest_point1 = rect.nearest(point1)
assert nearest_point1.x == pytest.approx(1.0, abs=k_epsilon)
assert nearest_point1.y == pytest.approx(2.5, abs=k_epsilon)
point2 = Translation2d(x=0, y=0)
nearest_point2 = rect.nearest(point2)
assert nearest_point2.x == pytest.approx(0.0, abs=k_epsilon)
assert nearest_point2.y == pytest.approx(0.0, abs=k_epsilon)
def test_equals():
center1 = Pose2d(x=2, y=3, rotation=Rotation2d.fromDegrees(0))
rect1 = Rectangle2d(center1, xWidth=5.0, yWidth=3.0)
center2 = Pose2d(x=2, y=3, rotation=Rotation2d.fromDegrees(0))
rect2 = Rectangle2d(center2, xWidth=5.0, yWidth=3.0)
center3 = Pose2d(x=2, y=3, rotation=Rotation2d.fromDegrees(0))
rect3 = Rectangle2d(center3, xWidth=3.0, yWidth=3.0)
assert rect1 == rect2
assert rect2 != rect3

7
wpimath/src/test/python/geometry/test_rotation2d.py
View File

@@ -57,6 +57,13 @@ def test_unary_minus() -> None:
assert math.isclose(-20.0, (-rot).degrees())
def test_multiply():
rot = Rotation2d.fromDegrees(10)
assert (rot * 3.0).degrees() == pytest.approx(30.0)
assert (rot * 41.0).degrees() == pytest.approx(50.0)
def test_equality() -> None:
rot1 = Rotation2d.fromDegrees(43.0)
rot2 = Rotation2d.fromDegrees(43.0)

308
wpimath/src/test/python/geometry/test_rotation3d.py Normal file
View File

@@ -0,0 +1,308 @@
import pytest
import math
import numpy as np
from wpimath import Rotation3d, Rotation2d, Quaternion
def test_gimbal_lock_accuracy():
rot1 = Rotation3d(roll=0, pitch=0, yaw=math.pi / 2)
rot2 = Rotation3d(roll=math.pi, pitch=0, yaw=0)
rot3 = Rotation3d(roll=-math.pi / 2, pitch=0, yaw=0)
result1 = rot1 + rot2 + rot3
expected1 = Rotation3d(roll=0, pitch=-math.pi / 2, yaw=math.pi / 2)
assert expected1 == result1
assert (result1.x + result1.z) == pytest.approx(math.pi / 2)
assert result1.y == pytest.approx(-math.pi / 2, abs=1e-7)
rot1 = Rotation3d(roll=0, pitch=0, yaw=math.pi / 2)
rot2 = Rotation3d(roll=-math.pi, pitch=0, yaw=0)
rot3 = Rotation3d(roll=math.pi / 2, pitch=0, yaw=0)
result2 = rot1 + rot2 + rot3
expected2 = Rotation3d(roll=0, pitch=math.pi / 2, yaw=math.pi / 2)
assert expected2 == result2
assert (result2.z - result2.x) == pytest.approx(math.pi / 2)
assert result2.y == pytest.approx(math.pi / 2, abs=1e-7)
rot1 = Rotation3d(roll=0, pitch=0, yaw=math.pi / 2)
rot2 = Rotation3d(roll=0, pitch=math.pi / 3, yaw=0)
rot3 = Rotation3d(roll=-math.pi / 2, pitch=0, yaw=0)
result3 = rot1 + rot2 + rot3
expected3 = Rotation3d(roll=0, pitch=math.pi / 2, yaw=math.pi / 6)
assert expected3 == result3
assert (result3.z - result3.x) == pytest.approx(math.pi / 6)
assert result3.y == pytest.approx(math.pi / 2)
def test_init_axis_angle_and_roll_pitch_yaw():
x_axis = np.array([1.0, 0.0, 0.0])
rot1 = Rotation3d(x_axis, math.pi / 3)
rot2 = Rotation3d(roll=math.pi / 3, pitch=0, yaw=0)
rvec1 = Rotation3d(x_axis * (math.pi / 3))
assert rot1 == rot2
assert rot1 == rvec1
y_axis = np.array([0.0, 1.0, 0.0])
rot3 = Rotation3d(y_axis, math.pi / 3)
rot4 = Rotation3d(roll=0, pitch=math.pi / 3, yaw=0)
rvec2 = Rotation3d(y_axis * (math.pi / 3))
assert rot3 == rot4
assert rot3 == rvec2
z_axis = np.array([0.0, 0.0, 1.0])
rot5 = Rotation3d(z_axis, math.pi / 3)
rot6 = Rotation3d(roll=0, pitch=0, yaw=math.pi / 3)
rvec3 = Rotation3d(z_axis * (math.pi / 3))
assert rot5 == rot6
assert rot5 == rvec3
def test_init_rotation_matrix():
# No rotation
R1 = np.identity(3)
rot1 = Rotation3d(R1)
assert Rotation3d() == rot1
# 90 degree CCW rotation around z-axis
R2 = np.zeros((3, 3))
R2[:, 0] = [0.0, 1.0, 0.0]
R2[:, 1] = [-1.0, 0.0, 0.0]
R2[:, 2] = [0.0, 0.0, 1.0]
rot2 = Rotation3d(R2)
expected2 = Rotation3d(roll=0, pitch=0, yaw=math.radians(90))
assert expected2 == rot2
# Matrix that isn't orthogonal
R3 = np.array([
[1.0, 0.0, 0.0],
[1.0, 0.0, 0.0],
[1.0, 0.0, 0.0]
])
with pytest.raises(ValueError):
Rotation3d(R3)
# Matrix that's orthogonal but not special orthogonal
R4 = np.identity(3) * 2.0
with pytest.raises(ValueError):
Rotation3d(R4)
def test_init_two_vector():
x_axis = np.array([1.0, 0.0, 0.0])
y_axis = np.array([0.0, 1.0, 0.0])
z_axis = np.array([0.0, 0.0, 1.0])
# 90 degree CW rotation around y-axis
rot1 = Rotation3d(x_axis, z_axis)
expected1 = Rotation3d(y_axis, -math.pi / 2.0)
assert expected1 == rot1
# 45 degree CCW rotation around z-axis
rot2 = Rotation3d(x_axis, np.array([1.0, 1.0, 0.0]))
expected2 = Rotation3d(z_axis, math.pi / 4.0)
assert expected2 == rot2
# 0 degree rotation of x-axes
rot3 = Rotation3d(x_axis, x_axis)
assert Rotation3d() == rot3
# 0 degree rotation of y-axes
rot4 = Rotation3d(y_axis, y_axis)
assert Rotation3d() == rot4
# 0 degree rotation of z-axes
rot5 = Rotation3d(z_axis, z_axis)
assert Rotation3d() == rot5
# 180 degree rotation tests. For 180 degree rotations, any quaternion with an
# orthogonal rotation axis is acceptable. The rotation axis and initial
# vector are orthogonal if their dot product is zero.
# 180 degree rotation of x-axes
rot6 = Rotation3d(x_axis, -x_axis)
q6 = rot6.getQuaternion()
assert q6.w == pytest.approx(0.0)
assert q6.x * x_axis[0] + q6.y * x_axis[1] + q6.z * x_axis[2] == pytest.approx(0.0)
# 180 degree rotation of y-axes
rot7 = Rotation3d(y_axis, -y_axis)
q7 = rot7.getQuaternion()
assert q7.w == pytest.approx(0.0)
assert q7.x * y_axis[0] + q7.y * y_axis[1] + q7.z * y_axis[2] == pytest.approx(0.0)
# 180 degree rotation of z-axes
rot8 = Rotation3d(z_axis, -z_axis)
q8 = rot8.getQuaternion()
assert q8.w == pytest.approx(0.0)
assert q8.x * z_axis[0] + q8.y * z_axis[1] + q8.z * z_axis[2] == pytest.approx(0.0)
def test_radians_to_degrees():
z_axis = np.array([0.0, 0.0, 1.0])
rot1 = Rotation3d(z_axis, math.pi / 3)
assert rot1.x == pytest.approx(0.0)
assert rot1.y == pytest.approx(0.0)
assert rot1.z == pytest.approx(math.radians(60))
rot2 = Rotation3d(z_axis, math.pi / 4)
assert rot2.x == pytest.approx(0.0)
assert rot2.y == pytest.approx(0.0)
assert rot2.z == pytest.approx(math.radians(45))
def test_degrees_to_radians():
z_axis = np.array([0.0, 0.0, 1.0])
rot1 = Rotation3d(z_axis, math.radians(45))
assert rot1.x == pytest.approx(0.0)
assert rot1.y == pytest.approx(0.0)
assert rot1.z == pytest.approx(math.pi / 4.0)
rot2 = Rotation3d(z_axis, math.radians(30))
assert rot2.x == pytest.approx(0.0)
assert rot2.y == pytest.approx(0.0)
assert rot2.z == pytest.approx(math.pi / 6.0)
def test_rotation_loop():
rot = Rotation3d()
rot = rot + Rotation3d(math.radians(90), 0, 0)
expected = Rotation3d(math.radians(90), 0, 0)
assert expected == rot
rot = rot + Rotation3d(0, math.radians(90), 0)
expected = Rotation3d(
np.array([1.0 / math.sqrt(3), 1.0 / math.sqrt(3), -1.0 / math.sqrt(3)]),
math.radians(120),
)
assert expected == rot
rot = rot + Rotation3d(0, 0, math.radians(90))
expected = Rotation3d(0, math.radians(90), 0)
assert expected == rot
rot = rot + Rotation3d(0, math.radians(-90), 0)
assert Rotation3d() == rot
def test_rotate_by_from_zero_x():
x_axis = np.array([1.0, 0.0, 0.0])
zero = Rotation3d()
rotated = zero + Rotation3d(x_axis, math.radians(90))
expected = Rotation3d(x_axis, math.radians(90))
assert expected == rotated
def test_rotate_by_from_zero_y():
y_axis = np.array([0.0, 1.0, 0.0])
zero = Rotation3d()
rotated = zero + Rotation3d(y_axis, math.radians(90))
expected = Rotation3d(y_axis, math.radians(90))
assert expected == rotated
def test_rotate_by_from_zero_z():
z_axis = np.array([0.0, 0.0, 1.0])
zero = Rotation3d()
rotated = zero + Rotation3d(z_axis, math.radians(90))
expected = Rotation3d(z_axis, math.radians(90))
assert expected == rotated
def test_rotate_by_non_zero_x():
x_axis = np.array([1.0, 0.0, 0.0])
rot = Rotation3d(x_axis, math.radians(90))
rot = rot + Rotation3d(x_axis, math.radians(30))
expected = Rotation3d(x_axis, math.radians(120))
assert expected == rot
def test_rotate_by_non_zero_y():
y_axis = np.array([0.0, 1.0, 0.0])
rot = Rotation3d(y_axis, math.radians(90))
rot = rot + Rotation3d(y_axis, math.radians(30))
expected = Rotation3d(y_axis, math.radians(120))
assert expected == rot
def test_rotate_by_non_zero_z():
z_axis = np.array([0.0, 0.0, 1.0])
rot = Rotation3d(z_axis, math.radians(90))
rot = rot + Rotation3d(z_axis, math.radians(30))
expected = Rotation3d(z_axis, math.radians(120))
assert expected == rot
def test_minus():
z_axis = np.array([0.0, 0.0, 1.0])
rot1 = Rotation3d(z_axis, math.radians(70))
rot2 = Rotation3d(z_axis, math.radians(30))
assert math.degrees((rot1 - rot2).z) == pytest.approx(40.0)
def test_axis_angle():
x_axis = np.array([1.0, 0.0, 0.0])
y_axis = np.array([0.0, 1.0, 0.0])
z_axis = np.array([0.0, 0.0, 1.0])
rot1 = Rotation3d(x_axis, math.radians(90))
assert np.allclose(x_axis, rot1.axis())
assert rot1.angle == pytest.approx(math.pi / 2.0)
rot2 = Rotation3d(y_axis, math.radians(45))
assert np.allclose(y_axis, rot2.axis())
assert rot2.angle == pytest.approx(math.pi / 4.0)
rot3 = Rotation3d(z_axis, math.radians(60))
assert np.allclose(z_axis, rot3.axis())
assert rot3.angle == pytest.approx(math.pi / 3.0)
def test_to_rotation2d():
rotation = Rotation3d(math.radians(20), math.radians(30), math.radians(40))
expected = Rotation2d(math.radians(40))
assert expected == rotation.toRotation2d()
def test_equality():
z_axis = np.array([0.0, 0.0, 1.0])
rot1 = Rotation3d(z_axis, math.radians(43))
rot2 = Rotation3d(z_axis, math.radians(43))
assert rot1 == rot2
rot3 = Rotation3d(z_axis, math.radians(-180))
rot4 = Rotation3d(z_axis, math.radians(180))
assert rot3 == rot4
def test_inequality():
z_axis = np.array([0.0, 0.0, 1.0])
rot1 = Rotation3d(z_axis, math.radians(43))
rot2 = Rotation3d(z_axis, math.radians(43.5))
assert rot1 != rot2
def test_to_matrix():
before = Rotation3d(math.radians(10), math.radians(20), math.radians(30))
after = Rotation3d(before.toMatrix())
assert before == after

34
wpimath/src/test/python/geometry/test_transform2d.py Normal file
View File

@@ -0,0 +1,34 @@
import pytest
import math
from wpimath import Rotation2d, Transform2d, Translation2d, Pose2d
def test_to_matrix():
before = Transform2d(x=1, y=2, rotation=Rotation2d.fromDegrees(20))
after = Transform2d.fromMatrix(before.toMatrix())
assert before == after
def test_inverse():
initial = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(45))
transform = Transform2d(Translation2d(x=5, y=0), Rotation2d.fromDegrees(5))
transformed = initial + transform
untransformed = transformed + transform.inverse()
assert initial == untransformed
def test_composition():
initial = Pose2d(x=1, y=2, rotation=Rotation2d.fromDegrees(45))
transform1 = Transform2d(
Translation2d(x=5, y=0), rotation=Rotation2d.fromDegrees(5)
)
transform2 = Transform2d(
Translation2d(x=0, y=2), rotation=Rotation2d.fromDegrees(5)
)
transformed_separate = initial + transform1 + transform2
transformed_combined = initial + (transform1 + transform2)
assert transformed_separate == transformed_combined

46
wpimath/src/test/python/geometry/test_transform3d.py Normal file
View File

@@ -0,0 +1,46 @@
import pytest
import math
import numpy as np
from wpimath import Pose3d, Rotation3d, Transform3d
def test_to_matrix():
before = Transform3d(
x=1,
y=2,
z=3,
rotation=Rotation3d(math.radians(10), math.radians(20), math.radians(30)),
)
after = Transform3d.fromMatrix(before.toMatrix())
assert before == after
def test_inverse():
z_axis = np.array([0.0, 0.0, 1.0])
initial = Pose3d(x=1, y=2, z=3, rotation=Rotation3d(z_axis, math.radians(45)))
transform = Transform3d(x=5, y=4, z=3, rotation=Rotation3d(z_axis, math.radians(5)))
transformed = initial + transform
untransformed = transformed + transform.inverse()
assert initial == untransformed
def test_composition():
z_axis = np.array([0.0, 0.0, 1.0])
initial = Pose3d(x=1, y=2, z=3, rotation=Rotation3d(z_axis, math.radians(45)))
transform1 = Transform3d(
x=5, y=0, z=0, rotation=Rotation3d(z_axis, math.radians(5))
)
transform2 = Transform3d(
x=0, y=2, z=0, rotation=Rotation3d(z_axis, math.radians(5))
)
transformed_separate = initial + transform1 + transform2
transformed_combined = initial + (transform1 + transform2)
assert transformed_separate == transformed_combined

135
wpimath/src/test/python/geometry/test_translation2d.py Normal file
View File

@@ -0,0 +1,135 @@
import pytest
import math
from wpimath import Translation2d, Rotation2d
def test_sum():
one = Translation2d(1.0, 3.0)
two = Translation2d(2.0, 5.0)
sum_ = one + two
assert sum_.x == pytest.approx(3.0)
assert sum_.y == pytest.approx(8.0)
def test_difference():
one = Translation2d(1.0, 3.0)
two = Translation2d(2.0, 5.0)
difference = one - two
assert difference.x == pytest.approx(-1.0)
assert difference.y == pytest.approx(-2.0)
def test_rotate_by():
another = Translation2d(3.0, 0.0)
rotated = another.rotateBy(Rotation2d.fromDegrees(90))
assert rotated.x == pytest.approx(0.0, abs=1e-9)
assert rotated.y == pytest.approx(3.0, abs=1e-9)
def test_rotate_around():
translation = Translation2d(2.0, 1.0)
other = Translation2d(3.0, 2.0)
rotated = translation.rotateAround(other, Rotation2d.fromDegrees(180))
assert rotated.x == pytest.approx(4.0, abs=1e-9)
assert rotated.y == pytest.approx(3.0, abs=1e-9)
def test_multiplication():
original = Translation2d(3.0, 5.0)
mult = original * 3
assert mult.x == pytest.approx(9.0)
assert mult.y == pytest.approx(15.0)
def test_division():
original = Translation2d(3.0, 5.0)
div = original / 2
assert div.x == pytest.approx(1.5)
assert div.y == pytest.approx(2.5)
def test_norm():
one = Translation2d(3.0, 5.0)
assert one.norm() == pytest.approx(math.hypot(3.0, 5.0))
def test_squared_norm():
one = Translation2d(3.0, 5.0)
assert one.squaredNorm() == pytest.approx(34.0)
def test_distance():
one = Translation2d(1.0, 1.0)
two = Translation2d(6.0, 6.0)
assert one.distance(two) == pytest.approx(5.0 * math.sqrt(2.0))
def test_squared_distance():
one = Translation2d(1.0, 1.0)
two = Translation2d(6.0, 6.0)
assert one.distance(two) ** 2 == pytest.approx(50.0)
def test_unary_minus():
original = Translation2d(-4.5, 7.0)
inverted = -original
assert inverted.x == pytest.approx(4.5)
assert inverted.y == pytest.approx(-7.0)
def test_equality():
one = Translation2d(9.0, 5.5)
two = Translation2d(9.0, 5.5)
assert one == two
def test_inequality():
one = Translation2d(9.0, 5.5)
two = Translation2d(9.0, 5.7)
assert one != two
def test_polar_constructor():
one = Translation2d(math.sqrt(2), Rotation2d.fromDegrees(45))
assert one.x == pytest.approx(1.0)
assert one.y == pytest.approx(1.0)
two = Translation2d(2.0, Rotation2d.fromDegrees(60))
assert two.x == pytest.approx(1.0)
assert two.y == pytest.approx(math.sqrt(3.0))
def test_nearest():
origin = Translation2d(0.0, 0.0)
translation1 = Translation2d(1.0, Rotation2d.fromDegrees(45))
translation2 = Translation2d(2.0, Rotation2d.fromDegrees(90))
translation3 = Translation2d(3.0, Rotation2d.fromDegrees(135))
translation4 = Translation2d(4.0, Rotation2d.fromDegrees(180))
translation5 = Translation2d(5.0, Rotation2d.fromDegrees(270))
assert origin.nearest([translation5, translation3, translation4]) == translation3
assert origin.nearest([translation1, translation2, translation3]) == translation1
assert origin.nearest([translation4, translation2, translation3]) == translation2
def test_dot():
one = Translation2d(2.0, 3.0)
two = Translation2d(3.0, 4.0)
assert one.dot(two) == pytest.approx(18.0)
def test_cross():
one = Translation2d(2.0, 3.0)
two = Translation2d(3.0, 4.0)
assert one.cross(two) == pytest.approx(-1.0)

205
wpimath/src/test/python/geometry/test_translation3d.py Normal file
View File

@@ -0,0 +1,205 @@
import pytest
import math
import numpy as np
from wpimath import (
Translation3d,
Rotation3d,
Translation2d,
)
def test_sum():
one = Translation3d(1, 3, 5)
two = Translation3d(2, 5, 8)
sum_ = one + two
assert sum_.x == pytest.approx(3.0, abs=1e-9)
assert sum_.y == pytest.approx(8.0, abs=1e-9)
assert sum_.z == pytest.approx(13.0, abs=1e-9)
def test_difference():
one = Translation3d(1, 3, 5)
two = Translation3d(2, 5, 8)
difference = one - two
assert difference.x == pytest.approx(-1.0, abs=1e-9)
assert difference.y == pytest.approx(-2.0, abs=1e-9)
assert difference.z == pytest.approx(-3.0, abs=1e-9)
def test_rotate_by():
x_axis = np.array([1.0, 0.0, 0.0])
y_axis = np.array([0.0, 1.0, 0.0])
z_axis = np.array([0.0, 0.0, 1.0])
translation = Translation3d(1, 2, 3)
rotated1 = translation.rotateBy(Rotation3d(x_axis, math.radians(90)))
assert rotated1.x == pytest.approx(1.0, abs=1e-9)
assert rotated1.y == pytest.approx(-3.0, abs=1e-9)
assert rotated1.z == pytest.approx(2.0, abs=1e-9)
rotated2 = translation.rotateBy(Rotation3d(y_axis, math.radians(90)))
assert rotated2.x == pytest.approx(3.0, abs=1e-9)
assert rotated2.y == pytest.approx(2.0, abs=1e-9)
assert rotated2.z == pytest.approx(-1.0, abs=1e-9)
rotated3 = translation.rotateBy(Rotation3d(z_axis, math.radians(90)))
assert rotated3.x == pytest.approx(-2.0, abs=1e-9)
assert rotated3.y == pytest.approx(1.0, abs=1e-9)
assert rotated3.z == pytest.approx(3.0, abs=1e-9)
def test_rotate_around():
x_axis = np.array([1.0, 0.0, 0.0])
y_axis = np.array([0.0, 1.0, 0.0])
z_axis = np.array([0.0, 0.0, 1.0])
translation = Translation3d(1, 2, 3)
around = Translation3d(3, 2, 1)
rotated1 = translation.rotateAround(around, Rotation3d(x_axis, math.radians(90)))
assert rotated1.x == pytest.approx(1.0, abs=1e-9)
assert rotated1.y == pytest.approx(0.0, abs=1e-9)
assert rotated1.z == pytest.approx(1.0, abs=1e-9)
rotated2 = translation.rotateAround(around, Rotation3d(y_axis, math.radians(90)))
assert rotated2.x == pytest.approx(5.0, abs=1e-9)
assert rotated2.y == pytest.approx(2.0, abs=1e-9)
assert rotated2.z == pytest.approx(3.0, abs=1e-9)
rotated3 = translation.rotateAround(around, Rotation3d(z_axis, math.radians(90)))
assert rotated3.x == pytest.approx(3.0, abs=1e-9)
assert rotated3.y == pytest.approx(0.0, abs=1e-9)
assert rotated3.z == pytest.approx(3.0, abs=1e-9)
def test_to_translation2d():
translation = Translation3d(1, 2, 3)
expected = Translation2d(1, 2)
assert expected == translation.toTranslation2d()
def test_multiplication():
original = Translation3d(3, 5, 7)
mult = original * 3
assert mult.x == pytest.approx(9.0, abs=1e-9)
assert mult.y == pytest.approx(15.0, abs=1e-9)
assert mult.z == pytest.approx(21.0, abs=1e-9)
def test_division():
original = Translation3d(3, 5, 7)
div = original / 2
assert div.x == pytest.approx(1.5, abs=1e-9)
assert div.y == pytest.approx(2.5, abs=1e-9)
assert div.z == pytest.approx(3.5, abs=1e-9)
def test_norm():
one = Translation3d(3, 5, 7)
assert one.norm() == pytest.approx(math.hypot(3, 5, 7), abs=1e-9)
def test_squared_norm():
one = Translation3d(3, 5, 7)
assert one.squaredNorm() == pytest.approx(83.0, abs=1e-9)
def test_distance():
one = Translation3d(1, 1, 1)
two = Translation3d(6, 6, 6)
assert one.distance(two) == pytest.approx(5 * math.sqrt(3), abs=1e-9)
def test_squared_distance():
one = Translation3d(1, 1, 1)
two = Translation3d(6, 6, 6)
assert one.squaredDistance(two) == pytest.approx(75.0, abs=1e-9)
def test_unary_minus():
original = Translation3d(-4.5, 7, 9)
inverted = -original
assert inverted.x == pytest.approx(4.5, abs=1e-9)
assert inverted.y == pytest.approx(-7.0, abs=1e-9)
assert inverted.z == pytest.approx(-9.0, abs=1e-9)
def test_equality():
one = Translation3d(9, 5.5, 3.5)
two = Translation3d(9, 5.5, 3.5)
assert one == two
def test_inequality():
one = Translation3d(9, 5.5, 3.5)
two = Translation3d(9, 5.7, 3.5)
assert one != two
def test_polar_constructor():
z_axis = np.array([0.0, 0.0, 1.0])
one = Translation3d(math.sqrt(2) * 1, Rotation3d(z_axis, math.radians(45)))
assert one.x == pytest.approx(1.0, abs=1e-9)
assert one.y == pytest.approx(1.0, abs=1e-9)
assert one.z == pytest.approx(0.0, abs=1e-9)
two = Translation3d(2, Rotation3d(z_axis, math.radians(60)))
assert two.x == pytest.approx(1.0, abs=1e-9)
assert two.y == pytest.approx(math.sqrt(3.0), abs=1e-9)
assert two.z == pytest.approx(0.0, abs=1e-9)
def test_to_vector():
vec = np.array([1.0, 2.0, 3.0])
translation = Translation3d(vec)
assert vec[0] == pytest.approx(translation.x)
assert vec[1] == pytest.approx(translation.y)
assert vec[2] == pytest.approx(translation.z)
assert np.allclose(vec, translation.toVector())
def test_nearest():
origin = Translation3d(0, 0, 0)
# Distance sort
# translations are in order of closest to farthest away from the origin at
# various positions in 3D space.
translation1 = Translation3d(1, 0, 0)
translation2 = Translation3d(0, 2, 0)
translation3 = Translation3d(0, 0, 3)
translation4 = Translation3d(2, 2, 2)
translation5 = Translation3d(3, 3, 3)
nearest1 = origin.nearest([translation5, translation3, translation4])
assert nearest1.x == pytest.approx(translation3.x)
assert nearest1.y == pytest.approx(translation3.y)
assert nearest1.z == pytest.approx(translation3.z)
nearest2 = origin.nearest([translation1, translation2, translation3])
assert nearest2.x == pytest.approx(translation1.x)
assert nearest2.y == pytest.approx(translation1.y)
assert nearest2.z == pytest.approx(translation1.z)
nearest3 = origin.nearest([translation4, translation2, translation3])
assert nearest3.x == pytest.approx(translation2.x)
assert nearest3.y == pytest.approx(translation2.y)
assert nearest3.z == pytest.approx(translation2.z)
def test_dot():
one = Translation3d(1, 2, 3)
two = Translation3d(4, 5, 6)
assert one.dot(two) == pytest.approx(32.0, abs=1e-9)

56
wpimath/src/test/python/geometry/test_twist2d.py Normal file
View File

@@ -0,0 +1,56 @@
import pytest
import math
from wpimath import Twist2d, Pose2d, Rotation2d
def test_straight():
straight = Twist2d(dx=5, dy=0, dtheta=0)
straight_pose = straight.exp()
assert straight_pose.x == pytest.approx(5.0)
assert straight_pose.y == pytest.approx(0.0)
assert straight_pose.rotation().radians() == pytest.approx(0.0)
def test_quarter_circle():
quarter_circle = Twist2d(dx=5 / 2.0 * math.pi, dy=0, dtheta=math.pi / 2.0)
quarter_circle_pose = quarter_circle.exp()
assert quarter_circle_pose.x == pytest.approx(5.0)
assert quarter_circle_pose.y == pytest.approx(5.0)
assert quarter_circle_pose.rotation().degrees() == pytest.approx(90.0)
def test_diagonal_no_dtheta():
diagonal = Twist2d(dx=2, dy=2, dtheta=0)
diagonal_pose = diagonal.exp()
assert diagonal_pose.x == pytest.approx(2.0)
assert diagonal_pose.y == pytest.approx(2.0)
assert diagonal_pose.rotation().degrees() == pytest.approx(0.0)
def test_equality():
one = Twist2d(dx=5, dy=1, dtheta=3)
two = Twist2d(dx=5, dy=1, dtheta=3)
assert one == two
def test_inequality():
one = Twist2d(dx=5, dy=1, dtheta=3)
two = Twist2d(dx=5, dy=1.2, dtheta=3)
assert one != two
def test_pose2d_log():
end = Pose2d(x=5, y=5, angle=math.radians(90))
start = Pose2d()
twist = (end - start).log()
expected = Twist2d(dx=5.0 / 2.0 * math.pi, dy=0, dtheta=math.pi / 2.0)
assert expected == twist
reapplied = start + twist.exp()
assert end == reapplied

56
wpimath/src/test/python/geometry/test_twist3d.py Normal file
View File

@@ -0,0 +1,56 @@
import pytest
import math
from wpimath import Twist2d, Pose2d
def test_straight():
straight = Twist2d(dx=5, dy=0, dtheta=0)
straight_pose = straight.exp()
assert straight_pose.x == pytest.approx(5.0)
assert straight_pose.y == pytest.approx(0.0)
assert straight_pose.rotation().radians() == pytest.approx(0.0)
def test_quarter_circle():
quarter_circle = Twist2d(dx=5 / 2.0 * math.pi, dy=0, dtheta=math.pi / 2.0)
quarter_circle_pose = quarter_circle.exp()
assert quarter_circle_pose.x == pytest.approx(5.0)
assert quarter_circle_pose.y == pytest.approx(5.0)
assert quarter_circle_pose.rotation().degrees() == pytest.approx(90.0)
def test_diagonal_no_dtheta():
diagonal = Twist2d(dx=2, dy=2, dtheta=math.radians(0))
diagonal_pose = diagonal.exp()
assert diagonal_pose.x == pytest.approx(2.0)
assert diagonal_pose.y == pytest.approx(2.0)
assert diagonal_pose.rotation().degrees() == pytest.approx(0.0)
def test_equality():
one = Twist2d(dx=5, dy=1, dtheta=3)
two = Twist2d(dx=5, dy=1, dtheta=3)
assert one == two
def test_inequality():
one = Twist2d(dx=5, dy=1, dtheta=3)
two = Twist2d(dx=5, dy=1.2, dtheta=3)
assert one != two
def test_pose2d_log():
end = Pose2d(x=5, y=5, angle=math.radians(90))
start = Pose2d()
twist = (end - start).log()
expected = Twist2d(dx=5.0 / 2.0 * math.pi, dy=0, dtheta=math.pi / 2.0)
assert expected == twist
reapplied = start + twist.exp()
assert end == reapplied