OneScience-Group/flex_ddG_tutorial

1#!/usr/bin/python2 3from __future__ import print_function4 5import socket6import sys7import os8import subprocess9 10script_dir = os.path.dirname(os.path.abspath(__file__))11project_root = os.path.dirname(script_dir)12conf_dir = os.path.join(project_root, 'conf')13inputs_dir = os.path.join(script_dir, 'inputs')14 15use_multiprocessing = True16if use_multiprocessing:17    import multiprocessing18    max_cpus = 2 # We might want to not run on the full number of cores, as Rosetta take about 2 Gb of memory per instance19 20###################################################################################################################################################################21# Important: The variables below are set to values that will make the run complete faster (as a tutorial example), but will not give scientifically valid results.22#            Please change them to the "normal" default values before a real run.23###################################################################################################################################################################24 25#rosetta_scripts_path = os.path.expanduser("~/rosetta/source/bin/rosetta_scripts")26rosetta_scripts_path = os.path.expanduser(27"/public/home/scnb9biwet/jiangqq/flex_ddG_tutorial-master/software/rosetta3.9/main/source/bin/rosetta_scripts.default.linuxgccrelease"28)29nstruct = 3 # Normally 3530max_minimization_iter = 5 # Normally 500031abs_score_convergence_thresh = 200.0 # Normally 1.032number_backrub_trials = 10 # Normally 3500033backrub_trajectory_stride = 5 # Can be whatever you want, if you would like to see results from earlier time points in the backrub trajectory. 7000 is a reasonable number, to give you three checkpoints for a 35000 step run, but you could also set it to 35000 for quickest run time (as the final minimization and packing steps will only need to be run one time).34path_to_script = os.path.join(conf_dir, 'ddG-backrub.xml')35 36# Additionally report the intramolecular energy of each individual chain into ddG.db3.37# This is reporting only: it adds no sampling, consumes no random numbers, and leaves the38# binding ddG bit-identical. Read the per-chain numbers off the *unbound* states, where the39# chains are 1000 A apart and the value is exactly that chain's intramolecular energy.40# See per_chain_protocol.py for the full explanation and caveats.41report_per_chain_energies = False42 43if not os.path.isfile(rosetta_scripts_path):44    print('ERROR: "rosetta_scripts_path" variable must be set to the location of the "rosetta_scripts" binary executable')45    print('This file might look something like: "rosetta_scripts.linuxgccrelease"')46    raise Exception('Rosetta scripts missing')47 48def protocol_for( input_pdb_path ):49    """Return the XML to run, generating the per-chain reporting variant if requested."""50    if not report_per_chain_energies:51        return os.path.abspath(path_to_script)52 53    import per_chain_protocol54    chains = per_chain_protocol.chains_in_pdb( input_pdb_path )55    derived = os.path.abspath( os.path.splitext(path_to_script)[0] + '-per_chain.xml' )56    per_chain_protocol.write_per_chain_protocol( path_to_script, chains, derived )57    return derived58 59def run_flex_ddg( name, input_path, input_pdb_path, chains_to_move, nstruct_i ):60    output_directory = os.path.join( project_root, 'output', os.path.join( name, '%02d' % nstruct_i ) )61    if not os.path.isdir(output_directory):62        os.makedirs(output_directory)63 64    flex_ddg_args = [65        os.path.abspath(rosetta_scripts_path),66        "-s %s" % os.path.abspath(input_pdb_path),67        '-parser:protocol', protocol_for( input_pdb_path ),68        '-parser:script_vars',69        'chainstomove=' + chains_to_move,70        'mutate_resfile_relpath=' + os.path.abspath( os.path.join( input_path, 'nataa_mutations.resfile' ) ),71        'number_backrub_trials=%d' % number_backrub_trials,72        'max_minimization_iter=%d' % max_minimization_iter,73        'abs_score_convergence_thresh=%.1f' % abs_score_convergence_thresh,74        'backrub_trajectory_stride=%d' % backrub_trajectory_stride ,75        '-restore_talaris_behavior',76        '-in:file:fullatom',77        '-ignore_unrecognized_res',78        '-ignore_zero_occupancy false',79        '-ex1',80        '-ex2',81    ]82 83    log_path = os.path.join(output_directory, 'rosetta.out')84 85    print( 'Running Rosetta with args:' )86    print( ' '.join(flex_ddg_args) )87    print( 'Output logged to:', os.path.abspath(log_path) )88    print()89 90    outfile = open(log_path, 'w')91    process = subprocess.Popen(flex_ddg_args, stdout=outfile, stderr=subprocess.STDOUT, close_fds = True, cwd = output_directory)92    returncode = process.wait()93    outfile.close()94 95if __name__ == '__main__':96    cases = []97    for nstruct_i in range(1, nstruct + 1 ):98        for case_name in os.listdir(inputs_dir):99            case_path = os.path.join( inputs_dir, case_name )100            for f in os.listdir(case_path):101                if f.endswith('.pdb'):102                    input_pdb_path = os.path.join( case_path, f )103                    break104 105            with open( os.path.join( case_path, 'chains_to_move.txt' ), 'r' ) as f:106                chains_to_move = f.readlines()[0].strip()107 108            cases.append( (case_name, case_path, input_pdb_path, chains_to_move, nstruct_i) )109 110    if use_multiprocessing:111        pool = multiprocessing.Pool( processes = min(max_cpus, multiprocessing.cpu_count()) )112 113    for args in cases:114        if use_multiprocessing:115            pool.apply_async( run_flex_ddg, args = args )116        else:117            run_flex_ddg( *args )118 119    if use_multiprocessing:120        pool.close()121        pool.join()122