Fix excluding v-sites from OpenMM positions (#390)

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.8
+-----
+
+Bugfixes
+""""""""
+
+* PR `#390 <https://github.com/openforcefield/openff-evaluator/pull/390>`_: Fix excluding v-sites from OpenMM positions
+
 0.3.7
 -----
 

integration-tests/virtual-sites/run-tip4p.py

@@ -0,0 +1,131 @@
+import shutil
+
+from openff.toolkit.typing.engines.smirnoff import ForceField, ParameterList
+from openmm import unit as openmm_unit
+
+from openff.evaluator import unit
+from openff.evaluator.forcefield import ParameterGradientKey, SmirnoffForceFieldSource
+from openff.evaluator.protocols.coordinates import BuildCoordinatesPackmol
+from openff.evaluator.protocols.forcefield import BuildSmirnoffSystem
+from openff.evaluator.protocols.openmm import OpenMMEnergyMinimisation, OpenMMSimulation
+from openff.evaluator.substances import Substance
+from openff.evaluator.thermodynamics import Ensemble, ThermodynamicState
+
+
+def tip4p_force_field() -> ForceField:
+
+    force_field = ForceField()
+
+    constraint_handler = force_field.get_parameter_handler("Constraints")
+    constraint_handler.add_parameter(
+        {
+            "smirks": "[#1:1]-[#8X2H2+0:2]-[#1]",
+            "distance": 0.9572 * openmm_unit.angstrom,
+        }
+    )
+    constraint_handler.add_parameter(
+        {
+            "smirks": "[#1:1]-[#8X2H2+0]-[#1:2]",
+            "distance": 1.5139 * openmm_unit.angstrom,
+        }
+    )
+
+    vdw_handler = force_field.get_parameter_handler("vdW")
+    vdw_handler.add_parameter(
+        {
+            "smirks": "[#1:1]-[#8X2H2+0]-[#1]",
+            "epsilon": (
+                78.0
+                * openmm_unit.kelvin
+                * openmm_unit.BOLTZMANN_CONSTANT_kB
+                * openmm_unit.AVOGADRO_CONSTANT_NA
+            ),
+            "sigma": 3.154 * openmm_unit.angstrom,
+        }
+    )
+    vdw_handler.add_parameter(
+        {
+            "smirks": "[#1]-[#8X2H2+0:1]-[#1]",
+            "epsilon": 0.0 * openmm_unit.kilojoules_per_mole,
+            "sigma": 1.0 * openmm_unit.angstrom,
+        }
+    )
+
+    force_field.get_parameter_handler("Electrostatics")
+    force_field.get_parameter_handler(
+        "ChargeIncrementModel",
+        {"version": "0.3", "partial_charge_method": "formal_charge"},
+    )
+
+    virtual_site_handler = force_field.get_parameter_handler("VirtualSites")
+    virtual_site_handler.add_parameter(
+        {
+            "smirks": "[#1:1]-[#8X2H2+0:2]-[#1:3]",
+            "type": "DivalentLonePair",
+            "distance": -0.15 * openmm_unit.angstrom,
+            "outOfPlaneAngle": 0.0 * openmm_unit.degrees,
+            "match": "once",
+            "charge_increment1": 0.52 * openmm_unit.elementary_charge,
+            "charge_increment2": 0.0 * openmm_unit.elementary_charge,
+            "charge_increment3": 0.52 * openmm_unit.elementary_charge,
+        }
+    )
+    # Currently required due to OpenFF issue #884
+    virtual_site_handler._parameters = ParameterList(virtual_site_handler._parameters)
+
+    return force_field
+
+
+def main():
+
+    force_field = tip4p_force_field()
+    substance = Substance.from_components("O")
+
+    with open("force-field.json", "w") as file:
+        file.write(SmirnoffForceFieldSource.from_object(force_field).json())
+
+    build_coordinates = BuildCoordinatesPackmol("")
+    build_coordinates.substance = substance
+    build_coordinates.max_molecules = 216
+    build_coordinates.execute("build-coords")
+
+    apply_parameters = BuildSmirnoffSystem("")
+    apply_parameters.force_field_path = "force-field.json"
+    apply_parameters.coordinate_file_path = build_coordinates.coordinate_file_path
+    apply_parameters.substance = substance
+    apply_parameters.execute("apply-params")
+
+    minimize = OpenMMEnergyMinimisation("")
+    minimize.input_coordinate_file = build_coordinates.coordinate_file_path
+    minimize.parameterized_system = apply_parameters.parameterized_system
+    minimize.execute("minimize-coords")
+
+    npt = OpenMMSimulation("")
+    npt.input_coordinate_file = minimize.output_coordinate_file
+    npt.parameterized_system = apply_parameters.parameterized_system
+    npt.ensemble = Ensemble.NPT
+    npt.thermodynamic_state = ThermodynamicState(
+        temperature=298.15 * unit.kelvin, pressure=1.0 * unit.atmosphere
+    )
+    npt.steps_per_iteration = 500
+    npt.total_number_of_iterations = 2
+    npt.gradient_parameters = [
+        ParameterGradientKey(
+            tag="vdW", smirks="[#1:1]-[#8X2H2+0]-[#1]", attribute="epsilon"
+        )
+    ]
+    npt.output_frequency = 50
+    npt.execute("run-npt")
+
+    shutil.copytree("run-npt", "run-npt-1")
+
+    npt.total_number_of_iterations = 4
+    npt.execute("run-npt")
+
+    assert len(npt.observables) == 40
+    assert len(npt.observables["PotentialEnergy"].gradients) == 1
+    assert len(npt.observables["PotentialEnergy"].gradients[0]) == 40
+
+
+if __name__ == "__main__":
+    main()

openff/evaluator/protocols/openmm.py

@@ -30,6 +30,8 @@
 )
 from openff.evaluator.utils.openmm import (
     disable_pbc,
+    extract_atom_indices,
+    extract_positions,
     openmm_quantity_to_pint,
     pint_quantity_to_openmm,
     setup_platform_with_resources,
@@ -356,7 +358,11 @@ def _execute(self, directory, available_resources):
             pint_quantity_to_openmm(self.tolerance), self.max_iterations
         )
 
-        positions = simulation.context.getState(getPositions=True).getPositions()
+        positions = extract_positions(
+            simulation.context.getState(getPositions=True),
+            # Discard any v-sites.
+            extract_atom_indices(system),
+        )
 
         self.output_coordinate_file = os.path.join(directory, "minimised.pdb")
 
@@ -425,6 +431,45 @@ def __init__(self, integrator, topology, system, current_step):
             self.system = system
             self.currentStep = current_step
 
+    class _DCDReporter:
+        def __init__(self, file, append=False):
+
+            self._append = append
+
+            mode = "r+b" if append else "wb"
+
+            self._out = open(file, mode)
+
+            self._dcd = None
+            self._atom_indices = None
+
+        def report(self, simulation, state):
+
+            from openmm import app
+
+            if self._dcd is None:
+
+                self._dcd = app.DCDFile(
+                    self._out,
+                    simulation.topology,
+                    simulation.integrator.getStepSize(),
+                    simulation.currentStep,
+                    0,
+                    self._append,
+                )
+
+                system = simulation.system
+
+                self._atom_indices = extract_atom_indices(system)
+
+            self._dcd.writeModel(
+                extract_positions(state, self._atom_indices),
+                periodicBoxVectors=state.getPeriodicBoxVectors(),
+            )
+
+        def __del__(self):
+            self._out.close()
+
     def __init__(self, protocol_id):
 
         super().__init__(protocol_id)
@@ -924,9 +969,7 @@ def _simulate(self, context, integrator):
         # Build the reporters which we will use to report the state
         # of the simulation.
         append_trajectory = is_file_and_not_empty(self._local_trajectory_path)
-        dcd_reporter = app.DCDReporter(
-            self._local_trajectory_path, 0, append_trajectory
-        )
+        dcd_reporter = self._DCDReporter(self._local_trajectory_path, append_trajectory)
 
         statistics_file = open(self._local_statistics_path, "a+")
 
@@ -988,7 +1031,7 @@ def _simulate(self, context, integrator):
 
         final_state = context.getState(getPositions=True)
 
-        positions = final_state.getPositions()
+        positions = extract_positions(final_state, extract_atom_indices(system))
         topology.setPeriodicBoxVectors(final_state.getPeriodicBoxVectors())
 
         with open(self.output_coordinate_file, "w+") as configuration_file:

openff/evaluator/tests/test_utils/test_openmm.py

@@ -25,6 +25,8 @@
 from openff.evaluator.utils import get_data_filename
 from openff.evaluator.utils.observables import ObservableArray, ObservableFrame
 from openff.evaluator.utils.openmm import (
+    extract_atom_indices,
+    extract_positions,
     openmm_quantity_to_pint,
     pint_quantity_to_openmm,
     system_subset,
@@ -453,6 +455,24 @@ def test_update_context_with_pdb(tmpdir):
         numpy.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.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]]),
+    )
+
     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)
+
+
+def test_extract_atom_indices():
+
+    force_field = hydrogen_chloride_force_field(True, False, True)
+
+    topology: Topology = Molecule.from_smiles("Cl").to_topology()
+    system = force_field.create_openmm_system(topology)
+
+    assert system.getNumParticles() == 3
+    assert extract_atom_indices(system) == [0, 1]

openff/evaluator/utils/openmm.py

@@ -3,7 +3,7 @@
 """
 import copy
 import logging
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import numpy
 import openmm
@@ -474,3 +474,25 @@ def update_context_with_pdb(
         box_vectors = context.getSystem().getDefaultPeriodicBoxVectors()
 
     update_context_with_positions(context, positions, box_vectors)
+
+
+def extract_atom_indices(system: openmm.System) -> List[int]:
+    """Returns the indices of atoms in a system excluding any virtual sites."""
+
+    return [i for i in range(system.getNumParticles()) if not system.isVirtualSite(i)]
+
+
+def extract_positions(
+    state: openmm.State,
+    particle_indices: Optional[List[int]] = None,
+) -> _openmm_unit.Quantity:
+    """Extracts the positions from an OpenMM context, optionally excluding any v-site
+    positions which should be uniquely defined by the atomic positions.
+    """
+
+    positions = state.getPositions(asNumpy=True)
+
+    if particle_indices is not None:
+        positions = positions[particle_indices]
+
+    return positions