四元数和旋转矩阵相互转换

四元数（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