#!/usr/bin/env python import os import numpy as np import sys from timeit import default_timer as timer import argparse import glob import uuid import jax import jax.numpy as jnp from jax.lib import xla_bridge from alphafold.common import residue_constants from alphafold.common import protein from alphafold.common import confidence from alphafold.data import pipeline from alphafold.model import data from alphafold.model import config from alphafold.model import model import af2_util parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.append(os.path.join(parent, 'include')) from silent_tools import silent_tools from pyrosetta import * from rosetta import * init( '-in:file:silent_struct_type binary -mute all' ) def range1(size): return range(1, size+1) ################################# # Parse Arguments ################################# parser = argparse.ArgumentParser() # I/O Arguments parser.add_argument( "-pdbdir", type=str, default="", help='The name of a directory of pdbs to run through the model' ) parser.add_argument( "-silent", type=str, default="", help='The name of a silent file to run through the model' ) parser.add_argument( "-outpdbdir", type=str, default="outputs", help='The directory to which the output PDB files will be written. Only used when -pdbdir is active' ) parser.add_argument( "-outsilent", type=str, default="out.silent", help='The name of the silent file to which output structs will be written. Only used when -silent is active' ) parser.add_argument( "-runlist", type=str, default='', help="The path of a list of pdb tags to run. Only used when -pdbdir is active (default: ''; Run all PDBs)" ) parser.add_argument( "-checkpoint_name", type=str, default='check.point', help="The name of a file where tags which have finished will be written (default: check.point)" ) parser.add_argument( "-scorefilename", type=str, default='out.sc', help="The name of a file where scores will be written (default: out.sc)" ) parser.add_argument( "-maintain_res_numbering", action="store_true", default=False, help='When active, the model will not renumber the residues when bad inputs are encountered (default: False)' ) parser.add_argument( "-debug", action="store_true", default=False, help='When active, errors will cause the script to crash and the error message to be printed out (default: False)') # AF2-Specific Arguments parser.add_argument( "-max_amide_dist", type=float, default=3.0, help='The maximum distance between an amide bond\'s carbon and nitrogen (default: 3.0)' ) parser.add_argument( "-recycle", type=int, default=3, help='The number of AF2 recycles to perform (default: 3)' ) parser.add_argument( "-no_initial_guess", action="store_true", default=False, help='When active, the model will not use an initial guess (default: False)' ) parser.add_argument( "-force_monomer", action="store_true", default=False, help='When active, the model will predict the structure of a monomer (default: False)' ) args = parser.parse_args() class FeatureHolder(): ''' This is a struct which holds the features for a single structure being run through the model ''' def __init__(self, pose, monomer, binderlen, tag): self.pose = pose self.tag = tag self.outtag = self.tag + '_af2pred' self.seq = pose.sequence() self.binderlen = binderlen self.monomer = monomer # Pre model features self.initial_all_atom_positions = None self.initial_all_atom_masks = None # Post model features self.outpose = None self.plddt_array = None self.score_dict = None class AF2_runner(): ''' This class handles generating features, running the model, and parsing outputs ''' def __init__(self, args, struct_manager): self.max_amide_dist = args.max_amide_dist # For timing self.t0 = None self.struct_manager = struct_manager # Other models may be run but their weights will also need to be downloaded self.model_name = "model_1_ptm" model_config = config.model_config(self.model_name) model_config.data.eval.num_ensemble = 1 model_config.data.common.num_recycle = args.recycle model_config.model.num_recycle = args.recycle model_config.model.embeddings_and_evoformer.initial_guess = False if args.no_initial_guess else True model_config.data.common.max_extra_msa = 5 model_config.data.eval.max_msa_clusters = 5 package_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) params_dir = os.path.join(package_root, 'weight', 'AlphaFold2') model_params = data.get_model_haiku_params(model_name=self.model_name, data_dir=params_dir) self.model_runner = model.RunModel(model_config, model_params) def featurize(self, feat_holder) -> None: all_atom_positions, all_atom_masks = af2_util.af2_get_atom_positions(feat_holder.pose, self.struct_manager.tmp_fn) feat_holder.initial_all_atom_positions = all_atom_positions feat_holder.initial_all_atom_masks = all_atom_masks initial_guess = af2_util.parse_initial_guess(feat_holder.initial_all_atom_positions) # Determine which residues to template if feat_holder.monomer: # For monomers predict all residues feat_holder.residue_mask = [False for i in range(len(feat_holder.seq))] else: # For interfaces fix the target and predict the binder feat_holder.residue_mask = [int(i) > feat_holder.binderlen for i in range(len(feat_holder.seq))] template_dict = af2_util.generate_template_features( feat_holder.seq, feat_holder.initial_all_atom_positions, feat_holder.initial_all_atom_masks, feat_holder.residue_mask ) # Gather features feature_dict = { **pipeline.make_sequence_features(sequence=feat_holder.seq, description="none", num_res=len(feat_holder.seq)), **pipeline.make_msa_features(msas=[[feat_holder.seq]], deletion_matrices=[[[0]*len(feat_holder.seq)]]), **template_dict } if feat_holder.monomer: breaks = [] else: breaks = af2_util.check_residue_distances( feat_holder.initial_all_atom_positions, feat_holder.initial_all_atom_masks, self.max_amide_dist ) feature_dict['residue_index'] = af2_util.insert_truncations(feature_dict['residue_index'], breaks) feature_dict = self.model_runner.process_features(feature_dict, random_seed=0) return feature_dict, initial_guess def generate_scoredict(self, feat_holder, confidences, rmsds) -> None: ''' Collect the confidence values, slicing them to the binder and target regions then add the parsed scores to the score_dict ''' binderlen = feat_holder.binderlen plddt_array = confidences['plddt'] plddt = np.mean( plddt_array ) if feat_holder.monomer: plddt_binder = np.mean( plddt_array ) plddt_target = float('nan') else: plddt_binder = np.mean( plddt_array[:binderlen] ) plddt_target = np.mean( plddt_array[binderlen:] ) pae = confidences['predicted_aligned_error'] if feat_holder.monomer: pae_binder = np.mean( pae ) pae_target = float('nan') pae_interaction_total = float('nan') else: pae_interaction1 = np.mean( pae[:binderlen,binderlen:] ) pae_interaction2 = np.mean( pae[binderlen:,:binderlen] ) pae_binder = np.mean( pae[:binderlen,:binderlen] ) pae_target = np.mean( pae[binderlen:,binderlen:] ) pae_interaction_total = ( pae_interaction1 + pae_interaction2 ) / 2 time = timer() - self.t0 score_dict = { "plddt_total" : plddt, "plddt_binder" : plddt_binder, "plddt_target" : plddt_target, "pae_binder" : pae_binder, "pae_target" : pae_target, "pae_interaction" : pae_interaction_total, "binder_aligned_rmsd": rmsds['binder_aligned_rmsd'], "target_aligned_rmsd": rmsds['target_aligned_rmsd'], "time" : time } # Store this in the feature holder for later use feat_holder.score_dict = score_dict # If we ever want to write strings to the score file we can do it here string_dict = None self.struct_manager.record_scores(feat_holder.outtag, score_dict, string_dict) print(score_dict) print(f"Tag: {feat_holder.outtag} reported success in {time} seconds") def process_output(self, feat_holder, feature_dict, prediction_result) -> None: ''' Take the AF2 output, parse the confidence scores from this and register the scores in the score file Also write out the structure ''' # First extract the structure and confidence scores from the prediction result structure_module = prediction_result['structure_module'] this_protein = protein.Protein( aatype=feature_dict['aatype'][0], atom_positions=structure_module['final_atom_positions'][...], atom_mask=structure_module['final_atom_mask'][...], residue_index=feature_dict['residue_index'][0] + 1, b_factors=np.zeros_like(structure_module['final_atom_mask'][...]) ) confidences = {} confidences['distogram'] = prediction_result['distogram'] confidences['plddt'] = confidence.compute_plddt( prediction_result['predicted_lddt']['logits'][...]) if 'predicted_aligned_error' in prediction_result: confidences.update(confidence.compute_predicted_aligned_error( prediction_result['predicted_aligned_error']['logits'][...], prediction_result['predicted_aligned_error']['breaks'][...])) feat_holder.plddt_array = confidences['plddt'] # Calculate the RMSDs target_mask = np.zeros(len(feat_holder.seq), dtype=bool) target_mask[feat_holder.binderlen:] = True rmsds = af2_util.calculate_rmsds( feat_holder.initial_all_atom_positions, this_protein.atom_positions, target_mask ) # Write the structure as a pdb file so Rosetta can read it unrelaxed_pdb_lines = protein.to_pdb(this_protein) with open(self.struct_manager.tmp_fn, 'w') as f: f.write(unrelaxed_pdb_lines) # Now read the structure to a Rosetta pose feat_holder.outpose = pyrosetta.pose_from_file(self.struct_manager.tmp_fn) os.remove(self.struct_manager.tmp_fn) # Now we can finally write the scores and the predicted structure to disk self.generate_scoredict(feat_holder, confidences, rmsds) self.struct_manager.dump_pose(feat_holder) def process_struct(self, tag) -> None: # Start the timer self.t0 = timer() # Load the structure pose, monomer, binderlen, usetag = self.struct_manager.load_pose(tag) # Store the pose in the feature holder feat_holder = FeatureHolder(pose, monomer, binderlen, usetag) print(f'Processing struct with tag: {feat_holder.tag}') # Generate features feature_dict, initial_guess = self.featurize(feat_holder) # Run model start = timer() print(f'Running {self.model_name}') prediction_result = self.model_runner.apply( self.model_runner.params, jax.random.PRNGKey(0), feature_dict, initial_guess ) print(f'Tag: {feat_holder.tag} finished AF2 prediction in {timer() - start} seconds') # Process outputs self.process_output(feat_holder, feature_dict, prediction_result) class StructManager(): ''' This class handles all of the input and output for the AF2 model. It deals with silent files vs. pdbs, checkpointing, and writing of outputs Note: This class could be moved to a separate file and shared with ProteinMPNN ''' def __init__(self, args): self.args = args self.maintain_res_numbering = args.maintain_res_numbering self.score_fn = args.scorefilename # Generate a random unique temporary filename self.tmp_fn = f'tmp_{uuid.uuid4()}.pdb' self.force_monomer = args.force_monomer self.silent = False if not args.silent == '': self.silent = True self.struct_iterator = silent_tools.get_silent_index(args.silent)['tags'] self.sfd_in = rosetta.core.io.silent.SilentFileData(rosetta.core.io.silent.SilentFileOptions()) self.sfd_in.read_file(args.silent) self.sfd_out = core.io.silent.SilentFileData(args.outsilent, False, False, "binary", core.io.silent.SilentFileOptions()) self.outsilent = args.outsilent self.pdb = False if not args.pdbdir == '': self.pdb = True self.pdbdir = args.pdbdir self.outpdbdir = args.outpdbdir self.struct_iterator = glob.glob(os.path.join(args.pdbdir, '*.pdb')) # Parse the runlist and determine which structures to process if args.runlist != '': with open(args.runlist, 'r') as f: self.runlist = set([line.strip() for line in f]) # Filter the struct iterator to only include those in the runlist self.struct_iterator = [struct for struct in self.struct_iterator if '.'.join(os.path.basename(struct).split('.')[:-1]) in self.runlist] print(f'After filtering by runlist, {len(self.struct_iterator)} structures remain') # Assert that either silent or pdb is true, but not both assert(self.silent ^ self.pdb), f'Both silent and pdb are set to {args.silent} and {args.pdb} respectively. Only one of these may be active at a time' # Setup checkpointing self.chkfn = args.checkpoint_name self.finished_structs = set() if os.path.isfile(self.chkfn): with open(self.chkfn, 'r') as f: for line in f: self.finished_structs.add(line.strip()) def input_check(self, pose, tag) -> bool: ''' This function checks that the given pose is valid for AF2. It specifically checks that all residue indices are unique ''' seen_indices = set() # Loop through the PDBinfo of the pose and check that all residue indices are unique pdbinfo = pose.pdb_info() for resi in range1(pose.size()): residx = pdbinfo.number(resi) if residx in seen_indices: print( f"\nNon-unique residue indices detected for tag: {tag}. " + "This will cause AF2 to yield garbage outputs." ) return False seen_indices.add(residx) return True def record_checkpoint(self, tag): ''' Record the fact that this tag has been processed. Write this tag to the list of finished structs ''' with open(self.chkfn, 'a') as f: f.write(f'{tag}\n') def iterate(self): ''' Iterate over the silent file or pdb directory and run the model on each structure ''' # Iterate over the structs and for each, check that the struct has not already been processed for struct in self.struct_iterator: if self.pdb: tag = '.'.join(os.path.basename(struct).split('.')[:-1]) else: tag = struct if tag in self.finished_structs: print(f'{tag} has already been processed. Skipping') continue yield struct def record_scores(self, tag, score_dict, string_dict): ''' Record the scores for this structure to the score file. Args: tag (str): The tag for this structure score_dict (dict): A dictionary of numerical scores to record string_dict (dict): A dictionary of string scores to record ''' # Check whether the score file exists write_header = False if not os.path.isfile(self.score_fn): write_header = True af2_util.add2scorefile(tag, self.score_fn, write_header, score_dict, string_dict) def dump_pose(self, feat_holder): ''' Dump this pose to either a silent file or a pdb file depending on the input arguments ''' if feat_holder.monomer: pose = feat_holder.outpose else: # Insert chainbreaks into the pose pose = af2_util.insert_Rosetta_chainbreaks(feat_holder.outpose, feat_holder.binderlen) # Add the plddt scores as b-factors to the pose info = pose.pdb_info() for resi in range1(pose.size()): info.add_reslabel(resi, f'{feat_holder.plddt_array[resi-1]}') for atom_i in range1(pose.residue(resi).natoms()): info.bfactor(resi, atom_i, feat_holder.plddt_array[resi-1]) # Assign the pose the updated pdb_info pose.pdb_info(info) if self.pdb: # If the outpdbdir does not exist, create it # If there are parents in the path that do not exist, create them as well if not os.path.exists(self.outpdbdir): os.makedirs(self.outpdbdir) pdbfile = os.path.join(self.outpdbdir, feat_holder.outtag + '.pdb') pose.dump_pdb(pdbfile) if self.silent: struct = self.sfd_out.create_SilentStructOP() struct.fill_struct(pose, feat_holder.outtag) # Write the scores to the silent file for scorename, value in feat_holder.score_dict.items(): struct.add_energy(scorename, value, 1) self.sfd_out.add_structure(struct) self.sfd_out.write_silent_struct(struct, self.outsilent) def load_pose(self, tag): ''' Load a pose from either a silent file or a pdb file depending on the input arguments. Also run input checking on the pose to ensure that it is valid for AF2. The pose will be automatically renumbered if -maintain_res_numbering is not False ''' if not self.pdb and not self.silent: raise Exception('Neither pdb nor silent is set to True. Cannot load pose') if self.pdb: pose = pose_from_pdb(tag) usetag = '.'.join(os.path.basename(tag).split('.')[:-1]) if self.silent: pose = Pose() usetag = tag self.sfd_in.get_structure(tag).fill_pose(pose) # Run input checking on the pose passes_check = self.input_check(pose, tag) if not passes_check: if not self.maintain_res_numbering: print( f"Renumbering {tag}" ) info = core.pose.PDBInfo(pose) pose.pdb_info(info) else: raise Exception( f"Pose {tag} failed input checking.") # Determine whether we are working with a monomer or a complex splits = pose.split_by_chain() if len(splits) > 2: raise Exception( f"Pose {tag} has more than two chains. This is not supported by this script." ) elif len(splits) < 1: raise Exception( f"Pose {tag} is empty. This is not supported by this script." ) elif len(splits) == 1: monomer = True binderlen = -1 elif len(splits) == 2: if self.force_monomer: print( "/" * 60 ) print( f"Pose {tag} has two chains. But force_monomer is set to True. Treating as monomer.") print( f"I am going to assume that the first chain is the binder and that is the chain I will predict") print( "/" * 60 ) monomer = True pose = splits[1] binderlen = -1 else: monomer = False binderlen = splits[1].size() return pose, monomer, binderlen, usetag #################### ####### Main ####### #################### device = xla_bridge.get_backend().platform if device == 'gpu': print('/' * 60) print('/' * 60) print('Found GPU and will use it to run AF2') print('/' * 60) print('/' * 60) print('\n') else: print('/' * 60) print('/' * 60) print('WARNING! No GPU detected running AF2 on CPU') print('/' * 60) print('/' * 60) print('\n') struct_manager = StructManager(args) af2_runner = AF2_runner(args, struct_manager) for pdb in struct_manager.iterate(): if args.debug: af2_runner.process_struct(pdb) else: # When not in debug mode the script will continue to run even when some poses fail t0 = timer() try: af2_runner.process_struct(pdb) except KeyboardInterrupt: sys.exit( "Script killed by Control+C, exiting" ) except: seconds = int(timer() - t0) print( "Struct with tag %s failed in %i seconds with error: %s"%( pdb, seconds, sys.exc_info()[0] ) ) # We are done with one pdb, record that we finished struct_manager.record_checkpoint(pdb)