diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c77fb9..b3e4c90 100644 --- a/CHANGELOG.md +++ b/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 ------------------------------------------------------------------------------------- diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 2724506..79166f0 100644 --- a/src/loch/_sampler.py +++ b/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,6 +878,55 @@ 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. @@ -879,8 +934,8 @@ def set_box(self, system: _Any) -> None: 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) diff --git a/tests/test_osmotic.py b/tests/test_osmotic.py new file mode 100644 index 0000000..840801a --- /dev/null +++ b/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