diff --git a/docs/releasehistory.rst b/docs/releasehistory.rst index 0de31538..c7d849fd 100644 --- a/docs/releasehistory.rst +++ b/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 `_: Fix v-site positions not set by OpenMM + 0.3.6 ----- diff --git a/openff/evaluator/protocols/openmm.py b/openff/evaluator/protocols/openmm.py index b3258e4e..4fde8ec0 100644 --- a/openff/evaluator/protocols/openmm.py +++ b/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 diff --git a/openff/evaluator/tests/test_utils/test_openmm.py b/openff/evaluator/tests/test_utils/test_openmm.py index b0cd7e09..81873d9d 100644 --- a/openff/evaluator/tests/test_utils/test_openmm.py +++ b/openff/evaluator/tests/test_utils/test_openmm.py @@ -1,8 +1,10 @@ +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 @@ -10,8 +12,10 @@ 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) diff --git a/openff/evaluator/utils/openmm.py b/openff/evaluator/utils/openmm.py index 18208c4f..9016b01c 100644 --- a/openff/evaluator/utils/openmm.py +++ b/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)