docs: def rot1(theta):
      """
      Args:
          theta (float): Angle in radians
      Return:
          Rotation matrix of angle theta around the X-axis
      """
      return np.array([
          [1, 0, 0],
          [0, np.cos(theta), np.sin(theta)],
          [0, -np.sin(theta), np.cos(theta)]
      ])