From 8faa0a901e19b60e123757d644a9a4260cf08476 Mon Sep 17 00:00:00 2001 From: cremebrule <84cremebrule@gmail.com> Date: Thu, 19 Dec 2024 12:04:16 -0800 Subject: [PATCH] verify that th and np transform utils have 1-to-1 corresponding functions --- omnigibson/utils/transform_utils.py | 117 ++---- omnigibson/utils/transform_utils_np.py | 546 +++++++++++-------------- omnigibson/utils/usd_utils.py | 8 +- 3 files changed, 276 insertions(+), 395 deletions(-) diff --git a/omnigibson/utils/transform_utils.py b/omnigibson/utils/transform_utils.py index c2301bb61..f9db861db 100644 --- a/omnigibson/utils/transform_utils.py +++ b/omnigibson/utils/transform_utils.py @@ -1,6 +1,8 @@ """ Utility functions of matrix and vector transformations. +NOTE: This file has a 1-to-1 correspondence to transform_utils_np.py + NOTE: convention for quaternions is (x, y, z, w) """ @@ -42,7 +44,7 @@ @torch.compile -def copysign(a, b): +def _copysign(a, b): # type: (float, torch.Tensor) -> torch.Tensor a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0]) return torch.abs(a) * torch.sign(b) @@ -308,6 +310,32 @@ def quat_slerp(quat0, quat1, frac, shortestpath=True, eps=1.0e-15): return val.reshape(list(quat_shape)) +@torch.compile +def random_quaternion(num_quaternions: int = 1) -> torch.Tensor: + """ + Generate random rotation quaternions, uniformly distributed over SO(3). + + Arguments: + num_quaternions (int): number of quaternions to generate (default: 1) + + Returns: + torch.Tensor: A tensor of shape (num_quaternions, 4) containing random unit quaternions. + """ + # Generate four random numbers between 0 and 1 + rand = torch.rand(num_quaternions, 4) + + # Use the formula from Ken Shoemake's "Uniform Random Rotations" + r1 = torch.sqrt(1.0 - rand[:, 0]) + r2 = torch.sqrt(rand[:, 0]) + t1 = 2 * torch.pi * rand[:, 1] + t2 = 2 * torch.pi * rand[:, 2] + + quaternions = torch.stack([r1 * torch.sin(t1), r1 * torch.cos(t1), r2 * torch.sin(t2), r2 * torch.cos(t2)], dim=1) + + return quaternions + + + @torch.compile def random_axis_angle(angle_limit: float = 2.0 * math.pi): """ @@ -444,6 +472,20 @@ def mat2quat(rmat: torch.Tensor) -> torch.Tensor: return quat +def mat2quat_batch(rmat: torch.Tensor) -> torch.Tensor: + """ + Converts given rotation matrix to quaternion. Version optimized for batch operations + + Args: + rmat (torch.Tensor): (3, 3) or (..., 3, 3) rotation matrix + + Returns: + torch.Tensor: (4,) or (..., 4) (x,y,z,w) float quaternion angles + """ + # For torch, no different than basic version + return mat2quat(rmat) + + @torch.compile def mat2pose(hmat): """ @@ -547,7 +589,7 @@ def quat2euler(q): roll = torch.atan2(sinr_cosp, cosr_cosp) # pitch (y-axis rotation) sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx]) - pitch = torch.where(torch.abs(sinp) >= 1, copysign(math.pi / 2.0, sinp), torch.asin(sinp)) + pitch = torch.where(torch.abs(sinp) >= 1, _copysign(math.pi / 2.0, sinp), torch.asin(sinp)) # yaw (z-axis rotation) siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy]) cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz] @@ -940,31 +982,6 @@ 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): """ @@ -1292,31 +1309,6 @@ def integer_spiral_coordinates(n: int) -> Tuple[int, int]: return int(x), int(y) -@torch.compile -def random_quaternion(num_quaternions: int = 1) -> torch.Tensor: - """ - Generate random rotation quaternions, uniformly distributed over SO(3). - - Arguments: - num_quaternions: int, number of quaternions to generate (default: 1) - - Returns: - torch.Tensor: A tensor of shape (num_quaternions, 4) containing random unit quaternions. - """ - # Generate four random numbers between 0 and 1 - rand = torch.rand(num_quaternions, 4) - - # Use the formula from Ken Shoemake's "Uniform Random Rotations" - r1 = torch.sqrt(1.0 - rand[:, 0]) - r2 = torch.sqrt(rand[:, 0]) - t1 = 2 * torch.pi * rand[:, 1] - t2 = 2 * torch.pi * rand[:, 2] - - quaternions = torch.stack([r1 * torch.sin(t1), r1 * torch.cos(t1), r2 * torch.sin(t2), r2 * torch.cos(t2)], dim=1) - - return quaternions - - @torch.compile def transform_points(points: torch.Tensor, matrix: torch.Tensor, translate: bool = True) -> torch.Tensor: """ @@ -1353,27 +1345,6 @@ def transform_points(points: torch.Tensor, matrix: torch.Tensor, translate: bool return torch.mm(matrix[:dim, :dim], points.t()).t() -@torch.compile -def quaternions_close(q1: torch.Tensor, q2: torch.Tensor, atol: float = 1e-3) -> bool: - """ - Whether two quaternions represent the same rotation, - allowing for the possibility that one is the negative of the other. - - Arguments: - q1: torch.Tensor - First quaternion - q2: torch.Tensor - Second quaternion - atol: float - Absolute tolerance for comparison - - Returns: - bool - Whether the quaternions are close - """ - return torch.allclose(q1, q2, atol=atol) or torch.allclose(q1, -q2, atol=atol) - - @torch.compile def orientation_error(desired, current): """ diff --git a/omnigibson/utils/transform_utils_np.py b/omnigibson/utils/transform_utils_np.py index 5c4997374..d0d67ea84 100644 --- a/omnigibson/utils/transform_utils_np.py +++ b/omnigibson/utils/transform_utils_np.py @@ -1,6 +1,9 @@ """ Utility functions of matrix and vector transformations. +NOTE: This file has a 1-to-1 correspondence to transform_utils.py. By default, we use scipy for most transform-related + operations, but we optionally implement numba versions for functions that are often called in "batch" mode + NOTE: convention for quaternions is (x, y, z, w) """ @@ -47,72 +50,90 @@ _TUPLE2AXES = dict((v, k) for k, v in _AXES2TUPLE.items()) -def ewma_vectorized(data, alpha, offset=None, dtype=None, order="C", out=None): - """ - Calculates the exponential moving average over a vector. - Will fail for large inputs. +def anorm(x, axis=None, keepdims=False): + """Compute L2 norms alogn specified axes.""" + return np.linalg.norm(x, axis=axis, keepdims=keepdims) - Args: - data (Iterable): Input data - alpha (float): scalar in range (0,1) - The alpha parameter for the moving average. - offset (None or float): If specified, the offset for the moving average. None defaults to data[0]. - dtype (None or type): Data type used for calculations. If None, defaults to float64 unless - data.dtype is float32, then it will use float32. - order (None or str): Order to use when flattening the data. Valid options are {'C', 'F', 'A'}. - None defaults to 'C'. - out (None or np.array): If specified, the location into which the result is stored. If provided, it must have - the same shape as the input. If not provided or `None`, - a freshly-allocated array is returned. - Returns: - np.array: Exponential moving average from @data +def normalize(v, axis=None, eps=1e-10): + """L2 Normalize along specified axes.""" + norm = anorm(v, axis=axis, keepdims=True) + return v / np.where(norm < eps, eps, norm) + + +def dot(v1, v2, dim=-1, keepdim=False): + return np.sum(v1 * v2, axis=dim, keepdims=keepdim) + + +def unit_vector(data, axis=None, out=None): """ - data = np.array(data, copy=False) + Returns ndarray normalized by length, i.e. eucledian norm, along axis. - if dtype is None: - if data.dtype == np.float32: - dtype = np.float32 - else: - dtype = np.float64 - else: - dtype = np.dtype(dtype) + E.g.: + >>> v0 = numpy.random.random(3) + >>> v1 = unit_vector(v0) + >>> numpy.allclose(v1, v0 / numpy.linalg.norm(v0)) + True - if data.ndim > 1: - # flatten input - data = data.reshape(-1, order) + >>> v0 = numpy.random.rand(5, 4, 3) + >>> v1 = unit_vector(v0, axis=-1) + >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=2)), 2) + >>> numpy.allclose(v1, v2) + True - if out is None: - out = np.empty_like(data, dtype=dtype) - else: - assert out.shape == data.shape - assert out.dtype == dtype + >>> v1 = unit_vector(v0, axis=1) + >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=1)), 1) + >>> numpy.allclose(v1, v2) + True - if data.size < 1: - # empty input, return empty array - return out + >>> v1 = numpy.empty((5, 4, 3), dtype=numpy.float32) + >>> unit_vector(v0, axis=1, out=v1) + >>> numpy.allclose(v1, v2) + True - if offset is None: - offset = data[0] + >>> list(unit_vector([])) + [] - alpha = np.array(alpha, copy=False).astype(dtype, copy=False) + >>> list(unit_vector([1.0])) + [1.0] + + Args: + data (np.array): data to normalize + axis (None or int): If specified, determines specific axis along data to normalize + out (None or np.array): If specified, will store computation in this variable - # scaling_factors -> 0 as len(data) gets large - # this leads to divide-by-zeros below - scaling_factors = np.power(1.0 - alpha, np.arange(data.size + 1, dtype=dtype), dtype=dtype) - # create cumulative sum array - np.multiply(data, (alpha * scaling_factors[-2]) / scaling_factors[:-1], dtype=dtype, out=out) - np.cumsum(out, dtype=dtype, out=out) + Returns: + None or np.array: If @out is not specified, will return normalized vector. Otherwise, stores the output in @out + """ + if out is None: + data = np.array(data, dtype=np.float32, copy=True) + if data.ndim == 1: + data /= math.sqrt(np.dot(data, data)) + return data + else: + if out is not data: + out[:] = np.array(data, copy=False) + data = out + length = np.atleast_1d(np.sum(data * data, axis)) + np.sqrt(length, length) + if axis is not None: + length = np.expand_dims(length, axis) + data /= length + if out is None: + return data - # cumsums / scaling - out /= scaling_factors[-2::-1] - if offset != 0: - offset = np.array(offset, copy=False).astype(dtype, copy=False) - # add offsets - out += offset * scaling_factors[1:] +def quat_apply(quat, vec): + """ + Apply a quaternion rotation to a vector (equivalent to R.from_quat(x).apply(y)) + Args: + quat (np.array): (4,) or (N, 4) or (N, 1, 4) quaternion in (x, y, z, w) format + vec (np.array): (3,) or (M, 3) or (1, M, 3) vector to rotate - return out + Returns: + np.array: (M, 3) or (N, M, 3) rotated vector + """ + return R.from_quat(quat).apply(vec) def convert_quat(q, to="xyzw"): @@ -270,38 +291,29 @@ def quat_slerp(quat0, quat1, fraction, shortestpath=True): return q0 -def random_quat(rand=None): +@jit(nopython=True) +def random_quaternion(num_quaternions=1): """ - Return uniform random unit quaternion. + Generate random rotation quaternions, uniformly distributed over SO(3). - E.g.: - >>> q = random_quat() - >>> np.allclose(1.0, vector_norm(q)) - True - >>> q = random_quat(np.random.random(3)) - >>> q.shape - (4,) - - Args: - rand (3-array or None): If specified, must be three independent random variables that are uniformly distributed - between 0 and 1. + Arguments: + num_quaternions (int): number of quaternions to generate (default: 1) Returns: - np.array: (x,y,z,w) random quaternion + np.array: A tensor of shape (num_quaternions, 4) containing random unit quaternions. """ - if rand is None: - rand = np.random.rand(3) - else: - assert len(rand) == 3 - r1 = np.sqrt(1.0 - rand[0]) - r2 = np.sqrt(rand[0]) - pi2 = math.pi * 2.0 - t1 = pi2 * rand[1] - t2 = pi2 * rand[2] - return np.array( - (np.sin(t1) * r1, np.cos(t1) * r1, np.sin(t2) * r2, np.cos(t2) * r2), - dtype=np.float32, - ) + # Generate four random numbers between 0 and 1 + rand = np.random.rand(num_quaternions, 4) + + # Use the formula from Ken Shoemake's "Uniform Random Rotations" + r1 = np.sqrt(1.0 - rand[:, 0]) + r2 = np.sqrt(rand[:, 0]) + t1 = 2 * np.pi * rand[:, 1] + t2 = 2 * np.pi * rand[:, 2] + + quaternions = np.stack((r1 * np.sin(t1), r1 * np.cos(t1), r2 * np.sin(t2), r2 * np.cos(t2)), axis=1) + + return quaternions def random_axis_angle(angle_limit=None, random_state=None): @@ -338,47 +350,53 @@ def random_axis_angle(angle_limit=None, random_state=None): return random_axis, random_angle -def vec(values): - """ - Converts value tuple into a numpy vector. - - Args: - values (n-array): a tuple of numbers - - Returns: - np.array: vector of given values - """ - return np.array(values, dtype=np.float32) +def quat2mat(quaternion): + if quaternion.dtype != np.float32: + quaternion = quaternion.astype(np.float32) + return _quat2mat(quaternion) -def mat4(array): +@jit(nopython=True) +def _quat2mat(quaternion): """ - Converts an array to 4x4 matrix. + Convert quaternions into rotation matrices. Args: - array (n-array): the array in form of vec, list, or tuple + quaternion (torch.Tensor): A tensor of shape (..., 4) representing batches of quaternions (w, x, y, z). Returns: - np.array: a 4x4 numpy matrix + torch.Tensor: A tensor of shape (..., 3, 3) representing batches of rotation matrices. """ - return np.array(array, dtype=np.float32).reshape((4, 4)) + # broadcast array is necessary to use numba parallel mode + q1, q2 = np.broadcast_arrays(quaternion[..., np.newaxis], quaternion[..., np.newaxis, :]) + outer = q1 * q2 + + # Extract the necessary components + xx = outer[..., 0, 0] + yy = outer[..., 1, 1] + zz = outer[..., 2, 2] + xy = outer[..., 0, 1] + xz = outer[..., 0, 2] + yz = outer[..., 1, 2] + xw = outer[..., 0, 3] + yw = outer[..., 1, 3] + zw = outer[..., 2, 3] + rotation_matrix = np.empty(quaternion.shape[:-1] + (3, 3), dtype=np.float32) -def mat2pose(hmat): - """ - Converts a homogeneous 4x4 matrix into pose. + rotation_matrix[..., 0, 0] = 1 - 2 * (yy + zz) + rotation_matrix[..., 0, 1] = 2 * (xy - zw) + rotation_matrix[..., 0, 2] = 2 * (xz + yw) - Args: - hmat (np.array): a 4x4 homogeneous matrix + rotation_matrix[..., 1, 0] = 2 * (xy + zw) + rotation_matrix[..., 1, 1] = 1 - 2 * (xx + zz) + rotation_matrix[..., 1, 2] = 2 * (yz - xw) - Returns: - 2-tuple: - - (np.array) (x,y,z) position array in cartesian coordinates - - (np.array) (x,y,z,w) orientation array in quaternion form - """ - pos = hmat[:3, 3] - orn = mat2quat(hmat[:3, :3]) - return pos, orn + rotation_matrix[..., 2, 0] = 2 * (xz - yw) + rotation_matrix[..., 2, 1] = 2 * (yz + xw) + rotation_matrix[..., 2, 2] = 1 - 2 * (xx + yy) + + return rotation_matrix def mat2quat(rmat): @@ -466,6 +484,23 @@ def mat2quat_batch(rmat): return quat +def mat2pose(hmat): + """ + Converts a homogeneous 4x4 matrix into pose. + + Args: + hmat (np.array): a 4x4 homogeneous matrix + + Returns: + 2-tuple: + - (np.array) (x,y,z) position array in cartesian coordinates + - (np.array) (x,y,z,w) orientation array in quaternion form + """ + pos = hmat[:3, 3] + orn = mat2quat(hmat[:3, :3]) + return pos, orn + + def vec2quat(vec, up=(0, 0, 1.0)): """ Converts given 3d-direction vector @vec to quaternion orientation with respect to another direction vector @up @@ -484,87 +519,70 @@ def vec2quat(vec, up=(0, 0, 1.0)): return mat2quat(np.array([vec_n, s_n, u_n]).T) -def euler2mat(euler): +def euler2quat(euler): """ - Converts euler angles into rotation matrix form + Converts euler angles into quaternion form Args: euler (np.array): (r,p,y) angles Returns: - np.array: 3x3 rotation matrix + np.array: (x,y,z,w) float quaternion angles Raises: AssertionError: [Invalid input shape] """ - - euler = np.asarray(euler, dtype=np.float64) - assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler) - - return R.from_euler("xyz", euler).as_matrix() + return R.from_euler("xyz", euler).as_quat() -def mat2euler(rmat): +def quat2euler(quat): """ - Converts given rotation matrix to euler angles in radian. + Converts euler angles into quaternion form Args: - rmat (np.array): 3x3 rotation matrix + quat (np.array): (x,y,z,w) float quaternion angles Returns: - np.array: (r,p,y) converted euler angles in radian vec3 float - """ - M = np.array(rmat, dtype=np.float32, copy=False)[:3, :3] - return R.from_matrix(M).as_euler("xyz") - + np.array: (r,p,y) angles -def quat2mat(quaternion): - if quaternion.dtype != np.float32: - quaternion = quaternion.astype(np.float32) - return _quat2mat(quaternion) + Raises: + AssertionError: [Invalid input shape] + """ + return R.from_quat(quat).as_euler("xyz") -@jit(nopython=True) -def _quat2mat(quaternion): +def euler2mat(euler): """ - Convert quaternions into rotation matrices. + Converts euler angles into rotation matrix form Args: - quaternion (torch.Tensor): A tensor of shape (..., 4) representing batches of quaternions (w, x, y, z). + euler (np.array): (r,p,y) angles Returns: - torch.Tensor: A tensor of shape (..., 3, 3) representing batches of rotation matrices. + np.array: 3x3 rotation matrix + + Raises: + AssertionError: [Invalid input shape] """ - # broadcast array is necessary to use numba parallel mode - q1, q2 = np.broadcast_arrays(quaternion[..., np.newaxis], quaternion[..., np.newaxis, :]) - outer = q1 * q2 - # Extract the necessary components - xx = outer[..., 0, 0] - yy = outer[..., 1, 1] - zz = outer[..., 2, 2] - xy = outer[..., 0, 1] - xz = outer[..., 0, 2] - yz = outer[..., 1, 2] - xw = outer[..., 0, 3] - yw = outer[..., 1, 3] - zw = outer[..., 2, 3] + euler = np.asarray(euler, dtype=np.float64) + assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler) - rotation_matrix = np.empty(quaternion.shape[:-1] + (3, 3), dtype=np.float32) + return R.from_euler("xyz", euler).as_matrix() - rotation_matrix[..., 0, 0] = 1 - 2 * (yy + zz) - rotation_matrix[..., 0, 1] = 2 * (xy - zw) - rotation_matrix[..., 0, 2] = 2 * (xz + yw) - rotation_matrix[..., 1, 0] = 2 * (xy + zw) - rotation_matrix[..., 1, 1] = 1 - 2 * (xx + zz) - rotation_matrix[..., 1, 2] = 2 * (yz - xw) +def mat2euler(rmat): + """ + Converts given rotation matrix to euler angles in radian. - rotation_matrix[..., 2, 0] = 2 * (xz - yw) - rotation_matrix[..., 2, 1] = 2 * (yz + xw) - rotation_matrix[..., 2, 2] = 1 - 2 * (xx + yy) + Args: + rmat (np.array): 3x3 rotation matrix - return rotation_matrix + Returns: + np.array: (r,p,y) converted euler angles in radian vec3 float + """ + M = np.array(rmat, dtype=np.float32, copy=False)[:3, :3] + return R.from_matrix(M).as_euler("xyz") @jit(nopython=True) @@ -613,38 +631,6 @@ def axisangle2quat(vec): return R.from_rotvec(vec).as_quat() -def euler2quat(euler): - """ - Converts euler angles into quaternion form - - Args: - euler (np.array): (r,p,y) angles - - Returns: - np.array: (x,y,z,w) float quaternion angles - - Raises: - AssertionError: [Invalid input shape] - """ - return R.from_euler("xyz", euler).as_quat() - - -def quat2euler(quat): - """ - Converts euler angles into quaternion form - - Args: - quat (np.array): (x,y,z,w) float quaternion angles - - Returns: - np.array: (r,p,y) angles - - Raises: - AssertionError: [Invalid input shape] - """ - return R.from_quat(quat).as_euler("xyz") - - def pose_in_A_to_pose_in_B(pose_A, pose_A_in_B): """ Converts a homogenous matrix corresponding to a point C in frame A @@ -983,64 +969,6 @@ def make_pose(translation, rotation): return pose -def unit_vector(data, axis=None, out=None): - """ - Returns ndarray normalized by length, i.e. eucledian norm, along axis. - - E.g.: - >>> v0 = numpy.random.random(3) - >>> v1 = unit_vector(v0) - >>> numpy.allclose(v1, v0 / numpy.linalg.norm(v0)) - True - - >>> v0 = numpy.random.rand(5, 4, 3) - >>> v1 = unit_vector(v0, axis=-1) - >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=2)), 2) - >>> numpy.allclose(v1, v2) - True - - >>> v1 = unit_vector(v0, axis=1) - >>> v2 = v0 / numpy.expand_dims(numpy.sqrt(numpy.sum(v0*v0, axis=1)), 1) - >>> numpy.allclose(v1, v2) - True - - >>> v1 = numpy.empty((5, 4, 3), dtype=numpy.float32) - >>> unit_vector(v0, axis=1, out=v1) - >>> numpy.allclose(v1, v2) - True - - >>> list(unit_vector([])) - [] - - >>> list(unit_vector([1.0])) - [1.0] - - Args: - data (np.array): data to normalize - axis (None or int): If specified, determines specific axis along data to normalize - out (None or np.array): If specified, will store computation in this variable - - Returns: - None or np.array: If @out is not specified, will return normalized vector. Otherwise, stores the output in @out - """ - if out is None: - data = np.array(data, dtype=np.float32, copy=True) - if data.ndim == 1: - data /= math.sqrt(np.dot(data, data)) - return data - else: - if out is not data: - out[:] = np.array(data, copy=False) - data = out - length = np.atleast_1d(np.sum(data * data, axis)) - np.sqrt(length, length) - if axis is not None: - length = np.expand_dims(length, axis) - data /= length - if out is None: - return data - - def get_orientation_error(target_orn, current_orn): """ Returns the difference between two quaternion orientations as a 3 DOF numpy array. @@ -1169,65 +1097,45 @@ def vecs2quat(vec0, vec1, normalized=False): return quat_unnormalized / np.linalg.norm(quat_unnormalized, axis=-1, keepdims=True) -def l2_distance(v1, v2): - """Returns the L2 distance between vector v1 and v2.""" - return np.linalg.norm(np.array(v1) - np.array(v2)) - +def align_vector_sets(vec_set1, vec_set2): + """ + Computes a single quaternion representing the rotation that best aligns vec_set1 to vec_set2. -def frustum(left, right, bottom, top, znear, zfar): - """Create view frustum matrix.""" - assert right != left - assert bottom != top - assert znear != zfar - - M = np.zeros((4, 4), dtype=np.float32) - M[0, 0] = +2.0 * znear / (right - left) - M[2, 0] = (right + left) / (right - left) - M[1, 1] = +2.0 * znear / (top - bottom) - # TODO: Put this back to 3,1 - # M[3, 1] = (top + bottom) / (top - bottom) - M[2, 1] = (top + bottom) / (top - bottom) - M[2, 2] = -(zfar + znear) / (zfar - znear) - M[3, 2] = -2.0 * znear * zfar / (zfar - znear) - M[2, 3] = -1.0 - return M + Args: + vec_set1 (np.array): (N, 3) tensor of N 3D vectors + vec_set2 (np.array): (N, 3) tensor of N 3D vectors + Returns: + np.array: (4,) Normalized quaternion representing the overall rotation + """ + # Compute the cross-covariance matrix + H = vec_set2.T @ vec_set1 -def ortho(left, right, bottom, top, znear, zfar): - """Create orthonormal projection matrix.""" - assert right != left - assert bottom != top - assert znear != zfar - - M = np.zeros((4, 4), dtype=np.float32) - M[0, 0] = 2.0 / (right - left) - M[1, 1] = 2.0 / (top - bottom) - M[2, 2] = -2.0 / (zfar - znear) - M[3, 0] = -(right + left) / (right - left) - M[3, 1] = -(top + bottom) / (top - bottom) - M[3, 2] = -(zfar + znear) / (zfar - znear) - M[3, 3] = 1.0 - return M + # Compute the elements for the quaternion + trace = H.trace() + w = trace + 1 + x = H[1, 2] - H[2, 1] + y = H[2, 0] - H[0, 2] + z = H[0, 1] - H[1, 0] + # Construct the quaternion + quat = np.stack([x, y, z, w]) -def perspective(fovy, aspect, znear, zfar): - """Create perspective projection matrix.""" - # fovy is in degree - assert znear != zfar - h = np.tan(fovy / 360.0 * np.pi) * znear - w = h * aspect - return frustum(-w, w, -h, h, znear, zfar) + # Handle the case where w is close to zero + if quat[3] < 1e-4: + quat[3] = 0 + max_idx = np.argmax(quat[:3].abs()) + 1 + quat[max_idx] = 1 + # Normalize the quaternion + quat = quat / (np.linalg.norm(quat) + 1e-8) # Add epsilon to avoid division by zero -def anorm(x, axis=None, keepdims=False): - """Compute L2 norms alogn specified axes.""" - return np.linalg.norm(x, axis=axis, keepdims=keepdims) + return quat -def normalize(v, axis=None, eps=1e-10): - """L2 Normalize along specified axes.""" - norm = anorm(v, axis=axis, keepdims=True) - return v / np.where(norm < eps, eps, norm) +def l2_distance(v1, v2): + """Returns the L2 distance between vector v1 and v2.""" + return np.linalg.norm(np.array(v1) - np.array(v2)) def cartesian_to_polar(x, y): @@ -1268,11 +1176,6 @@ def z_angle_from_quat(quat): return np.arctan2(rotated_X_axis[1], rotated_X_axis[0]) -def z_rotation_from_quat(quat): - """Get the quaternion for the rotation around the Z axis produced by the quaternion.""" - return R.from_euler("z", z_angle_from_quat(quat)).as_quat() - - def integer_spiral_coordinates(n): """A function to map integers to 2D coordinates in a spiral pattern around the origin.""" # Map integers from Z to Z^2 in a spiral pattern around the origin. @@ -1285,27 +1188,36 @@ def integer_spiral_coordinates(n): return int(x), int(y) -def calculate_xy_plane_angle(quaternion): +@jit(nopython=True) +def transform_points(points, matrix, translate=True): """ - Compute the 2D orientation angle from a quaternion assuming the initial forward vector is along the x-axis. + Returns points rotated by a homogeneous + transformation matrix. + If points are (n, 2) matrix must be (3, 3) + If points are (n, 3) matrix must be (4, 4) - Parameters: - quaternion : array_like - The quaternion (w, x, y, z) representing the rotation. + Arguments: + points (np.array): (n, dim) where `dim` is 2 or 3. + matrix (np.array): (3, 3) or (4, 4) homogeneous rotation matrix. + translate (bool): whether to apply translation from matrix or not. Returns: - float - The angle (in radians) of the projection of the forward vector onto the XY plane. - Returns 0.0 if the projected vector's magnitude is negligibly small. + np.array: (n, dim) transformed points. """ - fwd = R.from_quat(quaternion).apply([1, 0, 0]) - fwd[2] = 0.0 + if len(points) == 0 or matrix is None: + return points.copy() - if np.linalg.norm(fwd) < 1e-4: - return 0.0 + count, dim = points.shape + # Check if the matrix is close to an identity matrix + identity = np.eye(dim + 1) + if np.abs(matrix - identity[: dim + 1, : dim + 1]).max() < 1e-8: + return points.copy() - fwd /= np.linalg.norm(fwd) - return np.arctan2(fwd[1], fwd[0]) + if translate: + stack = np.ascontiguousarray(np.concatenate((points, np.ones((count, 1))), axis=1)) + return (matrix @ stack.T).T[:, :dim] + else: + return (matrix[:dim, :dim] @ points.T).T @jit(nopython=True) diff --git a/omnigibson/utils/usd_utils.py b/omnigibson/utils/usd_utils.py index 04e7fc651..0c2ff7dd5 100644 --- a/omnigibson/utils/usd_utils.py +++ b/omnigibson/utils/usd_utils.py @@ -15,7 +15,7 @@ import omnigibson.utils.transform_utils as T from omnigibson.macros import gm from omnigibson.utils.backend_utils import _compute_backend as cb -from omnigibson.utils.backend_utils import _ComputeBackend, _ComputeNumpyBackend, _ComputeTorchBackend +from omnigibson.utils.backend_utils import add_compute_function from omnigibson.utils.constants import PRIMITIVE_MESH_TYPES, JointType, PrimType from omnigibson.utils.numpy_utils import vtarray_to_torch from omnigibson.utils.python_utils import assert_valid_key @@ -1118,7 +1118,7 @@ def _get_relative_poses(self, prim_path): self._read_cache["link_transforms"] = cb.from_torch(self._view.get_link_transforms()) idx = self._idx[prim_path] - self._read_cache["relative_poses"][prim_path] = cb.compute_relative_poses( + self._read_cache["relative_poses"][prim_path] = cb.get_custom_method("compute_relative_poses")( idx, len(self._link_idx[idx]), self._read_cache["link_transforms"], @@ -1936,6 +1936,4 @@ def _compute_relative_poses_numpy(idx, n_links, all_tfs, base_pose): # Set these as part of the backend values -setattr(_ComputeBackend, "compute_relative_poses", None) -setattr(_ComputeTorchBackend, "compute_relative_poses", _compute_relative_poses_torch) -setattr(_ComputeNumpyBackend, "compute_relative_poses", _compute_relative_poses_numpy) +add_compute_function(name="compute_relative_poses", np_function=_compute_relative_poses_numpy, th_function=_compute_relative_poses_torch)