AF3 superimpose函数解读

AlphaFold3 superimpose函数通过使用SVD最小化RMSD,将坐标叠加到参考上,在蛋白质结构预测中用于比较预测结构与真实结构的相似性。

源代码:

复制代码
from src.utils.geometry.alignment import weighted_rigid_align
from src.utils.geometry.vector import Vec3Array
import torch


def compute_rmsd(tensor1, tensor2, mask, eps=1e-6):
    """Compute the RMSD between two tensors."""
    diff = tensor1 - tensor2
    squared_diff = diff ** 2
    sum_squared_diff = squared_diff.sum(dim=-1)

    # Mask out invalid positions
    sum_squared_diff = sum_squared_diff * mask

    # Average over valid positions
    denom = mask.sum(dim=-1) + eps
    mean_squared_diff = torch.sum(sum_squared_diff, dim=-1) / denom

    # Square root to get RMSD
    rmsd = torch.sqrt(mean_squared_diff + eps)
    return rmsd


def superimpose(reference, coords, mask):
    """
        Superimposes coordinates onto a reference by minimizing RMSD using SVD.

        Args:
            reference:
                [*, N, 3] reference tensor
            coords:
                [*, N, 3] tensor
            mask:
                [*, N] tensor
        Returns:
            A tuple of [*, N, 3] superimposed coords and [*] final RMSDs.
    """
    # To Vec3Array for alignment
    reference = Vec3Array.from_array(reference)
    coords = Vec3Array.from_array(coords)

    # Align the coordinates to the reference
    aligned_coords = weighted_rigid_align(coords, reference, weights=mask, mask=mask)
    aligned_coords = aligned_coords.to_tensor()
    reference = reference.to_tensor()
    # Compute RMSD
    rmsds = compute_rmsd(reference, aligned_coords)
    return aligned_coords, rmsds

代码解析

1. compute_rmsd 函数

功能: 计算两组坐标 tensor1tensor2 之间的 RMSD,常用于比较蛋白质结构的相似性。

复制代码
def compute_rmsd(tensor1, tensor2, mask, eps=1e-6):
  • tensor1: 真实坐标 (ground truth) [*, N, 3]
  • tensor2: 预测坐标 [*, N, 3]
  • mask: 指定哪些坐标是有效的 [*, N]
  • eps: 避免除零问题的小常数

步骤解析

复制代码
diff = tensor1 - tensor2  # 计算两者的差值
squared_diff = diff ** 2  # 差值平方
sum_squared_diff = squared_diff.sum(dim=-1)  # 在最后一个维度 (3D) 上求和,即每个点的平方误差和

计算的是欧几里得距离的平方,形状 [*, N]

复制代码
sum_squared_diff = sum_squared_diff * mask  # 只保留mask中有效的坐标

屏蔽无效坐标(未对齐的或者不感兴趣的原子)。

复制代码
denom = mask.sum(dim=-1) + eps  # 计算有效坐标数
mean_squared_diff = torch.sum(sum_squared_diff, dim=-1) / denom  # 计算均方误差

有效原子位置的均方误差(MSE)。

复制代码
rmsd = torch.sqrt(mean_squared_diff + eps)  # 开方得到 RMSD

最终 RMSD 计算完成,返回的是 [ * ] 形状的张量,即每个样本的 RMSD 值。


2. superimpose 函数

功能:

使用加权刚性对齐 (Weighted Rigid Alignment) 方法,通过奇异值分解 (SVD) 使预测坐标 coords 尽可能匹配 reference,然后计算 RMSD。

复制代码
def superimpose(reference, coords, mask):
  • reference: 真实的结构坐标 [*, N, 3]
  • coords: 需要对齐的预测坐标 [*, N, 3]
  • mask: 有效位置的掩码 [*, N]

步骤解析

复制代码
reference = Vec3Array.from_array(reference)
coords = Vec3Array.from_array(coords)

这里 Vec3Array 是一个三维向量的封装类,使得操作更符合几何向量计算需求。

复制代码
aligned_coords = weighted_rigid_align(coords, reference, weights=mask, mask=mask)

核心部分

  • weighted_rigid_align 是一个加权刚性对齐函数,使用SVD(奇异值分解)寻找最优旋转和平移 ,将 coordsreference 对齐。

  • weights=mask 代表对对齐过程加权,仅使用有效坐标对齐。

    aligned_coords = aligned_coords.to_tensor()
    reference = reference.to_tensor()

Vec3Array 转回 PyTorch 张量,方便计算 RMSD。

复制代码
rmsds = compute_rmsd(reference, aligned_coords)
  • 计算对齐后的 RMSD,衡量对齐质量。

最终返回:

  • aligned_coords: 对齐后的坐标
  • rmsds: 对齐后的 RMSD 值

代码核心逻辑

  1. compute_rmsd
    • 计算两组坐标的 RMSD(用于衡量预测误差)。
  2. superimpose
    • 通过 SVD 进行刚性对齐,使预测结构尽量匹配真实结构。
    • 计算对齐后的 RMSD,评估对齐质量。

这两个函数主要用于蛋白质结构比对,常用于 AlphaFold3 这种结构预测模型的评估。

相关推荐
Learn Beyond Limits43 分钟前
Transfer Learning|迁移学习
人工智能·python·深度学习·神经网络·机器学习·ai·吴恩达
程序员三明治44 分钟前
三、神经网络
人工智能·深度学习·神经网络
hundaxxx2 小时前
自演化大语言模型的技术背景
人工智能
数智顾问2 小时前
【73页PPT】美的简单高效的管理逻辑(附下载方式)
大数据·人工智能·产品运营
love530love2 小时前
【保姆级教程】阿里 Wan2.1-T2V-14B 模型本地部署全流程:从环境配置到视频生成(附避坑指南)
人工智能·windows·python·开源·大模型·github·音视频
木头左2 小时前
结合机器学习的Backtrader跨市场交易策略研究
人工智能·机器学习·kotlin
Coovally AI模型快速验证3 小时前
3D目标跟踪重磅突破!TrackAny3D实现「类别无关」统一建模,多项SOTA达成!
人工智能·yolo·机器学习·3d·目标跟踪·无人机·cocos2d
研梦非凡3 小时前
CVPR 2025|基于粗略边界框监督的3D实例分割
人工智能·计算机网络·计算机视觉·3d
MiaoChuAI3 小时前
秒出PPT vs 豆包AI PPT:实测哪款更好用?
人工智能·powerpoint
fsnine3 小时前
深度学习——残差神经网路
人工智能·深度学习