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: