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

wpimath/src/test/python/geometry/test_quaternion.py

@@ -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