From 7ab58b364c1be706a7b6c39305bd1f05c7b772bd Mon Sep 17 00:00:00 2001 From: Alexander Fabisch Date: Tue, 19 Nov 2024 17:46:36 +0100 Subject: [PATCH] Test 1D and 3D cases --- pytransform3d/batch_rotations.py | 8 +++++++- pytransform3d/test/test_batch_rotations.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/pytransform3d/batch_rotations.py b/pytransform3d/batch_rotations.py index c72fa5f5e..dc4412308 100644 --- a/pytransform3d/batch_rotations.py +++ b/pytransform3d/batch_rotations.py @@ -53,10 +53,14 @@ def norm_axis_angles(a): of the axis vector is 1 and the angle is in [0, pi). No rotation is represented by [1, 0, 0, 0]. """ + a = np.asarray(a) + only_one = a.ndim == 1 + + a = np.atleast_2d(a) + angles = a[..., 3] norm = np.linalg.norm(a[..., :3], axis=-1) - res = np.ones_like(a) # Create masks for elements where angle or norm is zero @@ -75,6 +79,8 @@ def norm_axis_angles(a): angle_normalized[negative_angle_mask] *= -1.0 res[non_zero_mask, 3] = angle_normalized[non_zero_mask] + if only_one: + res = res[0] return res diff --git a/pytransform3d/test/test_batch_rotations.py b/pytransform3d/test/test_batch_rotations.py index 93e0a4972..cfea89c04 100644 --- a/pytransform3d/test/test_batch_rotations.py +++ b/pytransform3d/test/test_batch_rotations.py @@ -48,11 +48,23 @@ def test_norm_axis_angles(): pbr.norm_vectors(rng.standard_normal(size=(10, 3))), np.pi + rng.uniform(size=(10, 1)) )) + + # 1D + assert_array_almost_equal( + pbr.norm_axis_angles(A[0]), pr.norm_axis_angle(A[0])) + + # 2D A_norm = pbr.norm_axis_angles(A) for a, a_norm in zip(A, A_norm): assert pytest.approx(a[3]) != a_norm[3] assert_array_almost_equal(a_norm, pr.norm_axis_angle(a)) + # 3D + A3D = A.reshape(5, 2, -1) + A3D_norm = pbr.norm_axis_angles(A3D) + for a, a_norm in zip(A3D.reshape(10, -1), A3D_norm.reshape(10, -1)): + assert_array_almost_equal(a_norm, pr.norm_axis_angle(a)) + def test_angles_between_vectors_0dims(): rng = np.random.default_rng(228)