initial fastfold

inference.py

@@ -227,9 +227,11 @@ def inference_multimer_model(args):
     )
 
     output_dir_base = args.output_dir
+    
     random_seed = args.data_random_seed
     if random_seed is None:
         random_seed = random.randrange(sys.maxsize)
+    # seed_torch(seed=1029)
     
     feature_processor = feature_pipeline.FeaturePipeline(
         config.data
@@ -312,21 +314,31 @@ def inference_multimer_model(args):
     with open(unrelaxed_output_path, 'w') as f:
         f.write(protein.to_pdb(unrelaxed_protein))
 
-    amber_relaxer = relax.AmberRelaxation(
-        use_gpu=True,
-        **config.relax,
-    )
+    if(args.relaxation):
+        amber_relaxer = relax.AmberRelaxation(
+            use_gpu=True,
+            **config.relax,
+        )
 
-    # Relax the prediction.
-    t = time.perf_counter()
-    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
-    print(f"Relaxation time: {time.perf_counter() - t}")
+        # Relax the prediction.
+        t = time.perf_counter()
+        relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
+        print(f"Relaxation time: {time.perf_counter() - t}")
+
+        # Save the relaxed PDB.
+        relaxed_output_path = os.path.join(args.output_dir,
+                                            f'{tag}_{args.model_name}_relaxed.pdb')
+        with open(relaxed_output_path, 'w') as f:
+            f.write(relaxed_pdb_str)
+        
+    if(args.save_outputs):
+            output_dict_path = os.path.join(
+                args.output_dir, f'{tag}_{args.model_name}_output_dict.pkl'
+            )
+            with open(output_dict_path, "wb") as fp:
+                pickle.dump(out, fp, protocol=pickle.HIGHEST_PROTOCOL)
 
-    # Save the relaxed PDB.
-    relaxed_output_path = os.path.join(args.output_dir,
-                                        f'{tag}_{args.model_name}_relaxed.pdb')
-    with open(relaxed_output_path, 'w') as f:
-        f.write(relaxed_pdb_str)
+            logger.info(f"Model output written to {output_dict_path}...")
 
 
 def inference_monomer_model(args):
@@ -454,21 +466,22 @@ def inference_monomer_model(args):
     with open(unrelaxed_output_path, 'w') as f:
         f.write(protein.to_pdb(unrelaxed_protein))
 
-    amber_relaxer = relax.AmberRelaxation(
-        use_gpu=True,
-        **config.relax,
-    )
-
-    # Relax the prediction.
-    t = time.perf_counter()
-    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
-    print(f"Relaxation time: {time.perf_counter() - t}")
+    if(args.relaxation):
+        amber_relaxer = relax.AmberRelaxation(
+            use_gpu=True,
+            **config.relax,
+        )
 
-    # Save the relaxed PDB.
-    relaxed_output_path = os.path.join(args.output_dir,
-                                        f'{tag}_{args.model_name}_relaxed.pdb')
-    with open(relaxed_output_path, 'w') as f:
-        f.write(relaxed_pdb_str)
+        # Relax the prediction.
+        t = time.perf_counter()
+        relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
+        print(f"Relaxation time: {time.perf_counter() - t}")
+
+        # Save the relaxed PDB.
+        relaxed_output_path = os.path.join(args.output_dir,
+                                            f'{tag}_{args.model_name}_relaxed.pdb')
+        with open(relaxed_output_path, 'w') as f:
+            f.write(relaxed_pdb_str)
     
     if(args.save_outputs):
             output_dict_path = os.path.join(
@@ -512,6 +525,9 @@ if __name__ == "__main__":
                         help="""Path to model parameters. If None, parameters are selected
              automatically according to the model name from 
              ./data/params""")
+    parser.add_argument(
+        "--relaxation", action="store_false", default=False,
+    )
     parser.add_argument(
         "--save_outputs", action="store_true", default=False,
         help="Whether to save all model outputs, including embeddings, etc."