changes_to_alphafold.txt

Only in alphafold_v.2.2_customized/alphafold/data: custom_templates.py


diff -b -u -r alphafold_v2.2_original/alphafold/data/feature_processing.py alphafold_v.2.2_customized/alphafold/data/feature_processing.py
--- alphafold_v2.2_original/alphafold/data/feature_processing.py	2022-05-09 17:08:36.348921075 -0400
+++ alphafold_v.2.2_customized/alphafold/data/feature_processing.py	2023-01-04 10:52:45.506208747 -0500
@@ -60,7 +60,6 @@
   process_unmerged_features(all_chain_features)
 
   np_chains_list = list(all_chain_features.values())
-
   pair_msa_sequences = not _is_homomer_or_monomer(np_chains_list)
 
   if pair_msa_sequences:


diff -b -u -r alphafold_v2.2_original/alphafold/data/parsers.py alphafold_v.2.2_customized/alphafold/data/parsers.py
--- alphafold_v2.2_original/alphafold/data/parsers.py	2022-12-14 15:58:19.970260694 -0500
+++ alphafold_v.2.2_customized/alphafold/data/parsers.py	2023-03-07 11:19:39.153783844 -0500
@@ -91,6 +91,11 @@
       continue  # Skip blank lines.
     sequences[index] += line
 
+  # Use gap to break a single chain
+  for index,sequence in enumerate(sequences):
+    if ":" in sequence:
+      sequences[index]=sequence.upper().split(":")
+      
   return sequences, descriptions
 
 
@@ -296,6 +301,25 @@
 
   return ''.join(filtered_lines)
 
+def truncate_stockholm_msa2(stockholm_msa: str, max_sequences: int) -> str:
+  """Reads + truncates a Stockholm alignment - BP 12/15/22."""
+  seqnames = set()
+  filtered_lines = []
+  
+  for line in stockholm_msa.splitlines():    
+    if line.strip() and not line.startswith(('#', '//')):
+      # Ignore blank lines, markup and end symbols - remainder are alignment
+      # sequence parts.
+      seqname = line.partition(' ')[0]
+      seqnames.add(seqname)
+      if len(seqnames) >= max_sequences:
+        break
+
+  for line in stockholm_msa.splitlines():
+    if _keep_line(line, seqnames):
+      filtered_lines.append(line)
+
+  return '\n'.join(filtered_lines) + '\n'
 
 def remove_empty_columns_from_stockholm_msa(stockholm_msa: str) -> str:
   """Removes empty columns (dashes-only) from a Stockholm MSA."""
@@ -410,7 +434,7 @@
   # Parse first 2 lines.
   number_of_hit = int(detailed_lines[0].split()[-1])
   name_hit = detailed_lines[1][1:]
-  print("detailed_line:%s " % str(detailed_lines[2]))
+
   # Parse the summary line.
   pattern = (
       'Probab=(.*)[\t ]*E-value=(.*)[\t ]*Score=(.*)[\t ]*Aligned_cols=(.*)[\t'


Only in alphafold_v.2.2_customized/alphafold/data: pipeline_custom_templates.py


Only in alphafold_v.2.2_customized/alphafold/data: pipeline_multimer_custom_templates.py


diff -b -u -r alphafold_v2.2_original/alphafold/data/pipeline_multimer.py alphafold_v.2.2_customized/alphafold/data/pipeline_multimer.py
--- alphafold_v2.2_original/alphafold/data/pipeline_multimer.py	2022-05-09 17:08:36.357921193 -0400
+++ alphafold_v.2.2_customized/alphafold/data/pipeline_multimer.py	2023-03-07 11:21:00.749011883 -0500
@@ -200,7 +200,10 @@
       sequence: str,
       description: str,
       msa_output_dir: str,
-      is_homomer_or_monomer: bool) -> pipeline.FeatureDict:
+      is_homomer_or_monomer: bool,
+      save_msa_fasta: bool,
+      save_template_names: bool,
+      msa_for_template_query_seq_only: bool) -> pipeline.FeatureDict:
     """Runs the monomer pipeline on a single chain."""
     chain_fasta_str = f'>chain_{chain_id}\n{sequence}\n'
     chain_msa_output_dir = os.path.join(msa_output_dir, chain_id)
@@ -211,7 +214,10 @@
                    chain_id, description)
       chain_features = self._monomer_data_pipeline.process(
           input_fasta_path=chain_fasta_path,
-          msa_output_dir=chain_msa_output_dir)
+          msa_output_dir=chain_msa_output_dir,
+          save_msa_fasta=save_msa_fasta,
+          save_template_names=save_template_names,
+          msa_for_template_query_seq_only=msa_for_template_query_seq_only)
 
       # We only construct the pairing features if there are 2 or more unique
       # sequences.
@@ -239,7 +245,10 @@
 
   def process(self,
               input_fasta_path: str,
-              msa_output_dir: str) -> pipeline.FeatureDict:
+              msa_output_dir: str,
+              save_msa_fasta: bool,
+              save_template_names: bool,
+              msa_for_template_query_seq_only: bool) -> pipeline.FeatureDict:
     """Runs alignment tools on the input sequences and creates features."""
     with open(input_fasta_path) as f:
       input_fasta_str = f.read()
@@ -266,7 +275,10 @@
           sequence=fasta_chain.sequence,
           description=fasta_chain.description,
           msa_output_dir=msa_output_dir,
-          is_homomer_or_monomer=is_homomer_or_monomer)
+          is_homomer_or_monomer=is_homomer_or_monomer,
+          save_msa_fasta=save_msa_fasta,
+          save_template_names=save_template_names,
+          msa_for_template_query_seq_only=msa_for_template_query_seq_only)
 
       chain_features = convert_monomer_features(chain_features,
                                                 chain_id=chain_id)


diff -b -u -r alphafold_v2.2_original/alphafold/data/pipeline.py alphafold_v.2.2_customized/alphafold/data/pipeline.py
--- alphafold_v2.2_original/alphafold/data/pipeline.py	2022-12-14 16:33:15.504019588 -0500
+++ alphafold_v.2.2_customized/alphafold/data/pipeline.py	2023-03-07 11:22:18.847187291 -0500
@@ -147,7 +147,11 @@
     self.uniref_max_hits = uniref_max_hits
     self.use_precomputed_msas = use_precomputed_msas
 
-  def process(self, input_fasta_path: str, msa_output_dir: str) -> FeatureDict:
+  def process(self, input_fasta_path: str, 
+              msa_output_dir: str, 
+              save_msa_fasta: bool, 
+              save_template_names: bool,
+              msa_for_template_query_seq_only: bool) -> FeatureDict:
     """Runs alignment tools on the input sequence and creates features."""
     with open(input_fasta_path) as f:
       input_fasta_str = f.read()
@@ -180,6 +184,9 @@
     msa_for_templates = parsers.deduplicate_stockholm_msa(msa_for_templates)
     msa_for_templates = parsers.remove_empty_columns_from_stockholm_msa(
         msa_for_templates)
+    if msa_for_template_query_seq_only:
+      '''BP 12/15/22 let's just keep the target (query) sequence'''
+      msa_for_templates = parsers.truncate_stockholm_msa2(msa_for_templates, 1)
 
     if self.template_searcher.input_format == 'sto':
       pdb_templates_result = self.template_searcher.query(msa_for_templates)
@@ -231,6 +238,19 @@
 
     msa_features = make_msa_features((uniref90_msa, bfd_msa, mgnify_msa))
 
+    if save_msa_fasta:
+      msa_outpath=os.path.join(msa_output_dir, 'msa_feat_gaptoU.fasta')
+      with open(msa_outpath, 'w+') as fh:
+        fh.write(">query"+"\n"+input_sequence+"\n")
+        counter=1
+        for seq in msa_features['msa']:
+            seq=[residue_constants.ID_TO_HHBLITS_AA[num] for num in seq]
+            # for x in range(len(seq)):
+            counter+=1
+            fh.write(">seq_"+str(counter)+"\n")
+            out="".join(seq).replace("-","U")
+            fh.write(out+"\n")
+            
     logging.info('Uniref90 MSA size: %d sequences.', len(uniref90_msa))
     logging.info('BFD MSA size: %d sequences.', len(bfd_msa))
     logging.info('MGnify MSA size: %d sequences.', len(mgnify_msa))
@@ -240,4 +260,12 @@
                  'templates and is later filtered to top 4): %d.',
                  templates_result.features['template_domain_names'].shape[0])
 
+    if save_template_names:
+      temp_name_fn=os.path.join(msa_output_dir, 'template_names.txt')
+      # with open(temp_name_fn, 'w+') as fh:
+      #     fh.write("\n".join((map(str,templates_result.features['template_domain_names']))))
+      template_names=[name.decode('utf-8') for name in templates_result.features['template_domain_names']]
+      with open(temp_name_fn, 'w+') as fh:
+          fh.write("\n".join(template_names))
+          
     return {**sequence_features, **msa_features, **templates_result.features}


diff -b -u -r  alphafold_v2.2_original/run_alphafold.py alphafold_v2.2_customized/run_alphafold_customized.py 
--- ../alphafold_v2.2/run_alphafold.py	2023-03-08 15:18:16.899697496 -0500
+++ run_alphafold_customized.py	2023-07-09 16:26:18.452582062 -0400
@@ -30,6 +30,8 @@
 from alphafold.common import residue_constants
 from alphafold.data import pipeline
 from alphafold.data import pipeline_multimer
+from alphafold.data import pipeline_custom_templates
+from alphafold.data import pipeline_multimer_custom_templates
 from alphafold.data import templates
 from alphafold.data.tools import hhsearch
 from alphafold.data.tools import hmmsearch
@@ -38,12 +40,7 @@
 from alphafold.model import model
 from alphafold.relax import relax
 import numpy as np
-
-root_dir="/piercehome/yinr/alphafold/alphafold_v2.2"
-sys.path.insert(1,root_dir)
-
-# Internal import (7716).
-FLAGS = flags.FLAGS
+import pickle as pkl
 
 logging.set_verbosity(logging.INFO)
 
@@ -132,8 +129,44 @@
                      'Relax on GPU can be much faster than CPU, so it is '
                      'recommended to enable if possible. GPUs must be available'
                      ' if this setting is enabled.')
+flags.DEFINE_boolean('use_custom_templates', False, 'Whether to use custom '
+                    'templates or not.')
+flags.DEFINE_string('template_alignfile', None, 'The path to the custom template'
+                    'files. If the target is a monomer, provide the template path '
+                    'as-is. If a multimer, provide all template alignment files '
+                    'the order they appear in the target, comma seperated. Leave '
+                    'the path blank if no template should be used for a chain. '
+                    'Write "UseDefaultTemplate" to use default alphafold pipeline '
+                    'for generating the template for that chain.')
+flags.DEFINE_string('msa_mode', None, 'Type "single_sequence" to not use any MSA')
+flags.DEFINE_integer('num_recycle', 3, 'How many recycles')
+flags.DEFINE_integer('num_ensemble', 1, 'How many ensembling iteractions')
+flags.DEFINE_enum('use_custom_MSA_database', "none", ["none", "add", "only"], 'Whether to use custom '
+                    'MSA database or not.')
+flags.DEFINE_string('MSA_database', None, 'The path to the custom MSA database'
+                    'files. If the target is a monomer, provide the template path '
+                    'as-is. If a multimer, provide all template alignment files '
+                    'the order they appear in the target, comma seperated.')
+flags.DEFINE_string('run_model_names', None, 'Specify parameter name to run. This'
+                    'is comma seperated alphafold parameter name. Only specified'
+                    'model names will be run.')
+flags.DEFINE_boolean('save_msa_fasta', False, 'Save msa features or not.')
+flags.DEFINE_boolean('gen_feats_only', False, 'Only generate features and do not'
+                    ' produce structure predictions.')
+flags.DEFINE_boolean('save_template_names', False, 'Save template id to txt file.')
+flags.DEFINE_boolean('has_gap_chn_brk', False, 'Have chain breaks introduced by ":".')
+flags.DEFINE_string('substitute_msa', None, 'Path to feature.pkl whose MSA will '
+                    'be used to substitute whatever MSA that will be generated by '
+                    'this prediction round.')
+flags.DEFINE_boolean('msa_for_template_query_seq_only', False, 'msa_for_template_query_seq_only')
+flags.DEFINE_string('iptm_interface', None, 'iptm_interface')
+flags.DEFINE_string('feature_prefix', None, 'Feature prefix')
+flags.DEFINE_boolean('save_ranked_pdb_only', False, 'Do not save result pkl files '
+                     'or unrelaxed pdbs, or relaxed pdbs that are not ranked.')
 
 
+FLAGS = flags.FLAGS
+
 MAX_TEMPLATE_HITS = 20
 RELAX_MAX_ITERATIONS = 0
 RELAX_ENERGY_TOLERANCE = 2.39
@@ -150,16 +183,43 @@
     raise ValueError(f'{flag_name} must {verb} set when running with '
                      f'"--{other_flag_name}={FLAGS[other_flag_name].value}".')
 
+def interface_parser(interfaces_string):
+  interfaces=[]
+  for interface in interfaces_string.split(","):
+    interfaces.append([int(i) for i in interface.split(":")])
+  return interfaces
+
+def gen_res_str(res):
+    res_len = len(str(res))
+    res_str = ""
+    for i in range(4-res_len):
+        res_str += " "
+    res_str += str(res) + " "
+    return res_str
 
 def predict_structure(
     fasta_path: str,
     fasta_name: str,
     output_dir_base: str,
-    data_pipeline: Union[pipeline.DataPipeline, pipeline_multimer.DataPipeline],
+    data_pipeline: Union[pipeline.DataPipeline, pipeline_multimer.DataPipeline, pipeline_custom_templates.DataPipeline],
     model_runners: Dict[str, model.RunModel],
     amber_relaxer: relax.AmberRelaxation,
     benchmark: bool,
-    random_seed: int):
+    random_seed: int,
+    use_custom_templates: bool,
+    template_alignfile: str,
+    msa_mode: str,
+    use_custom_MSA_database: str,
+    MSA_database: str,
+    save_msa_fasta: bool,
+    gen_feats_only: bool,
+    save_template_names: bool,
+    has_gap_chn_brk: bool,
+    substitute_msa: str,
+    msa_for_template_query_seq_only: bool,
+    iptm_interface: str,
+    feature_prefix: str,
+    save_ranked_pdb_only: bool):
   """Predicts structure using AlphaFold for the given sequence."""
   logging.info('Predicting %s', fasta_name)
   timings = {}
@@ -172,19 +232,62 @@
 
   # Get features.
   t_0 = time.time()
+  if use_custom_templates or use_custom_MSA_database!="none":
   feature_dict = data_pipeline.process(
       input_fasta_path=fasta_path,
-      msa_output_dir=msa_output_dir)
+        msa_output_dir=msa_output_dir,
+        use_custom_templates=use_custom_templates,
+        template_alignfile=template_alignfile,
+        msa_mode=msa_mode,
+        use_custom_MSA_database=use_custom_MSA_database,
+        MSA_database=MSA_database,
+        save_msa_fasta=save_msa_fasta,
+        save_template_names=save_template_names,
+        msa_for_template_query_seq_only=msa_for_template_query_seq_only)
+  else:
+    feature_dict = data_pipeline.process(
+      input_fasta_path=fasta_path,
+      msa_output_dir=msa_output_dir,
+      save_msa_fasta=save_msa_fasta,
+      save_template_names=save_template_names,
+      msa_for_template_query_seq_only=msa_for_template_query_seq_only)
   timings['features'] = time.time() - t_0
 
-  # Write out features as a pickled dictionary.
+  if save_msa_fasta:
+    with open(os.path.join(output_dir, "all_msa_feat_gaptoU.fasta"), 'w+') as fh:
+    # fh.write(">query"+"\n"+input_sequence+"\n")
+      counter=1
+      for seq in feature_dict['msa']:
+          seq=[residue_constants.ID_TO_HHBLITS_AA[num] for num in seq]
+          # for x in range(len(seq)):
+          counter+=1
+          fh.write(">seq_"+str(counter)+"\n")
+          out="".join(seq).replace("-","U")
+          fh.write(out+"\n")
+
+  if substitute_msa:
+    keys_substitute=['msa','deletion_matrix','cluster_bias_mask','bert_mask','msa_mask']
+    with open(substitute_msa, 'rb') as fh:
+      substitute_feature_dict = pkl.load(fh)
+    for key in keys_substitute:
+      feature_dict[key]=substitute_feature_dict[key]
+    print("loaded features from %s" % substitute_msa)
+
+  # # Write out features as a pickled dictionary.
+  if not save_ranked_pdb_only:
   features_output_path = os.path.join(output_dir, 'features.pkl')
+    if feature_prefix:
+      features_output_path = os.path.join(output_dir, '%s_features.pkl' % feature_prefix)
   with open(features_output_path, 'wb') as f:
     pickle.dump(feature_dict, f, protocol=4)
 
+  if gen_feats_only:
+    return 
+
   unrelaxed_pdbs = {}
   relaxed_pdbs = {}
   ranking_confidences = {}
+  model_scores = {}
 
   # Run the models.
   num_models = len(model_runners)
@@ -197,15 +300,22 @@
         feature_dict, random_seed=model_random_seed)
     timings[f'process_features_{model_name}'] = time.time() - t_0
 
+
+    interfaces=[]
+    if iptm_interface:
+      interfaces=interface_parser(iptm_interface)
+
     t_0 = time.time()
     prediction_result = model_runner.predict(processed_feature_dict,
-                                             random_seed=model_random_seed)
+                                             random_seed=model_random_seed,
+                                             interfaces=interfaces)
     t_diff = time.time() - t_0
     timings[f'predict_and_compile_{model_name}'] = t_diff
     logging.info(
         'Total JAX model %s on %s predict time (includes compilation time, see --benchmark): %.1fs',
         model_name, fasta_name, t_diff)
 
+
     if benchmark:
       t_0 = time.time()
       model_runner.predict(processed_feature_dict,
@@ -218,8 +328,17 @@
 
     plddt = prediction_result['plddt']
     ranking_confidences[model_name] = prediction_result['ranking_confidence']
+    model_scores[model_name] = [prediction_result['ranking_confidence'], np.mean(prediction_result['plddt'])]
+    if 'iptm' in prediction_result:
+      model_scores[model_name].append(prediction_result['ptm'])
+      model_scores[model_name].append(prediction_result['iptm'])
+    if "custom_iptm" in prediction_result:
+      print("prediction_result['custom_iptm'] is:" + str(prediction_result['custom_iptm']))
+      for score in prediction_result['custom_iptm']:
+        model_scores[model_name].append(score)
 
     # Save the model outputs.
+    if not save_ranked_pdb_only:
     result_output_path = os.path.join(output_dir, f'result_{model_name}.pkl')
     with open(result_output_path, 'wb') as f:
       pickle.dump(prediction_result, f, protocol=4)
@@ -234,7 +353,36 @@
         b_factors=plddt_b_factors,
         remove_leading_feature_dimension=not model_runner.multimer_mode)
 
+    if has_gap_chn_brk:
+      #break chain
+      prev_res=0
+      new_res=0
+      curr_chn=0
+      chn_idx_adj=[]
+      new_res_index=[]
+      for res in unrelaxed_protein.residue_index:
+          if res-prev_res>199:
+              prev_res=res
+              curr_chn+=1
+              chn_idx_adj.append(curr_chn)
+              new_res=1
+          else:
+              prev_res=res
+              chn_idx_adj.append(curr_chn)
+              new_res+=1
+          new_res_index.append(new_res)
+      chain_index=np.add(unrelaxed_protein.chain_index, np.array(chn_idx_adj))
+      unrelaxed_protein= protein.Protein(
+        aatype=unrelaxed_protein.aatype,
+        atom_positions=unrelaxed_protein.atom_positions,
+        atom_mask=unrelaxed_protein.atom_mask,
+        residue_index=np.array(new_res_index,dtype=np.int32),
+        chain_index=np.array(chain_index,dtype=np.int32),
+        b_factors=unrelaxed_protein.b_factors) 
+
     unrelaxed_pdbs[model_name] = protein.to_pdb(unrelaxed_protein)
+
+    if not save_ranked_pdb_only:
     unrelaxed_pdb_path = os.path.join(output_dir, f'unrelaxed_{model_name}.pdb')
     with open(unrelaxed_pdb_path, 'w') as f:
       f.write(unrelaxed_pdbs[model_name])
@@ -248,6 +396,7 @@
       relaxed_pdbs[model_name] = relaxed_pdb_str
 
       # Save the relaxed PDB.
+      if not save_ranked_pdb_only:
       relaxed_output_path = os.path.join(
           output_dir, f'relaxed_{model_name}.pdb')
       with open(relaxed_output_path, 'w') as f:
@@ -265,6 +414,14 @@
       else:
         f.write(unrelaxed_pdbs[model_name])
 
+  model_scores_output=""
+  for model in ranked_order:
+    model_scores_output+="%s\n" % "\t".join(map(str,model_scores[model]))
+  model_scores_output_path = os.path.join(output_dir, f'model_scores.txt')
+  with open(model_scores_output_path, 'w') as f:
+    f.write(model_scores_output)
+
+  if not save_ranked_pdb_only:
   ranking_output_path = os.path.join(output_dir, 'ranking_debug.json')
   with open(ranking_output_path, 'w') as f:
     label = 'iptm+ptm' if 'iptm' in prediction_result else 'plddts'
@@ -272,14 +429,13 @@
         {label: ranking_confidences, 'order': ranked_order}, indent=4))
 
   logging.info('Final timings for %s: %s', fasta_name, timings)
-
+  if not save_ranked_pdb_only:
   timings_output_path = os.path.join(output_dir, 'timings.json')
   with open(timings_output_path, 'w') as f:
     f.write(json.dumps(timings, indent=4))
 
 
 def main(argv):
-
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
 
@@ -288,7 +444,10 @@
     if not FLAGS[f'{tool_name}_binary_path'].value:
       raise ValueError(f'Could not find path to the "{tool_name}" binary. Make '
                        'sure it is installed on your system.')
-  # print(FLAGS.model_preset)
+
+  if FLAGS.gen_feats_only:
+    os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
   use_small_bfd = FLAGS.db_preset == 'reduced_dbs'
   # _check_flag('small_bfd_database_path', 'db_preset',
   #             should_be_set=use_small_bfd)
@@ -339,6 +498,20 @@
         release_dates_path=None,
         obsolete_pdbs_path=FLAGS.obsolete_pdbs_path)
 
+  if FLAGS.use_custom_templates or FLAGS.use_custom_MSA_database!="none":
+    monomer_data_pipeline = pipeline_custom_templates.DataPipeline(
+        jackhmmer_binary_path=FLAGS.jackhmmer_binary_path,
+        hhblits_binary_path=FLAGS.hhblits_binary_path,
+        uniref90_database_path=FLAGS.uniref90_database_path,
+        mgnify_database_path=FLAGS.mgnify_database_path,
+        bfd_database_path=FLAGS.bfd_database_path,
+        uniclust30_database_path=FLAGS.uniclust30_database_path,
+        small_bfd_database_path=FLAGS.small_bfd_database_path,
+        template_searcher=template_searcher,
+        template_featurizer=template_featurizer,
+        use_small_bfd=use_small_bfd,
+        use_precomputed_msas=FLAGS.use_precomputed_msas)
+  else:
   monomer_data_pipeline = pipeline.DataPipeline(
       jackhmmer_binary_path=FLAGS.jackhmmer_binary_path,
       hhblits_binary_path=FLAGS.hhblits_binary_path,
@@ -354,6 +527,13 @@
 
   if run_multimer_system:
     num_predictions_per_model = FLAGS.num_multimer_predictions_per_model
+    if FLAGS.use_custom_templates or FLAGS.use_custom_MSA_database!="none":
+      data_pipeline = pipeline_multimer_custom_templates.DataPipeline(
+          monomer_data_pipeline=monomer_data_pipeline,
+          jackhmmer_binary_path=FLAGS.jackhmmer_binary_path,
+          uniprot_database_path=FLAGS.uniprot_database_path,
+          use_precomputed_msas=FLAGS.use_precomputed_msas)
+    else:
     data_pipeline = pipeline_multimer.DataPipeline(
         monomer_data_pipeline=monomer_data_pipeline,
         jackhmmer_binary_path=FLAGS.jackhmmer_binary_path,
@@ -363,14 +543,25 @@
     num_predictions_per_model = 1
     data_pipeline = monomer_data_pipeline
 
+  num_recycle=FLAGS.num_recycle
+  num_ensemble=FLAGS.num_ensemble
+
   model_runners = {}
   model_names = config.MODEL_PRESETS[FLAGS.model_preset]
   for model_name in model_names:
+    if FLAGS.run_model_names and model_name not in FLAGS.run_model_names.split(","):
+        continue
     model_config = config.model_config(model_name)
     if run_multimer_system:
+      # model_config.model.num_ensemble_eval = num_ensemble
+      model_config.model.num_recycle = num_recycle
+      model_config.model.num_ensemble_train = num_ensemble
       model_config.model.num_ensemble_eval = num_ensemble
     else:
+      model_config.data.common.num_recycle = num_recycle
+      model_config.model.num_recycle = num_recycle
       model_config.data.eval.num_ensemble = num_ensemble
+    
     model_params = data.get_model_haiku_params(
         model_name=model_name, data_dir=FLAGS.data_dir)
     model_runner = model.RunModel(model_config, model_params)
@@ -380,7 +571,7 @@
   logging.info('Have %d models: %s', len(model_runners),
                list(model_runners.keys()))
 
-  if FLAGS.run_relax:
+  if FLAGS.run_relax==True:
     amber_relaxer = relax.AmberRelaxation(
         max_iterations=RELAX_MAX_ITERATIONS,
         tolerance=RELAX_ENERGY_TOLERANCE,
@@ -407,7 +598,21 @@
         model_runners=model_runners,
         amber_relaxer=amber_relaxer,
         benchmark=FLAGS.benchmark,
-        random_seed=random_seed)
+        random_seed=random_seed,
+        use_custom_templates=FLAGS.use_custom_templates,
+        template_alignfile=FLAGS.template_alignfile,
+        msa_mode=FLAGS.msa_mode,
+        use_custom_MSA_database=FLAGS.use_custom_MSA_database,
+        MSA_database=FLAGS.MSA_database,
+        save_msa_fasta=FLAGS.save_msa_fasta,
+        gen_feats_only=FLAGS.gen_feats_only,
+        save_template_names=FLAGS.save_template_names,
+        has_gap_chn_brk=FLAGS.has_gap_chn_brk,
+        substitute_msa=FLAGS.substitute_msa,
+        msa_for_template_query_seq_only=FLAGS.msa_for_template_query_seq_only,
+        iptm_interface=FLAGS.iptm_interface,
+        feature_prefix=FLAGS.feature_prefix,
+        save_ranked_pdb_only=FLAGS.save_ranked_pdb_only)
 
 
 if __name__ == '__main__':