AF3 drmsd函数解读

drmsd（distance Root Mean Square Deviation，距离均方根偏差）函数在AlphaFold3的 src.utils.validation_metrics模块中定义，用于计算两个蛋白质结构（或其他分子结构）之间的距离偏差。它衡量了两个结构的 成对原子间距离 差异，而不是直接比较原子坐标。这种度量方式比 RMSD（Root Mean Square Deviation，均方根偏差）更能反映全局结构差异，因为它不会受到全局对齐的影响。

源代码：

import torch
from typing import Optional


def drmsd(structure_1: torch.Tensor, structure_2: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """
    Calculate the distance Root Mean Square Deviation (dRMSD) between two structures.

    Args:
        structure_1 (torch.Tensor): First structure of shape [..., N, 3]
        structure_2 (torch.Tensor): Second structure of shape [..., N, 3]
        mask (Optional[torch.Tensor]): Mask of shape [..., N] indicating valid positions

    Returns:
        torch.Tensor: The dRMSD between the two structures
    """
    def pairwise_distances(structure: torch.Tensor) -> torch.Tensor:
        """Calculate pairwise distances within a structure."""
        diff = structure[..., :, None, :] - structure[..., None, :, :]
        return torch.norm(diff, dim=-1)

    d1 = pairwise_dis