Fix issue #25

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
 -------------------------------------------------------------------------------------

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()]

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}"
+    )