AF3 pair_sequences函数解读

AlphaFold3 msa_pairing模块的pair_sequences函数的核心目标是基于 MSA（多序列比对）中的物种信息，在多条链之间建立 MSA 配对索引，从而帮助 AlphaFold3 捕捉共进化信息，提升蛋白复合物预测的准确性。函数pair_sequences 通过调用 _make_msa_df、  _create_species_dict  以及_match_rows_by_sequence_similarity实现其功能。

源代码：

def pair_sequences(
        examples: List[Mapping[str, np.ndarray]],
) -> Dict[int, np.ndarray]:
    """Returns indices for paired MSA sequences across chains."""

    num_examples = len(examples)

    all_chain_species_dict = []
    common_species = set()
    for chain_features in examples:
        msa_df = _make_msa_df(chain_features)
        species_dict = _create_species_dict(msa_df)
        all_chain_species_dict.append(species_dict)
        common_species.update(set(species_dict))

    common_species = sorted(common_species)
    common_species.remove(b'')  # Remove target sequence species.

    all_paired_msa_rows = [np.zeros(len(examples), int)]
    all_paired_msa_rows_dict = {k: [] for k in range(num_examples)}
    all_paired_msa_rows_dict[num_examples] = [np.zeros(len(examples), int)]

    for species in common_species:
        if not species:
            continue
        this_species_msa_dfs = []
        species_dfs_present = 0
        for species_dict in all_chain_species_dict:
            if species in species_dict:
                this_species_msa_dfs.append(species_dict[species])
                species_dfs_present += 1
            else:
                this_species_msa_dfs.append(None)

        # Skip species that are present in only one chain.
        if species_dfs_present <= 1:
            continue

        if np.any(
                np.array([len(species_df) for species_df in
                          this_species_msa_dfs if
                          isinstance(species_df, pd.DataFrame)]) > 600):
            continue

        paired_msa_rows = _match_rows_by_sequence_similarity(this_species_msa_dfs)
        all_paired_msa_rows.extend(paired_msa_rows)
        all_paired_msa_rows_dict[species_dfs_present].extend(paired_msa_rows)
    all_paired_msa_rows_dict = {
        num_examples: np.array(paired_msa_rows) for
        num_examples, paired_msa_rows in all_paired_msa_rows_dict.items()
    }
    return all_paired_msa_rows_dict

代码解读：

函数输入

def pair_sequences(examples: List[Mapping[str, np.ndarray]]) -> Dict[int, np.ndarray]

examples：包含多条链的 MSA 信息，每个元素是一个字典，代表一条蛋白链的 MSA 相关特征。

• 例如，examples[0] 可能对应链 A 的 MSA 特征，examples[1] 可能对应链 B 的 MSA 特征。

代码执行过程

1️⃣ 提取所有链的 MSA 并构建物种索引

num_examples = len(examples)

all_chain_species_dict = []
common_species = set()
for chain_features in examples:
    msa_df = _make_msa_df(chain_features)
    species_dict = _create_species_dict(msa_df)
    all_chain_species_dict.append(species_dict)
    common_species.update(set(species_dict))

• 遍历 examples：
  • 使用 _make_msa_df 将每条链的 MSA 信息转换为 DataFrame（包含 msa_similarity、msa_row、msa_species_identifiers）。
  • 使用 _create_species_dict 按物种对 MSA 进行分组，得到 species_dict。
  • 将 species_dict 存入 all_chain_species_dict，并收集所有物种信息，存入 common_species。

2️⃣ 筛选出所有链共有的物种

common_species = sorted(common_species)
common_species.remove(b'')  # 移除目标序列所在的物种

• 由于 MSA 可能包含多个物种的序列，该步骤确保只保留所有链都出现过的物种。
• 移除 b''，它代表目标序列自身的物种，不需要配对。

3️⃣ 初始化存储配对 MSA 结果的字典

all_paired_msa_rows = [np