From 7e06ed9821926dc6ef5bae54048c813b331b3b58 Mon Sep 17 00:00:00 2001
From: Jennifer Wei <jennifer.wei@omsf.io>
Date: Thu, 24 Apr 2025 21:33:53 +0000
Subject: [PATCH] support openmm>8 and fix tolerance units in amber
 minimization

---
 environment.yml                     | 2 +-
 openfold/config.py                  | 2 +-
 openfold/np/relax/amber_minimize.py | 3 ++-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/environment.yml b/environment.yml
index bb510c2..1cc5867 100644
--- a/environment.yml
+++ b/environment.yml
@@ -8,7 +8,7 @@ dependencies:
   - python=3.10
   - setuptools=59.5.0
   - pip
-  - openmm=7.7
+  - openmm
   - pdbfixer
   - pytorch-lightning
   - biopython
diff --git a/openfold/config.py b/openfold/config.py
index 7bf30e3..5487f08 100644
--- a/openfold/config.py
+++ b/openfold/config.py
@@ -660,7 +660,7 @@ config = mlc.ConfigDict(
         },
         "relax": {
             "max_iterations": 0,  # no max
-            "tolerance": 2.39,
+            "tolerance": 10.0, 
             "stiffness": 10.0,
             "max_outer_iterations": 20,
             "exclude_residues": [],
diff --git a/openfold/np/relax/amber_minimize.py b/openfold/np/relax/amber_minimize.py
index 02816bb..c331778 100644
--- a/openfold/np/relax/amber_minimize.py
+++ b/openfold/np/relax/amber_minimize.py
@@ -34,6 +34,7 @@ from openmm import app as openmm_app
 from openmm.app.internal.pdbstructure import PdbStructure
 
 ENERGY = unit.kilocalories_per_mole
+FORCE = unit.kilojoules_per_mole / unit.nanometer 
 LENGTH = unit.angstroms
 
 
@@ -439,7 +440,7 @@ def _run_one_iteration(
     exclude_residues = exclude_residues or []
 
     # Assign physical dimensions.
-    tolerance = tolerance * ENERGY
+    tolerance = tolerance * FORCE 
     stiffness = stiffness * ENERGY / (LENGTH ** 2)
 
     start = time.perf_counter()