PSP - 基于开源框架 OpenFold Multimer 蛋白质复合物的结构预测与BugFix

欢迎关注我的CSDN：https://spike.blog.csdn.net/
本文地址：https://spike.blog.csdn.net/article/details/132410296

Multimer

AlphaFold2-Multimer 是一个基于 AlphaFold2 的神经网络模型，可以预测多链蛋白复合物的结构。该模型在训练和推理时都可以处理多链输入，并且考虑了链之间的对称性和遗传信息。

对于 AlphaFold2 的损失函数、特征编码、裁剪策略和模型架构进行了多项修改，以适应多链蛋白复合物的特点。该模型还提供了一个基于预测 TM-score 的置信度评估方法。
在两个数据集上进行了评估，一个是 Benchmark 2，包含 17 个低同源性的异二聚体；另一个是 Recent-PDB-Multimers，包含 4,433 个最近的蛋白复合物。该模型使用 DockQ 分数来衡量预测结构与真实结构之间的接触质量。
在Benchmark 2上显著优于其他基于 AlphaFold2 或 ClusPro 的方法，在 Recent-PDB-Multimers 上也表现出较大的提升，尤其是在异构二聚体上。该模型还能够准确地预测自身的置信度，并且给出一些高质量的结构示例。

其中 OpenFold 是 AlphaFold2-Multimer 的开源实现，即：

Paper: OpenFold: Retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization
GitHub: https://github.com/aqlaboratory/openfold

将 OpenFold 的分支切换至 multimer 分支，即可使用 Multimer 功能，目前是 Debug 版本，基本推理功能已经支持，MSA 部分支持使用 AF2 的推理结果，模型支持 DeepMind 提供的 Multimer v3 模型，其余使用 OpenFold 的相关源码。评估当前 OpenFold Multimer框架的有效性。

其他参考文章：

开源可训练的蛋白质结构预测框架 OpenFold 的环境配置
基于 OpenFold 训练的 Finetuning 模型与推理逻辑评估

1. 模型效果

测试序列是 H1106_A122_B114.fasta，来源于 CASP15，即：

>A
MSRIITAPHIGIEKLSAISLEELSCGLPDRYALPPDGHPVEPHLERLYPTAQSKRSLWDFASPGYTFHGLHRAQDYRRELDTLQSLLTTSQSSELQAAAALLKCQQDDDRLLQIILNLLHKV
>B
MNITLTKRQQEFLLLNGWLQLQCGHAERACILLDALLTLNPEHLAGRRCRLVALLNNNQGERAEKEAQWLISHDPLQAGNWLCLSRAQQLNGDLDKARHAYQHYLELKDHNESP

OpenFold Multimer 的 MSA 文件夹格式，与 Monomer 类似，位于 alignments 文件夹中，不同的链放入同名文件夹中，即文件夹 A 和 B，具体文件如下：

bfd_uniref_hits.a3m
mgnify_hits.sto
pdb_hits.sto
uniprot_hits.sto
uniref90_hits.sto

其中 bfd_uniref_hits.a3m、mgnify_hits.sto、uniref90_hits.sto 是 MSA 的搜索结果，uniprot_hits.sto 用于 MSA Pairing，pdb_hits.sto 是模版搜索的结果。

测试命令，如下：

因为使用已有的 AlphaFold2 Multimer 搜索的 MSA，因此 MSA 相关配置并未启用；
模型使用 AF2 的 params_model_1_multimer_v3.npz，配置使用 model_1_multimer_v3。

即：

python3 run_pretrained_openfold.py \
mydata/test-multimer \
af2-data-v230/pdb_mmcif/mmcif_files \
--uniref90_database_path af2-data-v230/uniref90/uniref90.fasta \
--mgnify_database_path af2-data-v230/mgnify/mgy_clusters_2022_05.fa \
--pdb70_database_path af2-data-v230/pdb70/pdb70 \
--uniclust30_database_path deepmsa2/uniclust30/uniclust30_2018_08 \
--uniref30_database_path af2-data-v230/uniref30/UniRef30_2021_03 \
--uniprot_database_path af2-data-v230/uniprot/uniprot.fasta \
--pdb_seqres_database_path af2-data-v230/pdb_seqres/pdb_seqres.txt \
--output_dir mydata/outputs-multimer/H1106_A122_B114/ \
--bfd_database_path af2-data-v230/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
--model_device "cuda:0" \
--jackhmmer_binary_path /opt/openfold/hhsuite-speed/jackhmmer \
--hmmsearch_binary_path /opt/openfold/hhsuite-speed/hmmsearch \
--hhblits_binary_path /opt/conda/envs/openfold/bin/hhblits \
--hhsearch_binary_path /opt/conda/envs/openfold/bin/hhsearch \
--kalign_binary_path /opt/conda/envs/openfold/bin/kalign \
--config_preset "model_1_multimer_v3" \
--jax_param_path af2-data-v230/params/params_model_1_multimer_v3.npz \
--max_template_date 2022-04-01

运行日志如下，整体推理速度较快：

INFO:openfold/utils/script_utils.py:Successfully loaded JAX parameters at af2-data-v230/params/params_model_1_multimer_v3.npz...
INFO:run_pretrained_openfold.py:Using precomputed alignments for A at mydata/outputs-multimer/H1106_A122_B114/alignments...
INFO:run_pretrained_openfold.py:Using precomputed alignments for B at mydata/outputs-multimer/H1106_A122_B114/alignments...
INFO:openfold/utils/script_utils.py:Running inference for A-B...
INFO:openfold/utils/script_utils.py:Inference time: 44.876936707645655
INFO:run_pretrained_openfold.py:Output written to mydata/outputs-multimer/H1106_A122_B114/predictions/A-B_model_1_multimer_v3_unrelaxed.pdb...
INFO:run_pretrained_openfold.py:Running relaxation on mydata/outputs-multimer/H1106_A122_B114/predictions/A-B_model_1_multimer_v3_unrelaxed.pdb...
INFO:openfold/utils/script_utils.py:Relaxation time: 26.89977646060288
INFO:openfold/utils/script_utils.py:Relaxed output written to mydata/outputs-multimer/H1106_A122_B114/predictions/A-B_model_1_multimer_v3_relaxed.pdb...

与 AlphaFold2 Multimer 的预测结果 unrelaxed_model_1_multimer_v3_pred_0.pdb，作为对比，效果在 H1106_A122_B114 中，略有提升，即：

[Info] {'TMScore': 0.8824, 'RMSD(local)': 1.92, 'Align.Len.': 173, 'DockQ': 0.613}
[Info] {'TMScore': 0.8803, 'RMSD(local)': 2.12, 'Align.Len.': 174, 'DockQ': 0.600}

其中，黄色是 Reference，蓝色是 AlphaFold2 Multimer 的预测结果，粉色是 OpenFold Multimer 的预测结果，如下：

Multimer

2. Bugfix

Bug: 在MSA 序列 (sequence) 中，存在无法解析的 "." 关键字，导致 KeyError，即：

Traceback (most recent call last):
  File "run_pretrained_openfold.py", line 477, in <module>
    main(args)
  File "run_pretrained_openfold.py", line 291, in main
    feature_dict = generate_feature_dict(
  File "run_pretrained_openfold.py", line 134, in generate_feature_dict
    feature_dict = data_processor.process_fasta(
  File "openfold/data/data_pipeline.py", line 1167, in process_fasta
    chain_features = self._process_single_chain(
  File "openfold/data/data_pipeline.py", line 1116, in _process_single_chain
    chain_features = self._monomer_data_pipeline.process_fasta(
  File "openfold/data/data_pipeline.py", line 860, in process_fasta
    msa_features = self._process_msa_feats(alignment_dir, input_sequence, alignment_index)
  File "openfold/data/data_pipeline.py", line 818, in _process_msa_feats
    msa_features = make_msa_features(
  File "openfold/data/data_pipeline.py", line 232, in make_msa_features
    [residue_constants.HHBLITS_AA_TO_ID[res] for res in sequence]
  File "openfold/data/data_pipeline.py", line 232, in <listcomp>
    [residue_constants.HHBLITS_AA_TO_ID[res] for res in sequence]
KeyError: '.'

源码位于 openfold/data/data_pipeline.py 中，即：

def _process_msa_feats(
    self,
    alignment_dir: str,
    input_sequence: Optional[str] = None,
    alignment_index: Optional[str] = None
) -> Mapping[str, Any]:
    msas = self._get_msas(
        alignment_dir, input_sequence, alignment_index
    )
    msa_features = make_msa_features(
        msas=msas
    )
    return msa_features

定义日志 logger，即：

import logging
logging.basicConfig()
logger = logging.getLogger(__file__)
logger.setLevel(level=logging.INFO)

定位 sequence，来源于 pdb_hits.sto 模版搜索结果，即：

INFO:openfold/data/data_pipeline.py:[CL] Error sequence: .MALLPDGQSI.EPHISR...LY...P....ERL.....ADRALLDFATPHR..GFHDLLRP.VD..FHQAMQ...G.LRSV.LAE.....GQSPELRAAA..ILLEQM.HADEQLMQMTLHLLHKV

原因：在 Multimer 中，Template 的搜索结果是 pdb_hits.sto，误解析成 MSA 文件，排除即可，同时，增加 pdb_hits.sto 的解析函数。

相关代码，各有 2 处，都需要修改，之前验证的是hmm_output，现修改成pdb_hits，即：

# ...
elif ext == ".sto" and "pdb_hits" not in filename:
    msa = parsers.parse_stockholm(read_msa(start, size))
# ...
elif name == "pdb_hits.sto":
    hits = parsers.parse_hmmsearch_sto(
        read_template(start, size),
        input_sequence,
    )
    all_hits[name] = hits
# ...