def r2quat(R): """ Convert a 3x3 rotation matrix to a quaternion. Parameters: R (np.array): 3x3 rotation matrix. Returns: q (np.array): Quaternion in the form [w, x, y, z]. """ R = np.asarray(R, dtype=np.float64) Qxx, Qyx, Qzx = R[..., 0, 0], R[..., 0, 1], R[..., 0, 2] Qxy, Qyy, Qzy = R[..., 1, 0], R[..., 1, 1], R[..., 1, 2] Qxz, Qyz, Qzz = R[..., 2, 0], R[..., 2, 1], R[..., 2, 2] # Fill only lower half of symmetric matrix K = np.zeros(R.shape[:-2] + (4, 4), dtype=np.float64) K[..., 0, 0] = Qxx - Qyy - Qzz K[..., 1, 0] = Qyx + Qxy K[..., 1, 1] = Qyy - Qxx - Qzz K[..., 2, 0] = Qzx + Qxz K[..., 2, 1] = Qzy + Qyz K[..., 2, 2] = Qzz - Qxx - Qyy K[..., 3, 0] = Qyz - Qzy K[..., 3, 1] = Qzx - Qxz K[..., 3, 2] = Qxy - Qyx K[..., 3, 3] = Qxx + Qyy + Qzz K /= 3.0 # TODO: vectorize this -- probably could be made faster q = np.empty(K.shape[:-2] + (4,)) it = np.nditer(q[..., 0], flags=['multi_index']) while not it.finished: # Use Hermitian eigenvectors, values for speed vals, vecs = np.linalg.eigh(K[it.multi_index]) # Select largest eigenvector, reorder to w,x,y,z quaternion q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)] # Prefer quaternion with positive w # (q * -1 corresponds to same rotation as q) if q[it.multi_index][0] < 0: q[it.multi_index] *= -1 it.iternext() return q