Test 1D and 3D cases

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
 
 

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)