Pythonで回転行列

回転行列をnumpyで実装します．回転行列については以下の記事を参照してください．

2次元

Python (numpy)

import numpy as np

def rot(vec, theta):
    vec = vec.reshape(-1, 1)
    cos, sin = np.cos(theta), np.sin(theta)
    R = np.array([
        [cos, - sin],
        [sin, cos]
    ])
    return R @ vec

実行例：

import matplotlib.pyplot as plt

theta = np.pi / 4
v = np.array([
    [1],
    [0]
])

Rv = rot(v, theta)

fig = plt.figure(figsize = (5, 5), tight_layout = True)
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)
ax.grid()
ax.quiver(0, 0, *v,
            angles = 'xy', scale_units = 'xy', scale = 1)
ax.text(*v, f'v=\n{v}')
ax.quiver(0, 0, *Rv,
            angles = 'xy', scale_units = 'xy', scale = 1)
ax.text(*Rv, f'Rv=\n{Rv}')
plt.show()

Pythonで回転行列