网站建设问题大全软文素材网

AlphaFold3  rigid_utils 模块的 quat_multiply 和 quat_multiply_by_vec 函数实现了四元数的乘法以及四元数与纯向量四元数的乘法。它使用了一个预定义的四元数乘法表 _QUAT_MULTIPLY 来简化计算。

源代码：

_QUAT_MULTIPLY = np.zeros((4, 4, 4))
_QUAT_MULTIPLY[:, :, 0] = [[ 1, 0, 0, 0],[ 0,-1, 0, 0],[ 0, 0,-1, 0],[ 0, 0, 0,-1]]_QUAT_MULTIPLY[:, :, 1] = [[ 0, 1, 0, 0],[ 1, 0, 0, 0],[ 0, 0, 0, 1],[ 0, 0,-1, 0]]_QUAT_MULTIPLY[:, :, 2] = [[ 0, 0, 1, 0],[ 0, 0, 0,-1],[ 1, 0, 0, 0],[ 0, 1, 0, 0]]_QUAT_MULTIPLY[:, :, 3] = [[ 0, 0, 0, 1],[ 0, 0, 1, 0],[ 0,-1, 0, 0],[ 1, 0, 0, 0]]_QUAT_MULTIPLY_BY_VEC = _QUAT_MULTIPLY[:, 1:, :]_CACHED_QUATS = {"_QTR_MAT": _QTR_MAT,"_QUAT_MULTIPLY": _QUAT_MULTIPLY,"_QUAT_MULTIPLY_BY_VEC": _QUAT_MULTIPLY_BY_VEC
}@lru_cache(maxsize=None)
def _get_quat(quat_key, dtype, device):return torch.tensor(_CACHED_QUATS[quat_key], dtype=dtype, device=device)def quat_multiply(quat1, quat2):"""Multiply a quaternion by another quaternion."""mat = _get_quat("_QUAT_MULTIPLY", dtype=quat1.dtype, device=quat1.device)reshaped_mat = mat.view((1,) * len(quat1.shape[:-1]) + mat.shape)return torch.sum(reshaped_mat *quat1[..., :, None, None] *quat2[..., None, :, None],dim=(-3, -2))def quat_multiply_by_vec(quat, vec):"""Multiply a quaternion by a pure-vector quaternion."