AF3 unify_template_features 函数解读

AlphaFold3 data_pipeline 模块的 unify_template_features 函数用于将多条链的模板特征整合为一个统一的 FeatureDict，以适应对多链复合物的处理。每条链的模板特征经过索引偏移处理后，拼接为一个完整的模板特征矩阵。

该方法的核心在于：

1. 序列对齐：根据每条链的长度，将模板特征填充到一个统一的序列矩阵中。
2. 链索引标记：为每条链的模板特征添加 template_chain_index，标识模板来源的链。
3. 特征拼接：将所有链的模板特征按模板数量拼接，并生成最终的特征字典。

源代码：

def unify_template_features(template_feature_list: Sequence[FeatureDict]
) -> FeatureDict:out_dicts = []seq_lens = [fd["template_aatype"].shape[1] for fd in template_feature_list]for i, fd in enumerate(template_feature_list):out_dict = {}n_templates, n_res = fd["template_aatype"].shape[:2]for k, v in fd.items():seq_keys = ["template_aatype","template_all_atom_positions","template_all_atom_mask",]if k in seq_keys:new_shape = list(v.shape)assert (new_shape[1] == n_res)new_shape[1] = sum(seq_lens)new_array = np.zeros(new_shape, dtype=v.dtype)if k == "template_aatype":new_array[..., residue_constants.HHBLITS_AA_TO_ID['-']] = 1offset = sum(seq_lens[:i])new_array[:, offset:offset + seq_lens[i]] = vout_dict[k] = new_arrayelse:out_dict[k] = vchain_indices = np.array(n_templates * [i])out_dict["template_chain_index"] = chain_indicesif n_templates != 0:out_dicts.append(out