support openmm >= 7.6

openfold/np/relax/amber_minimize.py

@@ -28,10 +28,18 @@ import openfold.utils.loss as loss
 from openfold.np.relax import cleanup, utils
 import ml_collections
 import numpy as np
-from simtk import openmm
-from simtk import unit
-from simtk.openmm import app as openmm_app
-from simtk.openmm.app.internal.pdbstructure import PdbStructure
+try:
+    # openmm >= 7.6
+    import openmm
+    from openmm import unit
+    from openmm import app as openmm_app
+    from openmm.app.internal.pdbstructure import PdbStructure
+except ImportError:
+    # openmm < 7.6 (requires DeepMind patch)
+    from simtk import openmm
+    from simtk import unit
+    from simtk.openmm import app as openmm_app
+    from simtk.openmm.app.internal.pdbstructure import PdbStructure
 
 ENERGY = unit.kilocalories_per_mole
 LENGTH = unit.angstroms

openfold/np/relax/cleanup.py

@@ -20,8 +20,14 @@ cases like removing chains of length one (see clean_structure).
 import io
 
 import pdbfixer
-from simtk.openmm import app
-from simtk.openmm.app import element
+try:
+    # openmm >= 7.6
+    from openmm import app
+    from openmm.app import element
+except ImportError:
+    # openmm < 7.6 (requires DeepMind patch)
+    from simtk.openmm import app
+    from simtk.openmm.app import element
 
 
 def fix_pdb(pdbfile, alterations_info):

openfold/np/relax/utils.py

@@ -18,8 +18,14 @@ import io
 from openfold.np import residue_constants
 from Bio import PDB
 import numpy as np
-from simtk.openmm import app as openmm_app
-from simtk.openmm.app.internal.pdbstructure import PdbStructure
+try:
+    # openmm >= 7.6
+    from openmm import app as openmm_app
+    from openmm.app.internal.pdbstructure import PdbStructure
+except ImportError:
+    # openmm < 7.6 (requires DeepMind patch)
+    from simtk.openmm import app as openmm_app
+    from simtk.openmm.app.internal.pdbstructure import PdbStructure
 
 
 def overwrite_pdb_coordinates(pdb_str: str, pos) -> str: