Add support for the osmotic ensemble

CHANGELOG.md

@@ -9,6 +9,7 @@ Changelog
 * Add support for long-range Lennard-Jones dispersion correction.
 * Add support for Beutler soft-core Lennard-Jones form.
 * Fixed type check for ``water_template``.
+* Add support for simulations in the osmotic ensemble.
 
 [2025.2.0](https://github.com/openbiosim/loch/compare/2025.1.0...2025.2.0) - Feb 2026
 -------------------------------------------------------------------------------------

src/loch/_sampler.py

@@ -39,6 +39,16 @@
 from ._platforms._rng import RNGManager as _RNGManager
 from ._softcore import SoftcoreForm as _SoftcoreForm
 
+import sys as _sys
+
+try:
+    "μ".encode(_sys.stdout.encoding or "utf-8")
+    _mu_sym = "μ"
+except (UnicodeEncodeError, LookupError):
+    _mu_sym = "mu"
+
+del _sys
+
 
 def _as_float32(arr: _np.ndarray) -> _np.ndarray:
     """Convert array to float32 only if not already float32."""
@@ -65,6 +75,7 @@ def __init__(
         excess_chemical_potential: str = "-6.09 kcal/mol",
         standard_volume: str = "30.543 A^3",
         temperature: str = "298 K",
+        pressure: _Optional[str] = None,
         adams_shift: _Union[int, float] = 0.0,
         num_ghost_waters: int = 20,
         batch_size: int = 1000,
@@ -126,6 +137,12 @@ def __init__(
         temperature: str
             The temperature of the system.
 
+        pressure: str
+            The pressure for NPT simulation, e.g. "1 atm". If None (default),
+            the simulation runs in the μVT ensemble with a fixed box. When set,
+            the box is updated from the context at each move and the Adams value
+            is recomputed accordingly, sampling the μpT (osmotic) ensemble.
+
         adams_shift: float
             The Adams shift.
 
@@ -330,6 +347,16 @@ def __init__(
         except Exception as e:
             raise ValueError(f"Could not validate the 'temperature': {e}")
 
+        if pressure is not None:
+            try:
+                self._pressure = self._validate_sire_unit(
+                    "pressure", pressure, _sr.u("atm")
+                )
+            except Exception as e:
+                raise ValueError(f"Could not validate the 'pressure': {e}")
+        else:
+            self._pressure = None
+
         if not isinstance(num_ghost_waters, int):
             raise ValueError("'num_ghost_waters' must be of type 'int'")
         self._num_ghost_waters = num_ghost_waters
@@ -685,27 +712,8 @@ def __init__(
             * _openmm.unit.kelvin
         )
 
-        # Work out the volume of the system and GCMC sphere.
-        volume = self._space.volume().value()
-        gcmc_volume = (4.0 * _np.pi * self._radius.value() ** 3) / 3.0
-
-        # Work out the Adams value.
-        B = (
-            self._beta * self._excess_chemical_potential.value()
-            + _np.log(gcmc_volume / self._standard_volume.value())
-        ) + self._adams_shift
-
-        # Work out the bulk Adams value.
-        B_bulk = (
-            self._beta * self._excess_chemical_potential.value()
-            + _np.log(volume / self._standard_volume.value())
-        ) + self._adams_shift
-
-        # Store the exponentials for the Adams values.
-        self._exp_B = _np.exp(B)
-        self._exp_minus_B = _np.exp(-B)
-        self._exp_B_bulk = _np.exp(B_bulk)
-        self._exp_minus_B_bulk = _np.exp(-B_bulk)
+        # Compute the Adams values from the current box and sphere volume.
+        self._compute_adams_values(self._system)
 
         # Coulomb energy prefactor.
         self._prefactor = 1.0 / (4.0 * _np.pi * _sr.units.epsilon0.value())
@@ -740,9 +748,6 @@ def __init__(
                     self._log_file, level=self._log_level.upper(), filter="loch"
                 )
 
-        # Log the Adams value.
-        _logger.debug(f"Adams value: {B:.6f}")
-
         import atexit
 
         # Register the cleanup function.
@@ -772,6 +777,7 @@ def __str__(self) -> str:
             f"excess_chemical_potential={self._excess_chemical_potential}, "
             f"standard_volume={self._standard_volume}, "
             f"temperature={self._temperature}, "
+            f"pressure={self._pressure}, "
             f"num_ghost_waters={self._num_ghost_waters}, "
             f"adams_shift={self._adams_shift}, "
             f"batch_size={self._batch_size}, "
@@ -872,15 +878,64 @@ def compiler_log(self) -> str:
         """
         return self._backend.compiler_log
 
+    def _compute_adams_values(self, system: _Any) -> None:
+        """
+        Compute the Adams values from the current box and sphere volume.
+
+        Updates self._exp_B, self._exp_minus_B, self._exp_B_bulk, and
+        self._exp_minus_B_bulk.
+
+        Parameters
+        ----------
+
+        system: sire.system.System, openmm.Context, openmm.State
+            The molecular system, OpenMM context, or OpenMM state.
+        """
+        if isinstance(system, _sr.system.System):
+            volume = system.property("space").volume().value()
+        elif isinstance(system, _openmm.Context):
+            box = system.getState().getPeriodicBoxVectors(asNumpy=True)
+            volume = _np.linalg.det(box / _openmm.unit.angstrom)
+        elif isinstance(system, _openmm.State):
+            box = system.getPeriodicBoxVectors(asNumpy=True)
+            volume = _np.linalg.det(box / _openmm.unit.angstrom)
+        else:
+            raise ValueError(
+                "'system' must be of type 'sire.system.System', 'openmm.Context', "
+                "or 'openmm.State'"
+            )
+
+        gcmc_volume = (4.0 * _np.pi * self._radius.value() ** 3) / 3.0
+
+        B = (
+            self._beta * self._excess_chemical_potential.value()
+            + _np.log(gcmc_volume / self._standard_volume.value())
+        ) + self._adams_shift
+
+        B_bulk = (
+            self._beta * self._excess_chemical_potential.value()
+            + _np.log(volume / self._standard_volume.value())
+        ) + self._adams_shift
+
+        self._exp_B = _np.exp(B)
+        self._exp_minus_B = _np.exp(-B)
+        self._exp_B_bulk = _np.exp(B_bulk)
+        self._exp_minus_B_bulk = _np.exp(-B_bulk)
+
+        _logger.debug(f"Adams value: {B:.6f}")
+
+        # Store the box volume in nm^3 for reuse in LRC corrections.
+        self._v_nm3 = volume / 1000.0
+
     def set_box(self, system: _Any) -> None:
         """
         Set the box information.
 
         Parameters
         ----------
 
-        system: sire.system.System, openmm.Context
-            The molecular system, or OpenMM context.
+        system: sire.system.System, openmm.Context, openmm.State
+            The molecular system, OpenMM context, or OpenMM state.
         """
 
         # Get the space property from the system.
@@ -890,8 +945,12 @@ def set_box(self, system: _Any) -> None:
             except Exception:
                 raise ValueError("'system' must contain a 'space' property")
         # Create a Sire TriclinicBox from the OpenMM box vectors.
-        elif isinstance(system, _openmm.Context):
-            box = system.getState().getPeriodicBoxVectors()
+        elif isinstance(system, (_openmm.Context, _openmm.State)):
+            box = (
+                system.getState().getPeriodicBoxVectors()
+                if isinstance(system, _openmm.Context)
+                else system.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]
@@ -900,7 +959,8 @@ def set_box(self, system: _Any) -> None:
             )
         else:
             raise ValueError(
-                "'system' must be of type 'sire.system.System' or 'openmm.Context'"
+                "'system' must be of type 'sire.system.System', 'openmm.Context', "
+                "or 'openmm.State'"
             )
 
         # Get the box information.
@@ -1304,10 +1364,12 @@ def move(self, context: _openmm.Context) -> list[int]:
                     else:
                         initial_energy = None
 
-                    # Cache the box volume (NVT, so constant throughout).
-                    if self._has_gcmc_lrc:
-                        box = state.getPeriodicBoxVectors(asNumpy=True)
-                        v_nm3 = _np.linalg.det(box / _openmm.unit.nanometer)
+                    # For NPT (μpT), update the box and recompute Adams values
+                    # from the current context state so that volume fluctuations
+                    # driven by the barostat are reflected in the acceptance criterion.
+                    if self._pressure is not None:
+                        self.set_box(state)
+                        self._compute_adams_values(state)
 
                     # Sample within the GCMC sphere.
                     if self._reference is not None and not self._is_bulk:
@@ -1583,7 +1645,7 @@ def move(self, context: _openmm.Context) -> list[int]:
                                     self._lrc_w_solute
                                     + 2.0 * n_w_before_insert * self._lrc_ww_half
                                 )
-                                / v_nm3
+                                / self._v_nm3
                                 * _openmm.unit.kilojoules_per_mole
                             )
                             dE_RF += dLRC
@@ -1686,7 +1748,7 @@ def move(self, context: _openmm.Context) -> list[int]:
                                     * (n_w_before_delete - 1.0)
                                     * self._lrc_ww_half
                                 )
-                                / v_nm3
+                                / self._v_nm3
                                 * _openmm.unit.kilojoules_per_mole
                             )
                             dE_RF += dLRC
@@ -2819,10 +2881,11 @@ def _set_nonbonded_forces(self, context):
                     self._nonbonded_force = force
                 elif self._is_fep and force.getName() == "GhostNonGhostNonbondedForce":
                     self._custom_nonbonded_force = force
-                elif "Barostat" in force.getName():
+                elif self._pressure is None and "Barostat" in force.getName():
                     msg = (
-                        f"GCMC must be used at constant volume: "
-                        f"'{force.getName()}' is not supported."
+                        f"A barostat was detected but no pressure was set: "
+                        f"'{force.getName()}' is not supported in the {_mu_sym}VT ensemble. "
+                        f"Pass pressure= to enable {_mu_sym}pT (osmotic) sampling."
                     )
                     _logger.error(msg)
                     raise TypeError(msg)

tests/test_osmotic.py

@@ -0,0 +1,70 @@
+import os
+
+import openmm
+import pytest
+
+from loch import GCMCSampler
+
+
+@pytest.mark.skipif(
+    "CUDA_VISIBLE_DEVICES" not in os.environ,
+    reason="Requires CUDA enabled GPU.",
+)
+@pytest.mark.parametrize("platform", ["cuda", "opencl"])
+def test_osmotic_ensemble(water_box, platform):
+    """
+    When pressure is set, move() updates box parameters and Adams values
+    from the current OpenMM context state after each call.
+
+    Two moves are performed with the box scaled between them. After the
+    second move, _v_nm3, _exp_B_bulk, and _exp_minus_B_bulk must all
+    reflect the new volume.
+    """
+    mols, _ = water_box
+
+    sampler = GCMCSampler(
+        mols,
+        cutoff_type="rf",
+        cutoff="10 A",
+        pressure="1 atm",
+        ghost_file=None,
+        log_file=None,
+        test=True,
+        platform=platform,
+        seed=42,
+    )
+
+    # NPT dynamics so the context carries a barostat (bypasses the muVT guard).
+    d = sampler.system().dynamics(
+        cutoff_type="rf",
+        cutoff="10 A",
+        temperature="298 K",
+        pressure="1 atm",
+        constraint="h_bonds",
+        timestep="2 fs",
+        platform=platform,
+    )
+    context = d.context()
+
+    # First move: record initial box-dependent values.
+    sampler.move(context)
+    v_nm3_before = sampler._v_nm3
+    exp_B_bulk_before = sampler._exp_B_bulk
+    exp_minus_B_bulk_before = sampler._exp_minus_B_bulk
+
+    # Scale the box uniformly to simulate a barostat volume move.
+    scale = 1.1
+    box = context.getState().getPeriodicBoxVectors(asNumpy=True)
+    box_nm = box.value_in_unit(openmm.unit.nanometer)
+    context.setPeriodicBoxVectors(
+        openmm.Vec3(*box_nm[0] * scale) * openmm.unit.nanometer,
+        openmm.Vec3(*box_nm[1] * scale) * openmm.unit.nanometer,
+        openmm.Vec3(*box_nm[2] * scale) * openmm.unit.nanometer,
+    )
+
+    # Second move: box parameters must now reflect the scaled volume.
+    sampler.move(context)
+
+    assert sampler._v_nm3 == pytest.approx(v_nm3_before * scale**3, rel=1e-5)
+    assert sampler._exp_B_bulk != exp_B_bulk_before
+    assert sampler._exp_minus_B_bulk != exp_minus_B_bulk_before