utils.py 文件源码

def rotate_vector(basis, to_rotate, psi):
    """Uses Rodrigues' rotation formula to rotate the given vector v by psi around k.

    If a matrix is given the operation will by applied on the last dimension.

    Args:
        basis: the unit vector defining the rotation axis (k)
        to_rotate: the vector to rotate by the angle psi (v)
        psi: the rotation angle (psi)

    Returns:
        vector: the rotated vector
    """
    cross_product = np.cross(basis, to_rotate)
    dot_product = np.sum(np.multiply(basis, to_rotate), axis=-1)[..., None]
    cos_psi = np.cos(psi)[..., None]
    sin_psi = np.sin(psi)[..., None]
    return to_rotate * cos_psi + cross_product * sin_psi + basis * dot_product * (1 - cos_psi)