diff --git a/devtools/conda-envs/dev.yaml b/devtools/conda-envs/dev.yaml index d4f4f59..26deb59 100644 --- a/devtools/conda-envs/dev.yaml +++ b/devtools/conda-envs/dev.yaml @@ -9,6 +9,7 @@ dependencies: - openff-qcsubmit - openmmforcefields - smirnoff-plugins =2023.08.0 + - espaloma - ipython - ipdb diff --git a/ibstore/_forcefields.py b/ibstore/_forcefields.py index ef24392..5d609e5 100644 --- a/ibstore/_forcefields.py +++ b/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) diff --git a/ibstore/_minimize.py b/ibstore/_minimize.py index cfd8103..6064d3b 100644 --- a/ibstore/_minimize.py +++ b/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] diff --git a/ibstore/_tests/unit_tests/test_forcefields.py b/ibstore/_tests/unit_tests/test_forcefields.py index ec505b6..8036bdd 100644 --- a/ibstore/_tests/unit_tests/test_forcefields.py +++ b/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")