Fix labelling molecules with virtual sites

devtools/conda-envs/test_env.yaml

@@ -12,7 +12,7 @@ dependencies:
   - codecov
 
     # Standard dependencies
-  - openff-toolkit >=0.10.0
+  - openff-toolkit >=0.10.6
   - openff-units >=0.1.4
   - smirnoff99frosst
   - numpy

docs/releasehistory.rst

@@ -8,6 +8,14 @@ Releases follow the ``major.minor.micro`` scheme recommended by
 * ``minor`` increments add features but do not break API compatibility
 * ``micro`` increments represent bugfix releases or improvements in documentation
 
+0.3.10
+-----
+
+Bugfixes
+""""""""
+
+* PR `#444 <https://github.com/openforcefield/openff-evaluator/pull/444>`_: Fix labelling molecules with virtual sites
+
 0.3.9
 -----
 

openff/evaluator/tests/test_utils/test_openmm.py

@@ -215,7 +215,7 @@ def hydrogen_chloride_force_field(
                 "smirks": "[#1:1]-[#17:2]",
                 "type": "BondCharge",
                 "distance": -0.2 * openmm_unit.nanometers,
-                "match": "once",
+                "match": "all_permutations",
                 "charge_increment1": 0.0 * openmm_unit.elementary_charge,
                 "charge_increment2": 0.0 * openmm_unit.elementary_charge,
             }
@@ -403,7 +403,7 @@ def test_update_context_with_positions(box_vectors):
 
     force_field = hydrogen_chloride_force_field(True, False, True)
 
-    topology: Topology = Molecule.from_smiles("Cl").to_topology()
+    topology: Topology = Molecule.from_mapped_smiles("[Cl:1][H:2]").to_topology()
     system = force_field.create_openmm_system(topology)
 
     context = openmm.Context(
@@ -419,7 +419,7 @@ def test_update_context_with_positions(box_vectors):
 
     assert numpy.allclose(
         context_positions.value_in_unit(openmm_unit.angstrom),
-        numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]),
+        numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]),
     )
 
     assert numpy.isclose(
@@ -437,7 +437,7 @@ def test_update_context_with_pdb(tmpdir):
 
     force_field = hydrogen_chloride_force_field(True, False, True)
 
-    topology: Topology = Molecule.from_smiles("Cl").to_topology()
+    topology: Topology = Molecule.from_mapped_smiles("[Cl:1][H:2]").to_topology()
     system = force_field.create_openmm_system(topology)
 
     context = openmm.Context(
@@ -458,14 +458,14 @@ def test_update_context_with_pdb(tmpdir):
 
     assert numpy.allclose(
         context_positions.value_in_unit(openmm_unit.angstrom),
-        numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]),
+        numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [-1.0, 0.0, 0.0]]),
     )
 
     assert numpy.allclose(
         extract_positions(context.getState(getPositions=True), [2]).value_in_unit(
             openmm_unit.angstrom
         ),
-        numpy.array([[2.0, 0.0, 0.0]]),
+        numpy.array([[-1.0, 0.0, 0.0]]),
     )
 
     assert numpy.isclose(context_box_vectors[0].x, 2.0)

openff/evaluator/tests/test_workflow/test_workflow.py

@@ -7,7 +7,7 @@
 
 import numpy
 import pytest
-from openff.toolkit.typing.engines.smirnoff import ForceField
+from openff.toolkit.typing.engines.smirnoff import ForceField, VirtualSiteHandler
 from openff.units import unit
 
 from openff.evaluator.attributes import UNDEFINED
@@ -299,13 +299,28 @@ def test_find_relevant_gradient_keys(tmpdir):
             "sigma": 1.0 * openmm_unit.angstrom,
         }
     )
+    vsite_handler = VirtualSiteHandler(version=0.3)
+    vsite_handler.add_parameter(
+        {
+            "smirks": "[#1:1][#17:2]",
+            "type": "BondCharge",
+            "distance": 0.1 * openmm_unit.nanometers,
+            "match": "all_permutations",
+            "charge_increment1": 0.0 * openmm_unit.elementary_charge,
+            "charge_increment2": 0.0 * openmm_unit.elementary_charge,
+        }
+    )
+    force_field.register_parameter_handler(vsite_handler)
 
     force_field_path = os.path.join(tmpdir, "ff.json")
     SmirnoffForceFieldSource.from_object(force_field).json(force_field_path)
 
     expected_gradient_keys = {
         ParameterGradientKey(tag="vdW", smirks=None, attribute="scale14"),
         ParameterGradientKey(tag="vdW", smirks="[#1:1]", attribute="epsilon"),
+        ParameterGradientKey(
+            tag="VirtualSites", smirks="[#1:1][#17:2]", attribute="distance"
+        ),
     }
 
     gradient_keys = Workflow._find_relevant_gradient_keys(

openff/evaluator/workflow/workflow.py

@@ -545,8 +545,39 @@ def replace_protocol(self, old_protocol, new_protocol, update_paths_only=False):
                 attribute_value.replace_protocol(old_protocol, new_protocol)
 
     @staticmethod
+    def label_molecules(force_field, topology):
+
+        from openff.toolkit.topology import Topology
+        from openff.toolkit.typing.engines.smirnoff.parameters import VirtualSiteHandler
+
+        molecule_labels = list()
+
+        for molecule_idx, molecule in enumerate(topology.reference_molecules):
+            top_mol = Topology.from_molecules([molecule])
+            current_molecule_labels = dict()
+            param_is_list = False
+            for tag, parameter_handler in force_field._parameter_handlers.items():
+
+                if type(parameter_handler) == VirtualSiteHandler:
+                    param_is_list = True
+
+                matches = parameter_handler.find_matches(top_mol, unique=True)
+
+                if param_is_list:
+                    parameter_matches = matches
+                else:
+                    parameter_matches = matches.__class__()
+                    for match in matches:
+                        parameter_matches[match] = matches[match].parameter_type
+
+                current_molecule_labels[tag] = parameter_matches
+
+            molecule_labels.append(current_molecule_labels)
+        return molecule_labels
+
+    @classmethod
     def _find_relevant_gradient_keys(
-        substance, force_field_path, parameter_gradient_keys
+        cls, substance, force_field_path, parameter_gradient_keys
     ):
         """Extract only those keys which may be applied to the
         given substance.
@@ -585,7 +616,7 @@ def _find_relevant_gradient_keys(
             all_molecules.append(Molecule.from_smiles(component.smiles))
 
         topology = Topology.from_molecules(all_molecules)
-        labelled_molecules = force_field.label_molecules(topology)
+        labelled_molecules = cls.label_molecules(force_field, topology)
 
         reduced_parameter_keys = []
 
@@ -601,7 +632,16 @@ def _find_relevant_gradient_keys(
 
                 contains_parameter = False
 
-                for parameter in labelled_molecule[parameter_key.tag].store.values():
+                labelled_parameters = (
+                    [
+                        match.parameter_type
+                        for match in labelled_molecule[parameter_key.tag]
+                    ]
+                    if isinstance(labelled_molecule[parameter_key.tag], list)
+                    else [*labelled_molecule[parameter_key.tag].store.values()]
+                )
+
+                for parameter in labelled_parameters:
 
                     if (
                         parameter_key.smirks is not None