Add missing directory AFfine

Author: AmirAsgary

AFfine/af2_util.py

@@ -242,6 +242,11 @@ def get_atom_positions_from_pdb(pdb_file_path: str, aligned_sequences: tuple[str
     if pep_len and anchors:
         mhc_len = all_positions.shape[0] - pep_len
         full_anchors = [i - 1 + mhc_len for i in anchors] #(2-1) + 180 = 181
+        full_anchors = sorted(full_anchors)
+        if len(full_anchors) == 2: # in case of MHC-I distance between two anchors is long and distrubs the folding, therefore, we add two other positions
+            full_anchors = [full_anchors[0], full_anchors[0]+2,  full_anchors[1]-2, full_anchors[1]]
+            print('mhc_1 anchor initial guess for positions:', full_anchors)
+            all_positions *= 0. # zero for mhc 1
         core_region = [i for i in range(mhc_len, mhc_len + pep_len) if i not in full_anchors] # (i in (180, 189) if not anchor)
         mask = np.ones([num_res_query, residue_constants.atom_type_num, 3], dtype=np.float32)
         mask[core_region] = 0.

AFfine/alphafold/model/model.py

@@ -51,7 +51,8 @@ class RunModel:
 
   def __init__(self,
                config: ml_collections.ConfigDict,
-               params: Optional[Mapping[str, Mapping[str, np.ndarray]]] = None):
+               params: Optional[Mapping[str, Mapping[str, np.ndarray]]] = None,
+               return_representations: bool = False): # added by https://github.com/AmirAsgary
     self.config = config
     self.params = params
 
@@ -73,7 +74,10 @@ def _forward_fn(batch,
     self.apply = jax.jit(hk.transform(partial(_forward_fn, is_training=True, compute_loss=True)).apply)
     self.init = jax.jit(hk.transform(partial(_forward_fn, is_training=True, compute_loss=False)).init)
     self.apply_infer = jax.jit(hk.transform(partial(_forward_fn, is_training=False, compute_loss=True)).apply)
-    self.apply_predict = jax.jit(hk.transform(partial(_forward_fn, is_training=False, compute_loss=False)).apply)
+    #self.apply_predict = jax.jit(hk.transform(partial(_forward_fn, is_training=False, compute_loss=False)).apply)
+    self.apply_predict = jax.jit(hk.transform(partial(_forward_fn, is_training=False, compute_loss=False, 
+        return_representations=return_representations)).apply, static_argnames=("return_representations",) # added by https://github.com/AmirAsgary
+      )
   def init_params(self, feat: features.FeatureDict, random_seed: int = 0):
     """Initializes the model parameters.
 

AFfine/predict_utils.py

@@ -222,9 +222,7 @@ def predict_structure(
 
     metric_tags = 'plddt ptm predicted_aligned_error'.split()
     if return_all_outputs:
-        metric_tags = ['distogram', 'masked_msa', 'predicted_aligned_error', 'predicted_lddt',
-                       'structure_module', 'plddt', 'aligned_confidence_probs',
-                       'max_predicted_aligned_error', 'ptm', 'representations']
+        metric_tags = 'plddt ptm predicted_aligned_error representations'.split() # added by Amir
     all_metrics = {} # eventual return value
 
     metrics = {} # stupid duplication
@@ -247,12 +245,11 @@ def predict_structure(
                                                    template_sequence=None, aln=aln, anchors=anchors,
                                                    peptide_seq=peptide_seq) #added by https://github.com/AmirAsgary/
             prediction_result = model_runner.predict(processed_feature_dict, initial_guess,
-                                                     return_representations=return_representations) #added by https://github.com/AmirAsgary/
+                                                     return_representations=return_all_outputs) #added by https://github.com/AmirAsgary/
         else:
-            prediction_result = model_runner.predict(processed_feature_dict, return_representations=return_representations) #added by https://github.com/AmirAsgary/
+            prediction_result = model_runner.predict(processed_feature_dict, return_representations=return_all_outputs) #added by https://github.com/AmirAsgary/
 
         ###
-        print('DEBUUUUUUUUUUUUG: \n',prediction_result.keys(), '\n ----------------------------')
         unrelaxed_protein = protein.from_prediction(
             processed_feature_dict, prediction_result)
         unrelaxed_pdb_lines.append(protein.to_pdb(unrelaxed_protein))
@@ -282,13 +279,16 @@ def predict_structure(
 
 
         #plddts_ranked[f"model_{n+1}"] = plddts[r]
-
         if dump_metrics:
             metrics_prefix = f'{prefix}_model_{n+1}_{model_names[r]}'
             for tag in metric_tags:
                 m = metrics[tag][r]
                 if m is not None:
-                    np.save(f'{metrics_prefix}_{tag}.npy', m)
+                    if tag != 'representations': # added by Amir
+                        np.save(f'{metrics_prefix}_{tag}.npy', m)
+                    else:
+                        with open(f'{metrics_prefix}_{tag}.pkl', 'wb') as f:
+                            pickle.dump(m, f)
 
     return all_metrics
 
@@ -348,7 +348,7 @@ def load_model_runners(
                 model_name=model_name, data_dir=data_dir)
 
         model_runners[model_name] = model.RunModel(
-            model_config, model_params)
+            model_config, model_params, args.return_all_outputs) #return_all_outputs added by https://github.com/AmirAsgary
     return model_runners
 
 

AFfine/run_finetuning.py

@@ -74,7 +74,7 @@
                     'data.get_model_haiku_params; should contain params/ subfolder')
 
 flags.DEFINE_bool('only_fit_binder', False, help='if True, dont fit alphafold params')
-flags.DEFINE_bool('freeze_binder', False, help='if True, dont fit binder params')
+flags.DEFINE_bool('freeze_binder', False, help='if True, dont fit binder params') # binder means --> logistic regression part
 flags.DEFINE_bool('freeze_everything', False, help='if True, dont fit anything')
 flags.DEFINE_bool('no_ramp', False, help='if True, dont ramp')
 flags.DEFINE_bool('no_valid', False, help='if True, dont compute valid stats')

AFfine/run_prediction.py

@@ -70,6 +70,7 @@
 
 parser.add_argument('--no_initial_guess', action='store_true', default=False, help='When active, no intial guess is used to direct modeling and only template is used.')
 parser.add_argument('--return_all_outputs', action='store_true', default=False, help='Save all alphafold outputs including evoformer output')
+parser.add_argument('--use_msa', action='store_true', default=False, help='If Enabled, use MSA for prediction. If not, only template is used.')
 args = parser.parse_args()
 
 import os
@@ -117,6 +118,9 @@
         template_pdb_dict = targetl.template_pdb_dict
         with open(template_pdb_dict, 'r') as f:
             template_pdb_dict = json.load(f)
+    else: # added by Amir for initial guess condition and getting dict from input tsv
+        template_pdb_dict = None
+            
 
     print(alignfile)
     assert exists(alignfile)
@@ -173,11 +177,21 @@
         )
         template_features_list.append(template_features)
 
+
+
     all_template_features = predict_utils.compile_template_features(
         template_features_list)
 
-    msa=[query_sequence]
-    deletion_matrix=[[0]*len(query_sequence)]
+    if not args.use_msa: # added by Amir for using msa or not
+        # if we are not using MSA, we need to set the deletion matrix
+        msa=[query_sequence]
+        deletion_matrix=[[0]*len(query_sequence)]
+    else: # added by Amir for using msa or not
+        # generate arbeitary msa from input sequence
+        msa = [query_sequence] + msa
+        
+
+   
 
     all_metrics = predict_utils.run_alphafold_prediction(
         query_sequence=query_sequence,