diff --git a/CHANGELOG.md b/CHANGELOG.md index ca9dbb1..87f18d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ Changelog -------------------------------------------------------------------------------------- * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR. +* Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29). [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026 ------------------------------------------------------------------------------------- diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 0c17461..c1ffc54 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -731,6 +731,8 @@ def __init__( # Null the nonbonded forces. self._nonbonded_force = None self._custom_nonbonded_force = None + self._integration_groups = None + self._pme_groups = None # Flag for whether the last move was a bulk sampling move. self._is_bulk = False @@ -1353,6 +1355,8 @@ def reset(self) -> None: # Clear the forces. self._nonbonded_force = None self._custom_nonbonded_force = None + self._integration_groups = None + self._pme_groups = None # Clear the OpenMM context. self._openmm_context = None @@ -1482,7 +1486,11 @@ def move(self, context: _openmm.Context) -> list[int]: self._init_gcmc_lrc(context) # Get the OpenMM state. - state = context.getState(getPositions=True, getEnergy=self._is_pme) + state = context.getState( + getPositions=True, + getEnergy=self._is_pme, + groups=self._pme_groups, + ) # Get the current positions in OpenMM format and in Angstrom. positions_openmm = state.getPositions(asNumpy=True) @@ -1768,7 +1776,7 @@ def move(self, context: _openmm.Context) -> list[int]: # Get the new energy. final_energy = context.getState( - getEnergy=True + getEnergy=True, groups=self._pme_groups ).getPotentialEnergy() # Add the analytic LRC delta so the PME correction sees only @@ -1870,7 +1878,7 @@ def move(self, context: _openmm.Context) -> list[int]: # Get the new energy. final_energy = context.getState( - getEnergy=True + getEnergy=True, groups=self._pme_groups ).getPotentialEnergy() # Add the analytic LRC delta. @@ -3007,12 +3015,19 @@ def _set_nonbonded_forces(self, context): context: openmm.Context The OpenMM context to use. """ + if self._integration_groups is None: + self._integration_groups = ( + context.getIntegrator().getIntegrationForceGroups() + ) + if self._nonbonded_force is None or ( self._is_fep and self._custom_nonbonded_force is None ): for force in context.getSystem().getForces(): if isinstance(force, _openmm.NonbondedForce): - self._nonbonded_force = force + # Only accept a force actually used for integration. + if self._integration_groups & (1 << force.getForceGroup()): + self._nonbonded_force = force elif self._is_fep and force.getName() == "GhostNonGhostNonbondedForce": self._custom_nonbonded_force = force elif self._pressure is None and "Barostat" in force.getName(): @@ -3034,6 +3049,12 @@ def _set_nonbonded_forces(self, context): _logger.error(msg) raise ValueError(msg) + if self._pme_groups is None: + groups = 1 << self._nonbonded_force.getForceGroup() + if self._is_fep: + groups |= 1 << self._custom_nonbonded_force.getForceGroup() + self._pme_groups = groups + def _get_target_position(self, positions): """ Get the current centre of the GCMC sphere.