PSP - AlphaFold2 适配不同来源搜索的 MSA 接口

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

MSA (Multiple Sequence Alignment) 在 AlphaFold2 中的工作方式如下：

使用搜索工具 (hhblits/hhsearch/jackhmmer)，从大型数据库中，搜索与目标蛋白质序列相关的蛋白质序列，与目标蛋白质序列对齐。
将对齐的序列作为其神经网络模型的输入特征，该模型由两个主要部分组成：一个表示MSA的行和列，另一个表示蛋白质模型中每个氨基酸之间的原子间距离。
使用神经网络模型预测蛋白质的三维结构，通过反复细化MSA和原子间距离的表示，直到收敛。

当使用 AlphaFold2 进行蛋白结构预测时，有时，比较复杂的序列，需要优化 MSA 搜索，再进行预测蛋白结构，即需要将 MSA 与结构推理两个部分解耦。

1. 修改运行脚本

即 run_alphafold.sh，增加运行参数 -s true/false，即可，修改如下：

...
echo "-s <use_saved_msa>    Use saved msas in the msa directory of output dir, i.e {output_dir}/{seq_name}/msas/*.a3m."
...
while getopts ":d:o:f:t:g:r:e:n:a:m:c:p:l:b:s:" i; do  # 注意 s 之后的冒号":"
...s)use_saved_msa=$OPTARG;;
...
if [[ "$use_saved_msa" == "" ]] ; then  # 默认设置use_saved_msa="false"
fi
...
command_args="--fasta_paths=$fasta_path --output_dir=$output_dir --max_template_date=$max_template_date --db_preset=$db_preset --model_preset=$model_preset --benchmark=$benchmark --use_precomputed_msas=$use_precomputed_msas --num_multimer_predictions_per_model=$num_multimer_predictions_per_model --use_gpu_relax=$use_gpu_relax --logtostderr --use_saved_msa=$use_saved_msa"  # 命令参数
...

2. 修改运行文件

接着，run_alphafold.sh 调用 run_alphafold.py，在flags中，增加参数use_saved_msa，修改如下：

...
flags.DEFINE_boolean('use_saved_msa', None, 'Use saved msas in the msa directory ''of output dir, i.e {output_dir}/{seq_name}/msas/*.a3m.')
...
if __name__ == '__main__':flags.mark_flags_as_required(['fasta_paths','output_dir','data_dir','uniref90_database_path','mgnify_database_path','template_mmcif_dir','max_template_date','obsolete_pdbs_path','use_gpu_relax','use_saved_msa',  # 添加FLAGS支持])app.run(main)
...predict_structure(fasta_path=fasta_path,fasta_name=fasta_name,output_dir_base=FLAGS.output_dir,data_pipeline=data_pipeline,model_runners=model_runners,amber_relaxer=amber_relaxer,benchmark=FLAGS.benchmark,random_seed=random_seed,models_to_relax=FLAGS.models_to_relax,use_saved_msa=FLAGS.use_saved_msa)  # 由 predict_structure 调用
...

接着，修改run_alphafold.py#predict_structure()函数，不同 MSA 设置，使用不同的处理(data_pipeline.process)函数，修改如下：

def predict_structure(fasta_path: str,fasta_name: str,output_dir_base: str,data_pipeline: Union[pipeline.DataPipeline, pipeline_multimer.DataPipeline],model_runners: Dict[str, model.RunModel],amber_relaxer: relax.AmberRelaxation,benchmark: bool,random_seed: int,models_to_relax: ModelsToRelax,use_saved_msa: bool):"""Predicts structure using AlphaFold for the given sequence."""logging.info('[CL] The flag of use_saved_msa is %s', str(use_saved_msa))...if not use_saved_msa:feature_dict = data_pipeline.process(input_fasta_path=fasta_path,msa_output_dir=msa_output_dir)else:  # 切换处理函数feature_dict = data_pipeline.process_with_saved_msa(input_fasta_path=fasta_path,msa_output_dir=msa_output_dir)

3. 修改 MSA 函数

即，由run_alphafold.py#predict_structure()调用alphafold/data/pipeline.py#process_with_saved_msa()。

在process()函数的基础之上，进行修改。主要改动如下：

添加 MSA 文件夹的读取函数，读取全部已搜索的 MSA 文件，调用 parse_a3m 和 parse_stockholm 函数。
去除 AF2 默认搜索，即 mgnify 和 bfd，两者库较大，但是，保留 uniref90，原因是 uniref90 是搜索模版 (Template) 的前置。
组合全部的 MSA ，作为之后的 MSA 输入。

完整源码修改，如下：

  def process_with_saved_msa(self, input_fasta_path: str, msa_output_dir: str) -> FeatureDict:"""Runs alignment tools on the input sequence and creates features. by chenlong"""with open(input_fasta_path) as f:input_fasta_str = f.read()input_seqs, input_descs = parsers.parse_fasta(input_fasta_str)if len(input_seqs) != 1:raise ValueError(f'More than one input sequence found in {input_fasta_path}.')# 读取函数def load_saved_msa(_msa_output_dir):"""load saved msa, only support a3m and sto."""from myutils.project_utils import traverse_dir_filesa3m_path_list = traverse_dir_files(_msa_output_dir, ext="a3m")sto_path_list = traverse_dir_files(_msa_output_dir, ext="sto")msa_list = []for a3m_path in a3m_path_list:with open(a3m_path) as f:a3m_string = f.read()msa_list.append(parsers.parse_a3m(a3m_string))for sto_path in sto_path_list:# use default stoif os.path.basename(sto_path) == 'uniref90_hits.sto':continuewith open(sto_path) as f:sto_string = f.read()msa_list.append(parsers.parse_stockholm(sto_string))logging.info("[CL] Load saved extra %d msas with %d a3m and %d sto (without uniref90_hits.sto).",len(msa_list), len(a3m_path_list), len(sto_path_list))return msa_list# 读取全部的已搜索文件msa_list = load_saved_msa(msa_output_dir)  # load saved msaassert len(msa_list) >= 1, \"[CL] If the flag of use_saved_msa is true, " \" %s must have 1 extra msa as least (>=1).".format(msa_output_dir)input_sequence = input_seqs[0]input_description = input_descs[0]num_res = len(input_sequence)uniref90_out_path = os.path.join(msa_output_dir, 'uniref90_hits.sto')jackhmmer_uniref90_result = run_msa_tool(msa_runner=self.jackhmmer_uniref90_runner,input_fasta_path=input_fasta_path,msa_out_path=uniref90_out_path,msa_format='sto',use_precomputed_msas=self.use_precomputed_msas,max_sto_sequences=self.uniref_max_hits)msa_for_templates = jackhmmer_uniref90_result['sto']msa_for_templates = parsers.deduplicate_stockholm_msa(msa_for_templates)msa_for_templates = parsers.remove_empty_columns_from_stockholm_msa(msa_for_templates)if self.template_searcher.input_format == 'sto':pdb_templates_result = self.template_searcher.query(msa_for_templates)elif self.template_searcher.input_format == 'a3m':uniref90_msa_as_a3m = parsers.convert_stockholm_to_a3m(msa_for_templates)pdb_templates_result = self.template_searcher.query(uniref90_msa_as_a3m)else:raise ValueError('Unrecognized template input format: 'f'{self.template_searcher.input_format}')pdb_hits_out_path = os.path.join(msa_output_dir, f'pdb_hits.{self.template_searcher.output_format}')with open(pdb_hits_out_path, 'w') as f:f.write(pdb_templates_result)uniref90_msa = parsers.parse_stockholm(jackhmmer_uniref90_result['sto'])pdb_template_hits = self.template_searcher.get_template_hits(output_string=pdb_templates_result, input_sequence=input_sequence)templates_result = self.template_featurizer.get_templates(query_sequence=input_sequence,hits=pdb_template_hits)sequence_features = make_sequence_features(sequence=input_sequence,description=input_description,num_res=num_res)msa_features = make_msa_features((uniref90_msa, *msa_list))logging.info('Uniref90 MSA size: %d sequences.', len(uniref90_msa))logging.info('Final (deduplicated) MSA size: %d sequences.',msa_features['num_alignments'][0])logging.info('Total number of templates (NB: this can include bad ''templates and is later filtered to top 4): %d.',templates_result.features['template_domain_names'].shape[0])return {**sequence_features, **msa_features, **templates_result.features}

4. 修改 MSA 函数

准备数据阶段：

在输出文件夹中，创建待搜索 FASTA 文件的同名文件夹，作为序列输出文件夹，例如 outputs/T1104-D1_A117。
在序列输出文件夹中，创建 msas 文件夹，例如 outputs/T1104-D1_A117/msas。
将已准备的 MSA 文件，支持 a3m 格式和 sto 格式，运行时，增加 -s true，即可。