四元数(Quaternions)和旋转矩阵是在三维空间中表示旋转的两种常见数学工具。
### 四元数和旋转矩阵相互转换
import jax.numpy as jnp
import numpy as np
import math
### 1. 角度和弧度相互转化
## 角度转弧度
degrees = 45
radians = math.radians(degrees)
print(f"{degrees} degrees is equal to {radians} radians")
## 弧度转角度
radians = math.pi / 4
degrees = math.degrees(radians)
print(f"{radians} radians is equal to {degrees} degrees")
### 2. 绕 x,y,z轴旋转一定角度的旋转矩阵
degrees = 45
θ = math.radians(degrees)
# 绕x轴45度旋转的旋转矩阵
Rx = jnp.stack([jnp.array([1, 0, 0]),
jnp.array([0, math.cos(θ), -math.sin(θ)]),
jnp.array([0, math.sin(θ), math.cos(θ)])])
print(f"绕x轴{degrees}度旋转的旋转矩阵为:{Rx}")
# 绕y轴45度旋转的旋转矩阵
Ry = jnp.stack([jnp.array([math.cos(θ), 0, math.sin(θ)]),
jnp.array([0, 1, 0]),
jnp.array([-math.sin(θ), 0, math.cos(θ)])])
print(f"绕y轴{degrees}度旋转的旋转矩阵为:{Ry}")
# 绕z轴45度旋转的旋转矩阵
Rz = jnp.stack([jnp.array([math.cos(θ), -math.sin(θ), 0]),
jnp.array([math.sin(θ), math.cos(θ), 0]),
jnp.array([0, 0, 1])])
print(f"绕z轴{degrees}度旋转的旋转矩阵为:{Rz}")
##### chatGPT代码 #####
### 3. 四元数和旋转矩阵相互转换-代码1
def quaternion_to_rotation_matrix(q):
"""
Convert a quaternion to a 3x3 rotation matrix.
"""
a, b, c, d = q
return np.array([
[a**2 + b**2 - c**2 - d**2, 2 * (b*c - a*d), 2 * (b*d + a*c)],
[2 * (b*c + a*d), a**2 - b**2 + c**2 - d**2, 2 * (c*d - a*b)],
[2 * (b*d - a*c), 2 * (c*d + a*b), a**2 - b**2 - c**2 + d**2]
])
def rotation_matrix_to_quaternion(R):
"""
Convert a 3x3 rotation matrix to a quaternion.
"""
tr = np.trace(R)
if tr > 0:
S = 2 * np.sqrt(1 + tr)
qw = 0.25 * S
qx = (R[2, 1] - R[1, 2]) / S
qy = (R[0, 2] - R[2, 0]) / S
qz = (R[1, 0] - R[0, 1]) / S
elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
S = 2 * np.sqrt(1 + R[0, 0] - R[1, 1] - R[2, 2])
qw = (R[2, 1] - R[1, 2]) / S
qx = 0.25 * S
qy = (R[0, 1] + R[1, 0]) / S
qz = (R[0, 2] + R[2, 0]) / S
elif R[1, 1] > R[2, 2]:
S = 2 * np.sqrt(1 + R[1, 1] - R[0, 0] - R[2, 2])
qw = (R[0, 2] - R[2, 0]) / S
qx = (R[0, 1] + R[1, 0]) / S
qy = 0.25 * S
qz = (R[1, 2] + R[2, 1]) / S
else:
S = 2 * np.sqrt(1 + R[2, 2] - R[0, 0] - R[1, 1])
qw = (R[1, 0] - R[0, 1]) / S
qx = (R[0, 2] + R[2, 0]) / S
qy = (R[1, 2] + R[2, 1]) / S
qz = 0.25 * S
return np.array([qw, qx, qy, qz])
# 相互转化
converted_quaternion = rotation_matrix_to_quaternion(Ry)
rotation_matrix = quaternion_to_rotation_matrix(converted_quaternion)
print("="*30)
print("Original Rotation Matrix:\n ", Ry)
print("Converted Quaternion:\n", converted_quaternion)
print("Converted Rotation Matrix:\n", rotation_matrix)
### 4. 四元数和旋转矩阵相互转换-代码2
## pylint: disable=bad-whitespace 是用于禁用 Pylint 代码检查器中的 "bad-whitespace" 警告的注释。
## Pylint 是一个用于检查 Python 代码风格和潜在错误的工具。
# pylint: disable=bad-whitespace
QUAT_TO_ROT = np.zeros((4, 4, 3, 3), dtype=np.float32)
QUAT_TO_ROT[0, 0] = [[ 1, 0, 0], [ 0, 1, 0], [ 0, 0, 1]] # rr
QUAT_TO_ROT[1, 1] = [[ 1, 0, 0], [ 0,-1, 0], [ 0, 0,-1]] # ii
QUAT_TO_ROT[2, 2] = [[-1, 0, 0], [ 0, 1, 0], [ 0, 0,-1]] # jj
QUAT_TO_ROT[3, 3] = [[-1, 0, 0], [ 0,-1, 0], [ 0, 0, 1]] # kk
QUAT_TO_ROT[1, 2] = [[ 0, 2, 0], [ 2, 0, 0], [ 0, 0, 0]] # ij
QUAT_TO_ROT[1, 3] = [[ 0, 0, 2], [ 0, 0, 0], [ 2, 0, 0]] # ik
QUAT_TO_ROT[2, 3] = [[ 0, 0, 0], [ 0, 0, 2], [ 0, 2, 0]] # jk
QUAT_TO_ROT[0, 1] = [[ 0, 0, 0], [ 0, 0,-2], [ 0, 2, 0]] # ir
QUAT_TO_ROT[0, 2] = [[ 0, 0, 2], [ 0, 0, 0], [-2, 0, 0]] # jr
QUAT_TO_ROT[0, 3] = [[ 0,-2, 0], [ 2, 0, 0], [ 0, 0, 0]] # kr
def rot_to_quat(rot, unstack_inputs=False):
"""Convert rotation matrix to quaternion.
Note that this function calls self_adjoint_eig which is extremely expensive on
the GPU. If at all possible, this function should run on the CPU.
Args:
rot: rotation matrix (see below for format).
unstack_inputs: If true, rotation matrix should be shape (..., 3, 3)
otherwise the rotation matrix should be a list of lists of tensors.
Returns:
Quaternion as (..., 4) tensor.
"""
if unstack_inputs:
rot = [jnp.moveaxis(x, -1, 0) for x in jnp.moveaxis(rot, -2, 0)]
[[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot
# pylint: disable=bad-whitespace
k = [[ xx + yy + zz, zy - yz, xz - zx, yx - xy,],
[ zy - yz, xx - yy - zz, xy + yx, xz + zx,],
[ xz - zx, xy + yx, yy - xx - zz, yz + zy,],
[ yx - xy, xz + zx, yz + zy, zz - xx - yy,]]
# pylint: enable=bad-whitespace
k = (1./3.) * jnp.stack([jnp.stack(x, axis=-1) for x in k],
axis=-2)
# Get eigenvalues in non-decreasing order and associated.
_, qs = jnp.linalg.eigh(k)
return qs[..., -1]
def quat_to_rot(normalized_quat):
"""Convert a normalized quaternion to a rotation matrix."""
rot_tensor = jnp.sum(
np.reshape(QUAT_TO_ROT, (4, 4, 9)) *
normalized_quat[..., :, None, None] *
normalized_quat[..., None, :, None],
axis=(-3, -2))
rot = jnp.moveaxis(rot_tensor, -1, 0) # Unstack.
return [[rot[0], rot[1], rot[2]],
[rot[3], rot[4], rot[5]],
[rot[6], rot[7], rot[8]]]
conv_quaternion = rot_to_quat(Ry)
rot_matrix = quat_to_rot(conv_quaternion)
print("Original Rotation Matrix:\n", Ry)
print("Converted Quaternion:\n", conv_quaternion)
print("Converted Rotation Matrix :\n", np.array(rot_matrix))
### 5. 旋转矩阵和四元数对向量进行旋转变换
def rotate_vector_with_matrix(rotation_matrix, vector):
"""
Rotate a 3D vector using a rotation matrix.
"""
rotated_vector = np.dot(rotation_matrix, vector)
return rotated_vector
def rotate_vector_with_quaternion(quaternion, vector):
"""
Rotate a 3D vector using a quaternion.
"""
quaternion_matrix = np.array([
[1 - 2*(quaternion[2]**2 + quaternion[3]**2), 2*(quaternion[1]*quaternion[2] - quaternion[0]*quaternion[3]), 2*(quaternion[1]*quaternion[3] + quaternion[0]*quaternion[2])],
[2*(quaternion[1]*quaternion[2] + quaternion[0]*quaternion[3]), 1 - 2*(quaternion[1]**2 + quaternion[3]**2), 2*(quaternion[2]*quaternion[3] - quaternion[0]*quaternion[1])],
[2*(quaternion[1]*quaternion[3] - quaternion[0]*quaternion[2]), 2*(quaternion[2]*quaternion[3] + quaternion[0]*quaternion[1]), 1 - 2*(quaternion[1]**2 + quaternion[2]**2)]
])
rotated_vector = np.dot(quaternion_matrix, vector)
return rotated_vector
test_vec = np.array([2, 5, 8])
rotated_vector_matrix = rotate_vector_with_matrix(rot_matrix, test_vec)
rotated_vector_quaternion = rotate_vector_with_quaternion(conv_quaternion, test_vec)
print("Original Vector:", test_vec)
print("Rotated Vector with Rotation Matrix:", rotated_vector_matrix)
print("Rotated Vector with Quaternion:", rotated_vector_quaternion)
## 两种计算旋转方法结果相同。
参考:
https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation