run_af_multimer_step3.py

#!/usr/bin/env python
#-*- coding:utf-8 -*-
###############################################
#
#
#   Run AlphaFold-Multimer step by step
#   (https://github.com/deepmind/alphafold)
#   Author: Pan Li (lipan@shuimubio.com)
#                    @ Shuimu BioScience
#        https://www.shuimubio.com/
#
#
################################################

#
#
#  AlphaFold-Multimer Step 3 -- Relax models
#  Usage: run_af_multimer_step3.py /path/to/features.pkl /path/to/result_model.pkl /path/to/output.pdb
#
#

import json
import os
import pathlib
import pickle
import random
import shutil
import sys
import time
import gzip
from typing import Dict, Union, Optional
import numpy as np
import argparse
import configparser

# os.environ['CUDA_VISIBLE_DEVICES'] = ''
os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'
os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '4.0'

cur_path = pathlib.Path(__file__).parent.resolve()

ini_config = configparser.ConfigParser(allow_no_value=True)
assert len(ini_config.read(os.path.join(cur_path, 'config.ini'))) > 0, "Read config.ini failed"

sys.path.insert(0, ini_config['ALPHAFOLD2']['alphafold_path'])

from alphafold.common import protein
from alphafold.common import residue_constants
from alphafold.relax import relax

parser = argparse.ArgumentParser(description='AlphaFold-Multimer Step 3 -- Relax models')

parser.add_argument('model_pdb', metavar='model_pdb', type=str, help='The result_model.pdb file generated by AlphaFold-Multimer step 2.')
parser.add_argument('output_pdb_file', metavar='output_pdb_file', type=str, help='The PDB file to save the relaxed models.')

args = parser.parse_args()

if os.path.exists(args.output_pdb_file):
    print(f"Info: {args.output_pdb_file} exists, please delete and try again")
    exit(0)

######################
## Define the relaxer
######################

RELAX_MAX_ITERATIONS = 0
RELAX_ENERGY_TOLERANCE = 2.39
RELAX_STIFFNESS = 10.0
RELAX_EXCLUDE_RESIDUES = []
RELAX_MAX_OUTER_ITERATIONS = 3
use_gpu_relax=False # It is easy to Minimization failed using GPU

amber_relaxer = relax.AmberRelaxation(
    max_iterations=RELAX_MAX_ITERATIONS,
    tolerance=RELAX_ENERGY_TOLERANCE,
    stiffness=RELAX_STIFFNESS,
    exclude_residues=RELAX_EXCLUDE_RESIDUES,
    max_outer_iterations=RELAX_MAX_OUTER_ITERATIONS,
    use_gpu=use_gpu_relax)

######################
## Define Protein object
######################

pdb_string = open(args.model_pdb).read()
unrelaxed_protein = protein.from_pdb_string(pdb_string)

######################
## Start to relax
######################

print(f"Relax model {args.model_pdb}", flush=True)
relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
print( relaxed_pdb_str, file=open(args.output_pdb_file, 'w') )