Fix v-site positions not set by OpenMM (#389)

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.7
+-----
+
+Bugfixes
+""""""""
+
+* PR `#389 <https://github.com/openforcefield/openff-evaluator/pull/389>`_: Fix v-site positions not set by OpenMM
+
 0.3.6
 -----
 

openff/evaluator/protocols/openmm.py

@@ -34,6 +34,8 @@
     pint_quantity_to_openmm,
     setup_platform_with_resources,
     system_subset,
+    update_context_with_pdb,
+    update_context_with_positions,
 )
 from openff.evaluator.utils.serialization import TypedJSONDecoder, TypedJSONEncoder
 from openff.evaluator.utils.utils import is_file_and_not_empty
@@ -99,12 +101,12 @@ def _evaluate_energies(
     for frame_index in range(trajectory.n_frames):
 
         positions = trajectory.xyz[frame_index]
+        box_vectors = None
 
         if enable_pbc:
             box_vectors = trajectory.openmm_boxes(frame_index)
-            openmm_context.setPeriodicBoxVectors(*box_vectors)
 
-        openmm_context.setPositions(positions)
+        update_context_with_positions(openmm_context, positions, box_vectors)
 
         state = openmm_context.getState(getEnergy=True)
 
@@ -348,13 +350,7 @@ def _execute(self, directory, available_resources):
             input_pdb_file.topology, system, integrator, platform
         )
 
-        box_vectors = input_pdb_file.topology.getPeriodicBoxVectors()
-
-        if box_vectors is None:
-            box_vectors = simulation.system.getDefaultPeriodicBoxVectors()
-
-        simulation.context.setPeriodicBoxVectors(*box_vectors)
-        simulation.context.setPositions(input_pdb_file.positions)
+        update_context_with_pdb(simulation.context, input_pdb_file)
 
         simulation.minimizeEnergy(
             pint_quantity_to_openmm(self.tolerance), self.max_iterations
@@ -605,6 +601,7 @@ def _setup_simulation_objects(self, temperature, pressure, available_resources):
 
         # Initialize the context with the correct positions etc.
         input_pdb_file = app.PDBFile(self.input_coordinate_file)
+        box_vectors = None
 
         if self.enable_pbc:
 
@@ -617,9 +614,10 @@ def _setup_simulation_objects(self, temperature, pressure, available_resources):
                     "The input file must contain box vectors when running with PBC."
                 )
 
-            context.setPeriodicBoxVectors(*box_vectors)
+        update_context_with_positions(
+            context, input_pdb_file.getPositions(asNumpy=True), box_vectors
+        )
 
-        context.setPositions(input_pdb_file.positions)
         context.setVelocitiesToTemperature(temperature)
 
         return context, integrator

openff/evaluator/tests/test_utils/test_openmm.py

@@ -1,17 +1,21 @@
+import os.path
 from random import randint, random
 
 import mdtraj
 import numpy
 import numpy as np
+import openmm
 import pytest
 from openff.toolkit.topology import Molecule, Topology
 from openff.toolkit.typing.engines.smirnoff import ForceField, vdWHandler
 from openff.toolkit.typing.engines.smirnoff.parameters import (
     ChargeIncrementModelHandler,
     ElectrostaticsHandler,
     LibraryChargeHandler,
+    VirtualSiteHandler,
 )
 from openmm import unit as openmm_unit
+from openmm.app import PDBFile
 
 from openff.evaluator import unit
 from openff.evaluator.backends import ComputeResources
@@ -24,6 +28,8 @@
     openmm_quantity_to_pint,
     pint_quantity_to_openmm,
     system_subset,
+    update_context_with_pdb,
+    update_context_with_positions,
 )
 
 
@@ -128,7 +134,9 @@ def test_constants():
 
 
 def hydrogen_chloride_force_field(
-    library_charge: bool, charge_increment: bool
+    library_charge: bool,
+    charge_increment: bool,
+    vsite: bool,
 ) -> ForceField:
     """Returns a SMIRNOFF force field which is able to parameterize hydrogen chloride."""
 
@@ -191,6 +199,21 @@ def hydrogen_chloride_force_field(
         )
         force_field.register_parameter_handler(charge_increment_handler)
 
+    if vsite:
+
+        vsite_handler = VirtualSiteHandler(version=0.3)
+        vsite_handler.add_parameter(
+            {
+                "smirks": "[#1:1]-[#17:2]",
+                "type": "BondCharge",
+                "distance": -0.2 * openmm_unit.nanometers,
+                "match": "once",
+                "charge_increment1": 0.0 * openmm_unit.elementary_charge,
+                "charge_increment2": 0.0 * openmm_unit.elementary_charge,
+            }
+        )
+        force_field.register_parameter_handler(vsite_handler)
+
     return force_field
 
 
@@ -202,7 +225,7 @@ def test_system_subset_vdw():
     # Create the system subset.
     system, parameter_value = system_subset(
         parameter_key=ParameterGradientKey("vdW", "[#1:1]", "epsilon"),
-        force_field=hydrogen_chloride_force_field(True, True),
+        force_field=hydrogen_chloride_force_field(True, True, False),
         topology=topology,
         scale_amount=0.5,
     )
@@ -233,7 +256,7 @@ def test_system_subset_vdw_cutoff():
     # Create the system subset.
     system, parameter_value = system_subset(
         parameter_key=ParameterGradientKey("vdW", None, "cutoff"),
-        force_field=hydrogen_chloride_force_field(True, True),
+        force_field=hydrogen_chloride_force_field(True, True, False),
         topology=topology,
         scale_amount=0.5,
     )
@@ -247,7 +270,7 @@ def test_system_subset_vdw_cutoff():
 
 def test_system_subset_library_charge():
 
-    force_field = hydrogen_chloride_force_field(True, False)
+    force_field = hydrogen_chloride_force_field(True, False, False)
 
     # Ensure a zero charge after perturbation.
     force_field.get_parameter_handler("LibraryCharges").parameters["[#1:1]"].charge1 = (
@@ -296,7 +319,7 @@ def test_system_subset_charge_increment():
         parameter_key=ParameterGradientKey(
             "ChargeIncrementModel", "[#1:1]-[#17:2]", "charge_increment1"
         ),
-        force_field=hydrogen_chloride_force_field(False, True),
+        force_field=hydrogen_chloride_force_field(False, True, False),
         topology=topology,
         scale_amount=0.5,
     )
@@ -363,3 +386,73 @@ def test_compute_gradients(tmpdir, smirks, all_zeros):
             observables["PotentialEnergy"].gradients[0].value,
             0.0 * observables["PotentialEnergy"].gradients[0].value.units,
         )
+
+
+@pytest.mark.parametrize(
+    "box_vectors", [None, (numpy.eye(3) * 3.0) * openmm_unit.nanometers]
+)
+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()
+    system = force_field.create_openmm_system(topology)
+
+    context = openmm.Context(
+        system, openmm.VerletIntegrator(0.1 * openmm_unit.femtoseconds)
+    )
+
+    positions = numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) * openmm_unit.angstrom
+
+    update_context_with_positions(context, positions, box_vectors)
+
+    context_positions = context.getState(getPositions=True).getPositions(asNumpy=True)
+    context_box_vectors = context.getState(getPositions=True).getPeriodicBoxVectors()
+
+    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]]),
+    )
+
+    assert numpy.isclose(
+        context_box_vectors[0].x, (2.0 if box_vectors is None else 3.0)
+    )
+    assert numpy.isclose(
+        context_box_vectors[1].y, (2.0 if box_vectors is None else 3.0)
+    )
+    assert numpy.isclose(
+        context_box_vectors[2].z, (2.0 if box_vectors is None else 3.0)
+    )
+
+
+def test_update_context_with_pdb(tmpdir):
+
+    force_field = hydrogen_chloride_force_field(True, False, True)
+
+    topology: Topology = Molecule.from_smiles("Cl").to_topology()
+    system = force_field.create_openmm_system(topology)
+
+    context = openmm.Context(
+        system, openmm.VerletIntegrator(0.1 * openmm_unit.femtoseconds)
+    )
+
+    positions = numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) * openmm_unit.angstrom
+
+    pdb_path = os.path.join(tmpdir, "tmp.pdb")
+    topology.to_file(pdb_path, positions)
+
+    pdb_file = PDBFile(pdb_path)
+
+    update_context_with_pdb(context, pdb_file)
+
+    context_positions = context.getState(getPositions=True).getPositions(asNumpy=True)
+    context_box_vectors = context.getState(getPositions=True).getPeriodicBoxVectors()
+
+    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]]),
+    )
+
+    assert numpy.isclose(context_box_vectors[0].x, 2.0)
+    assert numpy.isclose(context_box_vectors[1].y, 2.0)
+    assert numpy.isclose(context_box_vectors[2].z, 2.0)

openff/evaluator/utils/openmm.py

@@ -7,6 +7,7 @@
 
 import numpy
 import openmm
+from openmm import app
 from openmm import unit as _openmm_unit
 from pint import UndefinedUnitError
 
@@ -389,3 +390,87 @@ def system_subset(
     # Create the parameterized sub-system.
     system = force_field_subset.create_openmm_system(topology)
     return system, parameter_value
+
+
+def update_context_with_positions(
+    context: openmm.Context,
+    positions: _openmm_unit.Quantity,
+    box_vectors: Optional[_openmm_unit.Quantity],
+):
+    """Set a collection of positions and box vectors on an OpenMM context and compute
+    any extra positions such as v-site positions.
+
+    Parameters
+    ----------
+    context
+        The OpenMM context to set the positions on.
+    positions
+        A unit wrapped numpy array with shape=(n_atoms, 3) that contains the positions
+        to set.
+    box_vectors
+        An optional unit wrapped numpy array with shape=(3, 3) that contains the box
+        vectors to set.
+    """
+
+    system = context.getSystem()
+
+    n_vsites = sum(
+        1 for i in range(system.getNumParticles()) if system.isVirtualSite(i)
+    )
+
+    n_atoms = system.getNumParticles() - n_vsites
+
+    if len(positions) != n_atoms and len(positions) != (n_atoms + n_vsites):
+
+        raise ValueError(
+            "The length of the positions array does not match either the "
+            "the number of atoms or the number of atoms + v-sites."
+        )
+
+    if n_vsites > 0 and len(positions) != (n_atoms + n_vsites):
+
+        new_positions = numpy.zeros((system.getNumParticles(), 3))
+
+        i = 0
+
+        for j in range(system.getNumParticles()):
+
+            if not system.isVirtualSite(j):
+                # take an old position and update the index
+                new_positions[j] = positions[i].value_in_unit(_openmm_unit.nanometers)
+                i += 1
+
+        positions = new_positions * _openmm_unit.nanometers
+
+    if box_vectors is not None:
+        context.setPeriodicBoxVectors(*box_vectors)
+
+    context.setPositions(positions)
+
+    if n_vsites > 0:
+        context.computeVirtualSites()
+
+
+def update_context_with_pdb(
+    context: openmm.Context,
+    pdb_file: app.PDBFile,
+):
+    """Extracts the positions and box vectors from a PDB file object and set these
+    on an OpenMM context and compute any extra positions such as v-site positions.
+
+    Parameters
+    ----------
+    context
+        The OpenMM context to set the positions on.
+    pdb_file
+        The PDB file object to extract the positions and box vectors from.
+    """
+
+    positions = pdb_file.getPositions(asNumpy=True)
+
+    box_vectors = pdb_file.topology.getPeriodicBoxVectors()
+
+    if box_vectors is None:
+        box_vectors = context.getSystem().getDefaultPeriodicBoxVectors()
+
+    update_context_with_positions(context, positions, box_vectors)