Merge pull request #6 from ntBre/espaloma

Allow using espaloma as a force field

devtools/conda-envs/dev.yaml

@@ -9,6 +9,7 @@ dependencies:
   - openff-qcsubmit
   - openmmforcefields
   - smirnoff-plugins =2023.08.0
+  - espaloma
 
   - ipython
   - ipdb

ibstore/_forcefields.py

@@ -63,3 +63,34 @@ def _gaff(molecule: Molecule, force_field_name: str) -> openmm.System:
         nonbondedCutoff=0.9 * openmm.unit.nanometer,
         constraints=None,
     )
+
+
+def _espaloma(molecule: Molecule, force_field_name: str) -> openmm.System:
+    """Generate an OpenMM System for a molecule and force field name. The force
+    field name should be of the form espaloma-force-field-name, such as
+    espaloma-openff_unconstrained-2.1.0. Everything after the first dash is
+    passed as the forcefield argument to
+    espaloma.graphs.deploy.openmm_system_from_graph, where it will be appended
+    with .offxml. Raises a ValueError if there is no dash in force_field_name.
+    """
+    import espaloma as esp
+
+    if not force_field_name.startswith("espaloma"):
+        raise NotImplementedError(
+            f"Force field {force_field_name} not implemented."
+        )
+
+    ff = force_field_name.split("-", 1)[1:2]
+
+    if ff == []:
+        raise ValueError(
+            "espaloma force field must have an OpenFF force field too"
+        )
+    else:
+        ff = ff[0]
+
+    mol_graph = esp.Graph(molecule)
+    model = esp.get_model("latest")
+    model(mol_graph.heterograph)
+
+    return esp.graphs.deploy.openmm_system_from_graph(mol_graph, forcefield=ff)

ibstore/_minimize.py

@@ -114,6 +114,14 @@ def _run_openmm(
             force_field_name=input.force_field,
         )
 
+    elif input.force_field.startswith("espaloma"):
+        from ibstore._forcefields import _espaloma
+
+        system = _espaloma(
+            molecule=molecule,
+            force_field_name=input.force_field,
+        )
+
     else:
         try:
             force_field = FORCE_FIELDS[input.force_field]

ibstore/_tests/unit_tests/test_forcefields.py

@@ -2,7 +2,7 @@
 import pytest
 from openff.toolkit import Molecule
 
-from ibstore._forcefields import _gaff, _smirnoff
+from ibstore._forcefields import _espaloma, _gaff, _smirnoff
 
 
 @pytest.fixture()
@@ -27,3 +27,15 @@ def test_gaff_basic(molecule):
 def test_gaff_unsupported(molecule):
     with pytest.raises(NotImplementedError):
         _gaff(molecule, "foo")
+
+
+def test_espaloma_basic(molecule):
+    system = _espaloma(molecule, "espaloma-openff_unconstrained-2.1.0")
+
+    assert isinstance(system, openmm.System)
+    assert system.getNumParticles() == molecule.n_atoms
+
+
+def test_espaloma_unsupported(molecule):
+    with pytest.raises(NotImplementedError):
+        _espaloma(molecule, "foo")