add pose correctness tests

Author: mugamma

tests/python_tests/test_geometry.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 from scipy.spatial.transform import Rotation as Rot
+from scipy.spatial.transform import RigidTransform as Rigid
 
 from genmetaballs._genmetaballs_bindings import geometry as geometry
 
@@ -32,16 +33,18 @@ def test_vec3d_repr_returns_valid_string() -> None:
     ]
  )
 def test_vec3d_ops(rng: np.random.Generator, np_op, vec3d_op) -> None:
-    _a, _b = rng.uniform(size=3), rng.uniform(size=3)
-    a, b = geometry.Vec3D(*_a), geometry.Vec3D(*_b)
-    c = vec3d_op(a, b)
-    _c = np_op(_a, _b)
-    assert np.allclose(_c, np.array([c.x, c.y, c.z]))
+    _a, _b = rng.uniform(size=(100, 3)), rng.uniform(size=(100, 3))
+    for i in range(100):
+        a, b = geometry.Vec3D(*_a[i]), geometry.Vec3D(*_b[i])
+        c = vec3d_op(a, b)
+        _c = np_op(_a[i], _b[i])
+        assert np.allclose(_c, np.array([c.x, c.y, c.z]))
 
 def test_vec3d_dot(rng: np.random.Generator) -> None:
-    _a, _b = rng.uniform(size=3), rng.uniform(size=3)
-    a, b = geometry.Vec3D(*_a), geometry.Vec3D(*_b)
-    assert np.allclose(np.dot(_a, _b), geometry.dot(a, b))
+    _a, _b = rng.uniform(size=(100, 3)), rng.uniform(size=(100, 3))
+    for i in range(100):
+        a, b = geometry.Vec3D(*_a[i]), geometry.Vec3D(*_b[i])
+        assert np.allclose(np.dot(_a[i], _b[i]), geometry.dot(a, b))
 
 def test_rotation_apply(rng: np.random.Generator) -> None:
     quats = rng.uniform(-1, 1, size=(100, 4))
@@ -104,3 +107,76 @@ def test_rotation_inv(rng: np.random.Generator) -> None:
                rtol=1e-5,
                atol=1e-6,
            )
+
+def test_pose_apply(rng: np.random.Generator) -> None:
+    exp_coords = rng.uniform(size=(100, 6))
+    vecs = rng.uniform(size=(100, 3))
+
+    for i in range(100):
+        pose_scipy = Rigid.from_exp_coords(exp_coords[i])
+
+        pose_geom = geometry.Pose.from_components(
+            rot=geometry.Rotation.from_quat(*pose_scipy.rotation.as_quat()),
+            tran=geometry.Vec3D(*pose_scipy.translation),
+        )
+        vec_geom = geometry.Vec3D(*vecs[i])
+
+        result_scipy = pose_scipy.apply(vecs[i])
+        result_geom = pose_geom.apply(vec_geom)
+
+        assert np.allclose(result_scipy, np.array([result_geom.x, result_geom.y, result_geom.z]))
+
+def test_pose_compose(rng: np.random.Generator) -> None:
+    exp_coords1 = rng.uniform(size=(100, 6))
+    exp_coords2 = rng.uniform(size=(100, 6))
+    vecs = rng.uniform(size=(50, 3))
+
+    for i in range(100):
+        pose1_scipy = Rigid.from_exp_coords(exp_coords1[i])
+        pose2_scipy = Rigid.from_exp_coords(exp_coords2[i])
+        composed_scipy = pose1_scipy * pose2_scipy
+
+        pose1_geom = geometry.Pose.from_components(
+            rot=geometry.Rotation.from_quat(*pose1_scipy.rotation.as_quat()),
+            tran=geometry.Vec3D(*pose1_scipy.translation),
+        )
+        pose2_geom = geometry.Pose.from_components(
+            rot=geometry.Rotation.from_quat(*pose2_scipy.rotation.as_quat()),
+            tran=geometry.Vec3D(*pose2_scipy.translation),
+        )
+        composed_geom = pose1_geom.compose(pose2_geom)
+
+        for j in range(50):
+            vec_geom = geometry.Vec3D(*vecs[j])
+            result_scipy = composed_scipy.apply(vecs[j])
+            result_geom = composed_geom.apply(vec_geom)
+            assert np.allclose(
+               result_scipy,
+               np.array([result_geom.x, result_geom.y, result_geom.z]),
+               rtol=1e-5,
+               atol=1e-6,
+           )
+
+def test_pose_inv(rng: np.random.Generator) -> None:
+    exp_coords = rng.uniform(size=(100, 6))
+    vecs = rng.uniform(size=(50, 3))
+
+    for i in range(100):
+        pose_scipy = Rigid.from_exp_coords(exp_coords[i])
+
+        pose = geometry.Pose.from_components(
+            rot=geometry.Rotation.from_quat(*pose_scipy.rotation.as_quat()),
+            tran=geometry.Vec3D(*pose_scipy.translation),
+        )
+        poseinv = pose.inv()
+        composed = pose.compose(poseinv)
+
+        for j in range(50):
+            vec = geometry.Vec3D(*vecs[j])
+            vec_ = composed.apply(vec)
+            assert np.allclose(
+               np.array([vec.x, vec.y, vec.z]),
+               np.array([vec_.x, vec_.y, vec_.z]),
+               rtol=1e-5,
+               atol=1e-6,
+           )