generate rotation matrices for each 3D axis given an angle

get_rotation_mats.py

import numpy as np
def get_rotation_mat(theta=np.pi, axis=0):
  # https://en.wikipedia.org/wiki/Rotation_matrix

  sin = np.sin(theta)
  cos = np.cos(theta)

  if axis==0:
    R = np.asarray([[1,   0,    0],
                    [0, cos, -sin],
                    [0, sin,  cos]])
  elif axis==1:
    R = np.asarray([[cos,  0, sin],
                    [0,    1,   0],
                    [-sin, 0, cos]])
  elif axis==2:
    R = np.asarray([[cos, -sin, 0],
                    [sin,  cos, 0],
                    [0,      0, 1]])
  else:
    raise ValueError("axis should be 0, 1 or 2!")

  return R