xform: Add new utility

Author: rhamel3d

metafold/utils.py

@@ -0,0 +1,30 @@
+from numpy.typing import ArrayLike
+from scipy.spatial.transform import Rotation as R
+from typing import Optional
+import numpy as np
+
+
+def xform(
+    translation: Optional[ArrayLike] = None,
+    rotation: Optional[ArrayLike] = None,
+    scale: Optional[ArrayLike] = None,
+) -> np.array:
+    """Compose transformation matrix.
+
+       Args:
+            translation: Translation in the x, y, and z directions.
+            rotation: Euler angles in degrees. Rotation is applied in the order Rz, Ry, Rx.
+            scale: Scale factor applied in the x, y, z directions before rotation.
+
+        Returns:
+            4 x 4 affine transformation matrix.
+    """
+    translation = translation or np.array([0.0, 0.0, 0.0])
+    rotation = rotation or np.array([0.0, 0.0, 0.0])
+    scale = scale or np.array([1.0, 1.0, 1.0])
+
+    r = R.from_euler("xyz", rotation,  degrees=True).as_matrix()
+    m = np.eye(4)
+    m[:3, :3] = np.diag(scale) @ r
+    m[3, :3] = translation
+    return m

requirements.txt

@@ -1,3 +1,5 @@
 attrs~=23.2
 auth0-python~=4.7
+numpy ~= 1.23
 requests~=2.31
+scipy ~= 1.11

tests/test_utils.py

@@ -0,0 +1,53 @@
+from metafold.utils import xform
+from numpy.testing import assert_allclose
+
+
+def test_xform():
+    m = xform([1.0, 2.0, 3.0])  # T
+    assert_allclose(m, [
+        [1.0, 0.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 1.0, 0.0],
+        [1.0, 2.0, 3.0, 1.0],
+    ], atol=1.0e-7)
+
+    m = xform(
+        [1.0, 2.0, 3.0],      # T
+        [180.0, -90.0, 0.0])  # R
+    assert_allclose(m, [
+        [0.0,  0.0, 1.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [1.0,  0.0, 0.0, 0.0],
+        [1.0,  2.0, 3.0, 1.0],
+    ], atol=1.0e-7)
+
+    m = xform(
+        [1.0, 2.0, 3.0],        # T
+        scale=[2.0, 1.0, 1.0])  # S
+    assert_allclose(m, [
+        [2.0, 0.0, 0.0, 0.0],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 1.0, 0.0],
+        [1.0, 2.0, 3.0, 1.0],
+    ], atol=1.0e-7)
+
+    m = xform(
+        rotation=[180.0, -90.0, 0.0],  # R
+        scale=[2.0, 1.0, 1.0])         # S
+    assert_allclose(m, [
+        [0.0,  0.0, 2.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [1.0,  0.0, 0.0, 0.0],
+        [0.0,  0.0, 0.0, 1.0],
+    ], atol=1.0e-7)
+
+    m = xform(
+        [1.0, 2.0, 3.0],      # T
+        [180.0, -90.0, 0.0],  # R
+        [2.0, 1.0, 1.0])      # S
+    assert_allclose(m, [
+        [0.0,  0.0, 2.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [1.0,  0.0, 0.0, 0.0],
+        [1.0,  2.0, 3.0, 1.0],
+    ], atol=1.0e-7)