From 46c6ecb1d9a9aa5afd2bf44cf6bc34cc408ba7ac Mon Sep 17 00:00:00 2001
From: cremebrule <84cremebrule@gmail.com>
Date: Thu, 19 Dec 2024 12:41:43 -0800
Subject: [PATCH] re-add transformation_matrix

---
 omnigibson/utils/transform_utils.py    | 23 +++++++++++++++++++++++
 omnigibson/utils/transform_utils_np.py | 22 ++++++++++++++++++++++
 2 files changed, 45 insertions(+)

diff --git a/omnigibson/utils/transform_utils.py b/omnigibson/utils/transform_utils.py
index 98eb6bb73..985716e9d 100644
--- a/omnigibson/utils/transform_utils.py
+++ b/omnigibson/utils/transform_utils.py
@@ -981,6 +981,29 @@ def rotation_matrix(angle: float, direction: torch.Tensor) -> torch.Tensor:
     return R
 
 
+@torch.compile
+def transformation_matrix(angle: float, direction: torch.Tensor, point: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """
+    Returns a 4x4 homogeneous transformation matrix to rotate about axis defined by point and direction.
+    Args:
+        angle (float): Magnitude of rotation in radians
+        direction (torch.Tensor): (ax,ay,az) axis about which to rotate
+        point (Optional[torch.Tensor]): If specified, is the (x,y,z) point about which the rotation will occur
+    Returns:
+        torch.Tensor: 4x4 homogeneous transformation matrix
+    """
+    R = rotation_matrix(angle, direction)
+
+    M = torch.eye(4, dtype=torch.float32, device=direction.device)
+    M[:3, :3] = R
+
+    if point is not None:
+        # Rotation not about origin
+        point = point.to(dtype=torch.float32)
+        M[:3, 3] = point - R @ point
+    return M
+
+
 @torch.compile
 def clip_translation(dpos, limit):
     """
diff --git a/omnigibson/utils/transform_utils_np.py b/omnigibson/utils/transform_utils_np.py
index 6cd278322..b530f9a49 100644
--- a/omnigibson/utils/transform_utils_np.py
+++ b/omnigibson/utils/transform_utils_np.py
@@ -888,6 +888,28 @@ def rotation_matrix(angle, direction, point=None):
     return M
 
 
+def transformation_matrix(angle, direction, point=None):
+    """
+    Returns a 4x4 homogeneous transformation matrix to rotate about axis defined by point and direction.
+    Args:
+        angle (float): Magnitude of rotation in radians
+        direction (np.array): (ax,ay,az) axis about which to rotate
+        point (bool): If specified, is the (x,y,z) point about which the rotation will occur
+
+    Returns:
+        np.array: 4x4 homogeneous transformation matrix
+    """
+    R = rotation_matrix(angle, direction)
+
+    M = np.eye(4)
+    M[:3, :3] = R
+
+    if point is not None:
+        # Rotation not about origin
+        M[:3, 3] = point - R @ point
+    return M
+
+
 def clip_translation(dpos, limit):
     """
     Limits a translation (delta position) to a specified limit