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