AlphaFold3  data_transforms 模块的 make_msa_mask 函数  主要是 为 MSA (多序列比对) 生成初始化的 mask 掩码特征： msa_mask和 msa_row_mask。用于后续的特征处理，确保模型在处理 MSA 时能够识别哪些数据是有效的，哪些需要被忽略（比如填充的 0）。

源代码：

def make_msa_mask(protein):
    """Mask features are all ones, but will later be zero-padded."""
    protein["msa_mask"] = torch.ones(protein["msa"].shape, dtype=torch.float32)
    protein["msa_row_mask"] = torch.ones(
        (protein["msa"].shape[0]), dtype=torch.float32
    )
    return protein复制

源码解读：

protein["msa_mask"] = torch.ones(protein["msa"].shape, dtype=torch.float32)
复制

• protein["msa"] 是一个 (num_seq, num_res) 形状的张量：
  • num_seq：MSA 的序列数。
  • num_res：每条 MSA 序列的氨基酸个数。
• torch.ones(protein["msa"].shape, dtype=torch.float32) 生成一个全 1 的掩码：
  • 形状：(num_seq, num_res)
  • 作用：指示哪些 MSA 位置是有效的。
  • 后续可能的变化：
    • 在 MSA 被 zero-padding（填充 0） 之后，填充部分的 msa_mask 可能会被置为 0。

protein["msa_row_mask"] = torch.ones(
    (protein["msa"].shape[0]), dtype=torch.float32
)
复制

• protein["msa_row_mask"] 是一个 (num_seq,) 形状的 1D 张量：
  • 作用：指示哪些 MSA 序列是有效的。
  • 后续可能的变化：
    • 在 MSA 被 zero-padding 或者裁剪 之后，可能会有 0 置入其中。

代码意义

目的

• msa_mask：用于 逐个氨基酸 位置的掩码，确保后续计算时忽略填充区域。
• msa_row_mask：用于 整条 MSA 序列 的掩码，可能用于批处理时过滤无效序列。

作用

• 确保 MSA 计算时不会受 填充值 (zero-padding) 影响。
• 在后续的 extra_msa_mask 计算、聚类、特征提取时，利用 msa_mask 来过滤无效数据。
• 在 DataPipeline 处理过程中，MSA 会被填充到固定长度，msa_mask 可以帮助 标记填充区域。

代码关系

变量	形状	作用
msa_mask	(num_seq, num_res)	逐个氨基酸位置的掩码，标记 MSA 位置是否有效
msa_row_mask	(num_seq,)	逐行掩码，标记 MSA 序列是否有效

最终，make_msa_mask 主要是 初始化 MSA 掩码，确保后续处理可以正确忽略填充区域，在 AlphaFold3 处理中起到 屏蔽无效数据 的作用。