#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File : rotation_utils.py
# Author : Jiayuan Mao
# Email : maojiayuan@gmail.com
# Date : 02/18/2020
#
# This file is part of Project Concepts.
# Distributed under terms of the MIT license.
import numpy as np
from typing import Tuple, TYPE_CHECKING
from concepts.math.rotationlib_wxyz import euler2quat as euler2quat_wxyz, quat2axisangle as quat2axisangle_wxyz
from concepts.utils.typing_utils import Vec3f, Vec4f
if TYPE_CHECKING:
from concepts.simulator.pybullet.client import BulletClient
# For testing whether a number is close to zero
_FLOAT_EPS = np.finfo(np.float64).eps
_EPS4 = _FLOAT_EPS * 4.0
[docs]
def normalize_vector(v: np.ndarray) -> np.ndarray:
"""Normalize a vector."""
return v / np.linalg.norm(v)
[docs]
def find_orthogonal_vector(v: np.ndarray) -> np.ndarray:
"""Find an orthogonal vector to the given vector.
The returned vector is guaranteed to be normalized.
"""
v = v / np.linalg.norm(v, axis=-1, keepdims=True)
mask = np.abs(v[..., 0]) < 0.5
return np.cross(v, np.array([1, 0, 0])) * mask + np.cross(v, np.array([0, 1, 0])) * (1 - mask)
[docs]
def getQuaternionFromMatrix(mat):
""" Convert Rotation Matrix to Quaternion. See rotation.py for notes """
mat = np.asarray(mat, dtype=np.float64)
assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
# Fill only lower half of symmetric matrix
K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
K[..., 0, 0] = Qxx - Qyy - Qzz
K[..., 1, 0] = Qyx + Qxy
K[..., 1, 1] = Qyy - Qxx - Qzz
K[..., 2, 0] = Qzx + Qxz
K[..., 2, 1] = Qzy + Qyz
K[..., 2, 2] = Qzz - Qxx - Qyy
K[..., 3, 0] = Qyz - Qzy
K[..., 3, 1] = Qzx - Qxz
K[..., 3, 2] = Qxy - Qyx
K[..., 3, 3] = Qxx + Qyy + Qzz
K /= 3.0
# TODO: vectorize this -- probably could be made faster
q = np.empty(K.shape[:-2] + (4,))
it = np.nditer(q[..., 0], flags=['multi_index'])
while not it.finished:
# Use Hermitian eigenvectors, values for speed
vals, vecs = np.linalg.eigh(K[it.multi_index])
# Select largest eigenvector, reorder to w,x,y,z quaternion
q[it.multi_index] = vecs[:, np.argmax(vals)]
# Prefer quaternion with positive w
# (q * -1 corresponds to same rotation as q)
if q[it.multi_index][3] < 0:
q[it.multi_index] *= -1
it.iternext()
return q
[docs]
def get_quaternion_from_matrix(mat):
"""Convert a rotation matrix to a quaternion.
The quaternion is represented as a 4-element array in the order of (x, y, z, w),
which follows the convention of pybullet.
"""
mat = np.asarray(mat, dtype=np.float64)
assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
return getQuaternionFromMatrix(mat)
[docs]
def get_quaternion_from_axes(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> np.ndarray:
"""Converts a rotation matrix to a quaternion."""
m = np.stack([x, y, z], axis=1)
return get_quaternion_from_matrix(m)
[docs]
def quat_diff_in_axis_angle(quat1, quat2):
"""Compute the difference between two quaternions (q1 - q2) in axis-angle representation.
Note that since the difference is computed as q1 - q2, for conventional usage of computing
the difference between the source and target, the order should be quat_diff_in_axis_angle(target, source).
"""
quat1 = np.asarray(quat1)
quat2 = np.asarray(quat2)
assert quat1.shape == quat2.shape
assert quat1.shape[-1] == 4
assert quat2.shape[-1] == 4
quat_diff = quat_mul(quat1, quat_conjugate(quat2))
quat_diff = quat_diff[..., [3, 0, 1, 2]] # x, y, z, w -> w, x, y, z
axis, angle = quat2axisangle_wxyz(quat_diff)
return axis * angle
[docs]
def frame_mul(pos_a, quat_a, a_to_b):
"""Multiply a frame with a transformation.
The frame is represented as a tuple of position and quaternion.
The transformation is represented as a tuple of position and quaternion.
The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
"""
pos_a = np.asarray(pos_a, dtype=np.float32)
quat_a = np.asarray(quat_a, dtype=np.float32)
transform_pos = np.asarray(a_to_b[0], dtype=np.float32)
transform_quat = np.asarray(a_to_b[1], dtype=np.float32)
pos_b = pos_a + rotate_vector(transform_pos, quat_a)
quat_b = quat_mul(quat_a, transform_quat)
return pos_b, quat_b
[docs]
def frame_inv(pos_b, quat_b, a_to_b):
"""Inverse a frame with a transformation.
The frame is represented as a tuple of position and quaternion.
The transformation is represented as a tuple of position and quaternion.
The quaternion is represented as a 4-element array in the order of (x, y, z, w), which follows the convention of pybullet.
"""
pos_b = np.asarray(pos_b, dtype=np.float32)
quat_b = np.asarray(quat_b, dtype=np.float32)
transform_pos = np.asarray(a_to_b[0], dtype=np.float32)
transform_quat = np.asarray(a_to_b[1], dtype=np.float32)
inv_transform_quat = quat_conjugate(transform_quat)
quat_a = quat_mul(quat_b, inv_transform_quat)
pos_a = pos_b - rotate_vector(transform_pos, quat_a)
return pos_a, quat_a
[docs]
def quat_mul(q0, q1, *args):
if len(args) > 0:
return quat_mul(q0, quat_mul(q1, *args))
q0 = np.asarray(q0)
q1 = np.asarray(q1)
assert q0.shape == q1.shape
assert q0.shape[-1] == 4
assert q1.shape[-1] == 4
w0 = q0[..., 3]
x0 = q0[..., 0]
y0 = q0[..., 1]
z0 = q0[..., 2]
w1 = q1[..., 3]
x1 = q1[..., 0]
y1 = q1[..., 1]
z1 = q1[..., 2]
w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
y = w0 * y1 + y0 * w1 + z0 * x1 - x0 * z1
z = w0 * z1 + z0 * w1 + x0 * y1 - y0 * x1
q = np.array([x, y, z, w])
if q.ndim == 2:
q = q.swapaxes(0, 1)
assert q.shape == q0.shape
if w0 > 0:
return q
else:
return -q
[docs]
def quat_conjugate(q):
q = np.asarray(q)
assert q.shape[-1] == 4
q = np.array([-q[..., 0], -q[..., 1], -q[..., 2], q[..., 3]])
return q
[docs]
def rpy(r, p, y, degree=True):
"""Create a quaternion from euler angles."""
if degree:
r = np.deg2rad(r)
p = np.deg2rad(p)
y = np.deg2rad(y)
quat_wxyz = euler2quat_wxyz((r, p, y))
return quat_wxyz[..., [1, 2, 3, 0]]
[docs]
def rotate_vector(vec, quat):
"""Rotate a vector by a quaternion."""
vec = np.asarray(vec)
quat = np.asarray(quat)
u = quat[:3]
s = quat[3]
return 2.0 * np.dot(u, vec) * u + (s * s - np.dot(u, u)) * vec + 2.0 * s * np.cross(u, vec)
[docs]
def rotate_vector_batch(batch_vector, quat):
"""Rotate a batch of vectors by a quaternion."""
vec = np.asarray(batch_vector)
quat = np.asarray(quat)
u = quat[:3]
s = quat[3]
return 2.0 * np.dot(vec, u)[..., np.newaxis] * u + (s * s - np.dot(u, u)) * vec + 2.0 * s * np.cross(u, vec)
[docs]
def quaternion_from_vectors(vec1, vec2):
"""Create a rotation quaternion q such that q * vec1 = vec2."""
vec1 = np.asarray(vec1)
vec2 = np.asarray(vec2)
assert vec1.shape == vec2.shape, f'Vector shapes do not match: {vec1.shape} != {vec2.shape}'
assert vec1.shape[-1] == 3, f'Vector shape must be 3, got {vec1.shape[-1]}'
assert vec2.shape[-1] == 3, f'Vector shape must be 3, got {vec2.shape[-1]}'
vec1 = vec1 / np.linalg.norm(vec1, axis=-1, keepdims=True)
vec2 = vec2 / np.linalg.norm(vec2, axis=-1, keepdims=True)
u = np.cross(vec1, vec2)
s = np.dot(vec1, vec2)
if np.linalg.norm(u) < 1e-6:
return np.array([0, 0, 0, 1])
opposite_pairs = (s < -1 + 1e-6)
u = u * (1 - opposite_pairs) + opposite_pairs * find_orthogonal_vector(u)
if len(u.shape) == 1:
s = np.array([s], dtype=u.dtype)
q = np.concatenate([u, s + 1], axis=-1)
q = q / np.linalg.norm(q, axis=-1, keepdims=True)
return q
[docs]
def enumerate_quaternion_from_vectors(input_normal, target_normal, nr_samples: int = 4):
base_quat = quaternion_from_vectors(input_normal, target_normal)
yaw_quat = quaternion_from_vectors(target_normal, np.array([0, 0, 1]))
for yaw in np.arange(0, 2 * np.pi, 2 * np.pi / nr_samples):
quat = quat_mul(
quat_conjugate(yaw_quat),
rpy(0, 0, yaw, degree=False),
yaw_quat,
base_quat
)
yield quat
[docs]
def patch_pybullet():
import pybullet as p
p.getQuaternionFromMatrix = getQuaternionFromMatrix
[docs]
def compute_ee_rotation_mat_from_directions(u: Vec3f, v: Vec3f):
"""Compute the rotation matrix from two directions (the "down" direction for the end effector and the "forward" direction for the end effector).
Args:
u: The "down" direction for the end effector.
v: The "forward" direction for the end effector.
"""
u = np.asarray(u)
u = u / np.linalg.norm(u)
v = np.asarray(v)
v = v - np.dot(u, v) * u
v = v / np.linalg.norm(v)
w = np.cross(u, v)
return np.array([u, v, w]).T
[docs]
def compute_ee_quat_from_directions(u: Vec3f, v: Vec3f, default_quat: Vec4f = (0, 1, 0, 0)):
"""Compute the quaternion from two directions (the "down" direction for the end effector and the "forward" direction for the end effector).
Args:
u: the "down" direction for the end effector.
v: the "forward" direction for the end effector.
default_quat: the default quaternion to be multiplied. This is the quaternion that represents the rotation for the default end effector orientation,
facing downwards and the forward direction is along the x-axis.
"""
mat = compute_ee_rotation_mat_from_directions(u, v)
mat_reference = compute_ee_rotation_mat_from_directions((0, 0, -1), (1, 0, 0))
quat = get_quaternion_from_matrix(np.dot(mat, np.linalg.inv(mat_reference)))
return quat_mul(quat, default_quat)