add amber minimize

fastfold/relax/amber_minimize.py

@@ -94,8 +94,8 @@ def _openmm_minimize(
         _add_restraints(system, pdb, stiffness, restraint_set, exclude_residues)
 
     integrator = openmm.LangevinIntegrator(0, 0.01, 0.0)
-    platform = openmm.Platform.getPlatformByName("CUDA" if use_gpu else "CPU")
-    simulation = openmm_app.Simulation(pdb.topology, system, integrator, platform)
+    # platform = openmm.Platform.getPlatformByName("CUDA" if use_gpu else "CPU")
+    simulation = openmm_app.Simulation(pdb.topology, system, integrator)
     simulation.context.setPositions(pdb.positions)
 
     ret = {}
@@ -536,7 +536,6 @@ def get_initial_energies(
         openmm_pdbs[0].topology,
         system,
         openmm.LangevinIntegrator(0, 0.01, 0.0),
-        openmm.Platform.getPlatformByName("CPU"),
     )
     energies = []
     for pdb in openmm_pdbs:

inference.py

@@ -454,21 +454,21 @@ 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}")
-
-    # # 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)
+    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}")
+
+    # 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(