run_af2_step3.py

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

#
#
#  AlphaFold2 Step 3 -- Relax models
#  Usage: run_af2_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['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='AlphaFold2 Step 3 -- Relax models')

parser.add_argument('model_pdb', metavar='model_pdb', type=str, help='The result_model.pdb file generated by AlphaFold2 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') )