Adds the from_angle_axis_translation function to the dual_quaternion module.

PiperOrigin-RevId: 411801030

Author: G4G

tensorflow_graphics/geometry/transformation/dual_quaternion.py

@@ -335,5 +335,63 @@ def to_rotation_translation(
     return rotation, translation
 
 
+def from_axis_angle_translation(axis: type_alias.TensorLike,
+                                angle: type_alias.TensorLike,
+                                translation_vector: type_alias.TensorLike,
+                                name: str = "dual_quat_from_axis_angle_trans"
+                                ) -> type_alias.TensorLike:
+  """Converts an axis-angle rotation and translation to a dual quaternion.
+
+  Note:
+    In the following, A1 to An are optional batch dimensions.
+
+  Args:
+    axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension
+      represents a normalized axis.
+    angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension
+      represents an angle.
+    translation_vector: A `[A1, ..., An, 3]`-tensor, where the last dimension
+      represents a translation vector.
+    name: A name for this op that defaults to "dual_quat_from_axis_angle_trans".
+
+  Returns:
+    A `[A1, ..., An, 8]`-tensor, where the last dimension represents a
+    normalized dual quaternion.
+
+  Raises:
+    ValueError: If the shape of `axis`, `angle`, or `translation_vector`
+    is not supported.
+  """
+  with tf.name_scope(name):
+    axis = tf.convert_to_tensor(value=axis)
+    angle = tf.convert_to_tensor(value=angle)
+    translation_vector = tf.convert_to_tensor(value=translation_vector)
+
+    shape.check_static(tensor=axis,
+                       tensor_name="axis",
+                       has_dim_equals=(-1, 3))
+    shape.check_static(tensor=angle,
+                       tensor_name="angle",
+                       has_dim_equals=(-1, 1))
+    shape.check_static(tensor=translation_vector,
+                       tensor_name="translation_vector",
+                       has_dim_equals=(-1, 3))
+    shape.compare_batch_dimensions(tensors=(axis, angle, translation_vector),
+                                   last_axes=-2,
+                                   broadcast_compatible=True)
+
+    scalar_shape = tf.concat((tf.shape(translation_vector)[:-1], (1,)), axis=-1)
+    dtype = translation_vector.dtype
+
+    quaternion_rotation = quaternion.from_axis_angle(axis, angle)
+    quaternion_translation = tf.concat(
+        (translation_vector, tf.zeros(scalar_shape, dtype)), axis=-1)
+
+    dual_quaternion_dual_part = 0.5 * quaternion.multiply(
+        quaternion_translation, quaternion_rotation)
+
+    return tf.concat((quaternion_rotation, dual_quaternion_dual_part), axis=-1)
+
+
 # API contains all public functions and classes.
 __all__ = export_api.get_functions_and_classes()

tensorflow_graphics/geometry/transformation/tests/dual_quaternion_test.py

@@ -234,7 +234,7 @@ def test_from_rotation_translation_jacobian_random(self):
     self.assert_jacobian_is_finite_fn(dual_quaternion.from_rotation_translation,
                                       [rotation_init, translation_init])
 
-  def test_from_rotation_matrix_normalized_random(self):
+  def test_from_rotation_translation_normalized_random(self):
     (euler_angles, translation
     ) = test_helpers.generate_random_test_euler_angles_translations()
     rotation = quaternion.from_euler(euler_angles)
@@ -246,22 +246,19 @@ def test_from_rotation_matrix_normalized_random(self):
         dual_quaternion.is_normalized(random_dual_quaternion),
         np.ones(shape=rotation.shape[:-1] + (1,), dtype=bool))
 
-  def test_from_rotation_matrix_random(self):
+  def test_from_rotation_translation_random(self):
     (euler_angles_gt, translation_gt
     ) = test_helpers.generate_random_test_euler_angles_translations()
     rotation_gt = quaternion.from_euler(euler_angles_gt)
 
     dual_quaternion_output = dual_quaternion.from_rotation_translation(
         rotation_gt, translation_gt)
-    dual_quaternion_real = dual_quaternion_output[..., 0:4]
-    dual_quaternion_dual = dual_quaternion_output[..., 4:8]
-    rotation = rotation_matrix_3d.from_quaternion(dual_quaternion_real)
-    translation = 2.0 * quaternion.multiply(
-        dual_quaternion_dual, quaternion.inverse(dual_quaternion_real))
-    translation = translation[..., 0:3]
+
+    rotation, translation = dual_quaternion.to_rotation_translation(
+        dual_quaternion_output)
 
     self.assertAllClose(rotation_matrix_3d.from_quaternion(rotation_gt),
-                        rotation)
+                        rotation_matrix_3d.from_quaternion(rotation))
     self.assertAllClose(translation_gt, translation)
 
   @flagsaver.flagsaver(tfg_add_asserts_to_graph=False)
@@ -286,7 +283,7 @@ def to_translation(input_dual_quaternion):
 
     self.assert_jacobian_is_finite_fn(to_translation, [rnd_dual_quaternion])
 
-  def test_to_rotation_matrix_random(self):
+  def test_to_rotation_translation_random(self):
     (euler_angles_gt, translation_gt
     ) = test_helpers.generate_random_test_euler_angles_translations()
     rotation_gt = quaternion.from_euler(euler_angles_gt)
@@ -300,6 +297,39 @@ def test_to_rotation_matrix_random(self):
                         rotation_matrix_3d.from_quaternion(rotation))
     self.assertAllClose(translation_gt, translation)
 
+  def test_from_axis_angle_translation_normalized_random(self):
+    (random_axis,
+     random_angle,
+     random_translation
+     ) = test_helpers.generate_random_test_axis_angle_translation()
+
+    random_dual_quaternion = dual_quaternion.from_axis_angle_translation(
+        random_axis,
+        random_angle,
+        random_translation)
+
+    self.assertAllEqual(
+        dual_quaternion.is_normalized(random_dual_quaternion),
+        np.ones(shape=random_dual_quaternion.shape[:-1] + (1,), dtype=bool))
+
+  def test_from_axis_angle_translation_random(self):
+    (axis_gt,
+     angle_gt,
+     translation_gt
+    ) = test_helpers.generate_random_test_axis_angle_translation()
+
+    rotation_gt = quaternion.from_axis_angle(axis_gt, angle_gt)
+    dual_quaternion_output = dual_quaternion.from_axis_angle_translation(
+        axis_gt,
+        angle_gt,
+        translation_gt)
+
+    rotation, translation = dual_quaternion.to_rotation_translation(
+        dual_quaternion_output)
+
+    self.assertAllClose(rotation_gt, rotation)
+    self.assertAllClose(translation_gt, translation)
+
 
 if __name__ == "__main__":
   test_case.main()

tensorflow_graphics/geometry/transformation/tests/test_helpers.py

@@ -253,3 +253,14 @@ def generate_preset_test_lbs_blend():
                                [0.88587099, -0.09324637, -0.45012815]]]])
 
   return points, weights, rotations, translations, blended_points
+
+
+def generate_random_test_axis_angle_translation():
+  """Generates random test angles, axes, translations."""
+  tensor_dimensions = np.random.randint(3)
+  tensor_shape = np.random.randint(1, 10, size=(tensor_dimensions)).tolist()
+  random_axis = np.random.uniform(size=tensor_shape + [3])
+  random_axis /= np.linalg.norm(random_axis, axis=-1, keepdims=True)
+  random_angle = np.random.uniform(size=tensor_shape + [1])
+  random_translation = np.random.uniform(size=tensor_shape + [3])
+  return random_axis, random_angle, random_translation