diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b17936..392c880 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ Changelog * Fixed type check for ``water_template``. * Add support for simulations in the osmotic ensemble. * Fixed non-uniform bulk insertion positions caused by use of normal rather than uniform random numbers. +* Add methods to update the system with the current water state and return system without ghost waters. [2025.2.0](https://github.com/openbiosim/loch/compare/2025.1.0...2025.2.0) - Feb 2026 ------------------------------------------------------------------------------------- diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 0fec54b..0c17461 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -863,6 +863,136 @@ def system(self) -> _Any: """ return self._system.clone() + def _system_from_context(self, context: _Optional[_openmm.Context]) -> _Any: + """ + Clone the internal system and, if a context is provided, update + atomic coordinates and the periodic box from it. + """ + system = self._system.clone() + + if context is not None: + if not isinstance(context, _openmm.Context): + raise ValueError("'context' must be of type 'openmm.Context'") + + from sire.legacy.IO import setCoordinates + + state = context.getState(getPositions=True) + positions = ( + state.getPositions(asNumpy=True) / _openmm.unit.angstrom + ).tolist() + + try: + system._system = setCoordinates(self._system._system, positions) + except Exception as e: + raise ValueError( + f"Could not update the system with the current positions: {e}" + ) + + # Update the periodic box from the context. + box = state.getPeriodicBoxVectors() + v0 = [10 * box[0].x, 10 * box[0].y, 10 * box[0].z] + v1 = [10 * box[1].x, 10 * box[1].y, 10 * box[1].z] + v2 = [10 * box[2].x, 10 * box[2].y, 10 * box[2].z] + space = _sr.vol.TriclinicBox( + _sr.maths.Vector(*v0), + _sr.maths.Vector(*v1), + _sr.maths.Vector(*v2), + ) + system.set_property("space", space) + + return system + + def update_system(self, context: _Optional[_openmm.Context] = None) -> _Any: + """ + Update the system with the current water state and (optionally) positions. + + Parameters + ---------- + + context: openmm.Context + The OpenMM context containing the current positions. + + Returns + ------- + + system: sire.system.System + The updated GCMC system. + """ + + # Clone the system and update coordinates and box from the context. + system = self._system_from_context(context) + + # Flag the ghost waters in the system according to the current water state. + system = self._flag_ghost_waters(system) + + # Turn OFF interactions for ghost waters in the system. This is done + # by setting the charges and Lennard-Jones parameters of ghost waters + # to zero. + for mol in system["property is_ghost_water"].molecules(): + cursor = mol.cursor() + for atom in cursor.atoms(): + LJ = atom["LJ"] + atom["charge"] = 0.0 * _sr.units.mod_electron + atom["LJ"] = _sr.legacy.MM.LJParameter( + LJ.sigma(), 0.0 * _sr.units.kcal_per_mol + ) + mol = cursor.commit() + system.update(mol) + + # Turn ON any ghost buffer waters that were inserted during sampling. + # These are waters in the last _num_ghost_waters slots that now have state=1. + first_ghost_buffer = self._num_waters - self._num_ghost_waters + for i in range(first_ghost_buffer, self._num_waters): + if self._water_state[i] == 1: + mol = system[ + system.atoms()[int(self._water_indices_sire[i])].molecule() + ] + cursor = mol.cursor() + for j, atom in enumerate(cursor.atoms()): + atom["charge"] = self._water_charge[j] * _sr.units.mod_electron + atom["LJ"] = _sr.legacy.MM.LJParameter( + self._water_sigma[j] * _sr.units.angstrom, + self._water_epsilon[j] * _sr.units.kcal_per_mol, + ) + mol = cursor.commit() + system.update(mol) + + return system + + def finalise_system(self, context: _Optional[_openmm.Context] = None) -> _Any: + """ + Return the system with ghost waters removed and (optionally) positions + updated from an OpenMM context. + + Unlike :meth:`update_system`, which retains ghost waters with zeroed + parameters, this method deletes them entirely so the returned system is + suitable for use as input to non-GCMC simulations or analyses. + + Parameters + ---------- + + context: openmm.Context + The OpenMM context containing the current positions. + + Returns + ------- + + system: sire.system.System + The system with ghost waters removed. + """ + + # Clone the system and update coordinates and box from the context. + system = self._system_from_context(context) + + # Collect all ghost water molecules before removing any, since each + # removal shifts atom indices. + ghost_oxygens = self._water_indices_sire[self._get_ghost_waters()] + ghost_mols = [system[system.atoms()[int(i)].molecule()] for i in ghost_oxygens] + for mol in ghost_mols: + system.remove(mol) + + return system + def compiler_log(self) -> str: """ Return the GPU kernel compiler log. @@ -3091,6 +3221,9 @@ def _flag_ghost_waters(self, system): if not isinstance(system, _sr.system.System): raise ValueError("'system' must be a Sire system") + # Clone the system so that we don't modify the original. + system = system.clone() + # Use the Sire atom indices (no vsite offset) so that lookups into the # input topology are correct regardless of virtual sites in the context. ghost_oxygens = self._water_indices_sire[self._get_ghost_waters()] diff --git a/tests/test_energy.py b/tests/test_energy.py index 1dc878b..703a8a2 100644 --- a/tests/test_energy.py +++ b/tests/test_energy.py @@ -1,8 +1,10 @@ import math import openmm +import openmm.unit as omm_unit import os import pytest +import numpy as np import sire as sr from loch import GCMCSampler, SoftcoreForm @@ -458,3 +460,159 @@ def _create_and_run(seed): assert math.isclose(energy1_lj, energy2_lj, abs_tol=1e-4), ( f"LJ energy mismatch: {energy1_lj!r} vs {energy2_lj!r}" ) + + +@pytest.mark.skipif( + "CUDA_VISIBLE_DEVICES" not in os.environ, + reason="Requires CUDA enabled GPU.", +) +@pytest.mark.parametrize("platform", ["cuda", "opencl"]) +def test_update_system_insertion(platform, water_box): + """ + After an accepted insertion of a ghost buffer water, update_system() must: + - restore real LJ/charge parameters on the inserted water, and + - reflect the current OpenMM coordinates for that water. + """ + mols, reference = water_box + + sampler = GCMCSampler( + mols, + cutoff_type="rf", + cutoff="10 A", + reference=reference, + log_level="debug", + ghost_file=None, + log_file=None, + test=True, + platform=platform, + ) + + d = sampler.system().dynamics( + cutoff_type="rf", + cutoff="10 A", + temperature="298 K", + pressure=None, + constraint="h_bonds", + timestep="2 fs", + platform=platform, + ) + + ctx = d.context() + first_ghost_buffer = sampler._num_waters - sampler._num_ghost_waters + + # Loop until a ghost buffer water is inserted. + inserted_idx = None + while inserted_idx is None: + state_before = sampler.water_state() + moves = sampler.move(ctx) + if len(moves) == 0 or moves[0] != 0: + continue + state_after = sampler.water_state() + # Find a buffer-slot water that changed from 0 → 1. + for i in range(first_ghost_buffer, sampler._num_waters): + if state_before[i] == 0 and state_after[i] == 1: + inserted_idx = i + break + + assert inserted_idx is not None, "No buffer water was inserted" + + # Call update_system and get the returned system. + system = sampler.update_system(ctx) + + # Get the sire atom index (no vsite offset) of the inserted water's oxygen. + o_idx = int(sampler._water_indices_sire[inserted_idx]) + inserted_mol = system[system.atoms()[o_idx].molecule()] + + # Check parameters are real (non-zero). + for j, atom in enumerate(inserted_mol.atoms()): + charge = atom.property("charge").value() + lj = atom.property("LJ") + epsilon = lj.epsilon().value() + + assert charge == pytest.approx(sampler._water_charge[j], abs=1e-6), ( + f"Atom {j} charge not restored: got {charge}, expected {sampler._water_charge[j]}" + ) + assert epsilon == pytest.approx(sampler._water_epsilon[j], abs=1e-6), ( + f"Atom {j} epsilon not restored: got {epsilon}, expected {sampler._water_epsilon[j]}" + ) + + # Check coordinates match the OpenMM context. + omm_positions = ( + ctx.getState(getPositions=True, enforcePeriodicBox=False) + .getPositions(asNumpy=True) + .value_in_unit(omm_unit.angstrom) + ) + + # _water_indices uses the vsite-offset index into the OpenMM atom list. + omm_o_idx = int(sampler._water_indices[inserted_idx]) + + for j, atom in enumerate(inserted_mol.atoms()): + sire_coord = atom.property("coordinates") + sire_xyz = np.array( + [sire_coord.x().value(), sire_coord.y().value(), sire_coord.z().value()] + ) + omm_xyz = omm_positions[omm_o_idx + j] + assert np.allclose(sire_xyz, omm_xyz, atol=1e-3), ( + f"Atom {j} coordinates mismatch: sire={sire_xyz}, omm={omm_xyz}" + ) + + +@pytest.mark.skipif( + "CUDA_VISIBLE_DEVICES" not in os.environ, + reason="Requires CUDA enabled GPU.", +) +@pytest.mark.parametrize("platform", ["cuda", "opencl"]) +def test_finalise_system(platform, water_box): + """ + After an accepted insertion of a ghost buffer water, finalise_system() must + return a system whose water count equals the number of real (non-ghost) + waters in the sampler. + """ + mols, reference = water_box + + sampler = GCMCSampler( + mols, + cutoff_type="rf", + cutoff="10 A", + reference=reference, + log_level="debug", + ghost_file=None, + log_file=None, + test=True, + platform=platform, + ) + + d = sampler.system().dynamics( + cutoff_type="rf", + cutoff="10 A", + temperature="298 K", + pressure=None, + constraint="h_bonds", + timestep="2 fs", + platform=platform, + ) + + ctx = d.context() + first_ghost_buffer = sampler._num_waters - sampler._num_ghost_waters + + # Loop until a ghost buffer water is inserted. + inserted = False + while not inserted: + state_before = sampler.water_state() + moves = sampler.move(ctx) + if len(moves) == 0 or moves[0] != 0: + continue + state_after = sampler.water_state() + for i in range(first_ghost_buffer, sampler._num_waters): + if state_before[i] == 0 and state_after[i] == 1: + inserted = True + break + + system = sampler.finalise_system(ctx) + + expected_waters = int(sampler._get_non_ghost_waters().size) + actual_waters = system["water"].num_molecules() + + assert actual_waters == expected_waters, ( + f"Water count mismatch: got {actual_waters}, expected {expected_waters}" + )