AlphaFold3 superimpose函数通过使用SVD最小化RMSD，将坐标叠加到参考上，在蛋白质结构预测中用于比较预测结构与真实结构的相似性。
源代码：
from src.utils.geometry.alignment import weighted_rigid_align
from src.utils.geometry.vector import Vec3Array
import torchdef compute_rmsd(tensor1, tensor2, mask, eps=1e-6):"""Compute the RMSD between two tensors."""diff = tensor1 - tensor2squared_diff = diff ** 2sum_squared_diff = squared_diff.sum(dim=-1)# Mask out invalid positionssum_squared_diff = sum_squared_diff * mask# Average over valid positionsdenom = mask.sum(dim=-1) + epsmean_squared_diff = torch.sum(sum_squared_diff, dim=-1) / denom# Square root to get RMSDrmsd = torch.sqrt(mean_squared_diff + eps)return rmsddef superimpose(reference, coords, mask):"""Superimposes coordinates onto a reference by minimizing RMSD using SVD.Args:reference:[*, N, 3] reference tensorcoords:[*, N, 3] tensormask:[*, N] tensorReturns:A tuple of [*, N, 3] superimposed coords and [*] final RMSDs."""# To Vec3Array for alignmentreference = Vec3Array.from_array(reference)coords = Vec3Array.from_array(coords)# Align the coordinates to the referencealigned_coords = weighted_rigid_align(coords, reference, weights=mask, mask=mask)aligned_coords = aligned_coords.to_tensor()reference = reference.to_tensor()# Compute RMSDrmsds = compute_rmsd(reference, aligned_coords)return aligned_coords, rmsds
代码解析
1. compute_rmsd 函数
功能： 计算两组坐标 tensor1 和 tensor2 之间的 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（奇异值分解）寻找最优旋转和平移，将coords与reference对齐。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 值
代码核心逻辑
compute_rmsd- 计算两组坐标的 RMSD（用于衡量预测误差）。
superimpose- 通过 SVD 进行刚性对齐，使预测结构尽量匹配真实结构。
- 计算对齐后的 RMSD，评估对齐质量。
这两个函数主要用于蛋白质结构比对，常用于 AlphaFold3 这种结构预测模型的评估。
