From 40f79b3becd613f21be4205424bd536935e17507 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 29 Apr 2026 14:48:43 +0100 Subject: [PATCH 1/7] Add support for long-range LJ dispersion correction. --- src/loch/_sampler.py | 183 +++++++++++++++++++++++++++++++++++-------- 1 file changed, 150 insertions(+), 33 deletions(-) diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 0b9e04e..992b3f9 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -287,6 +287,11 @@ def __init__( else: self._is_pme = False + # LRC state: initialised lazily on first move() call. + self._has_gcmc_lrc = False + self._lrc_w_solute = 0.0 + self._lrc_ww_half = 0.0 + try: self._radius = self._validate_sire_unit("radius", radius, _sr.u("A")) except Exception as e: @@ -707,7 +712,8 @@ def __init__( # Null the nonbonded forces. self._nonbonded_force = None - self._custom_nonbonded_force = None + self._custom_coulomb_force = None + self._custom_lj_force = None # Flag for whether the last move was a bulk sampling move. self._is_bulk = False @@ -1148,7 +1154,8 @@ def reset(self) -> None: # Clear the forces. self._nonbonded_force = None - self._custom_nonbonded_force = None + self._custom_coulomb_force = None + self._custom_lj_force = None # Clear the OpenMM context. self._openmm_context = None @@ -1273,6 +1280,10 @@ def move(self, context: _openmm.Context) -> list[int]: # We only need to get the positions and initial energy for the first # batch. These will be updated dynamically as moves are accepted. if num_batches == 1: + # Detect GCMC LRC parameters from context on first call. + if not self._has_gcmc_lrc: + self._init_gcmc_lrc(context) + # Get the OpenMM state. state = context.getState(getPositions=True, getEnergy=self._is_pme) @@ -1286,6 +1297,11 @@ 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) + # Sample within the GCMC sphere. if self._reference is not None and not self._is_bulk: target = self._get_target_position(positions_angstrom).astype( @@ -1521,6 +1537,10 @@ def move(self, context: _openmm.Context) -> list[int]: # Insertion move. if is_deletion[idx] == 0: + # Capture n_w before the insertion for LRC delta calculation. + if self._has_gcmc_lrc: + n_w_before_insert = context.getParameter("n_w") + # Accept the move. self._accept_insertion( idx, idx_water, positions_openmm, positions_angstrom, context @@ -1548,6 +1568,19 @@ def move(self, context: _openmm.Context) -> list[int]: getEnergy=True ).getPotentialEnergy() + # Add the analytic LRC delta so the PME correction sees only + # the RF→PME electrostatic difference. + if self._has_gcmc_lrc: + dLRC = ( + ( + self._lrc_w_solute + + 2.0 * n_w_before_insert * self._lrc_ww_half + ) + / v_nm3 + * _openmm.unit.kilojoules_per_mole + ) + dE_RF += dLRC + # Compute the PME acceptance correction. acc_prob = _np.exp( min( @@ -1608,6 +1641,10 @@ def move(self, context: _openmm.Context) -> list[int]: # Deletion move. else: + # Capture n_w before the deletion for LRC delta calculation. + if self._has_gcmc_lrc: + n_w_before_delete = context.getParameter("n_w") + # Accept the move. self._accept_deletion(candidates[idx], context) @@ -1633,6 +1670,20 @@ def move(self, context: _openmm.Context) -> list[int]: getEnergy=True ).getPotentialEnergy() + # Add the analytic LRC delta. + if self._has_gcmc_lrc: + dLRC = ( + -( + self._lrc_w_solute + + 2.0 + * (n_w_before_delete - 1.0) + * self._lrc_ww_half + ) + / v_nm3 + * _openmm.unit.kilojoules_per_mole + ) + dE_RF += dLRC + # Compute the PME acceptance correction. acc_prob = _np.exp( min( @@ -2125,31 +2176,37 @@ def _initialise_gpu_memory(self): ) # Flags for the required force. - has_gng = False + has_gng_coul = False + has_gng_lj = False # Find the required forces. for force in d.context().getSystem().getForces(): - if force.getName() == "GhostNonGhostNonbondedForce": - gng_force = force - has_gng = True - break + if force.getName() == "GhostNonGhostCoulombForce": + gng_coul_force = force + has_gng_coul = True + elif force.getName() == "GhostNonGhostLJForce": + gng_lj_force = force + has_gng_lj = True # Make sure the force was found. - if not has_gng: + if not has_gng_coul: raise ValueError( - "Could not find the GhostNonGhostNonbondedForce in the system" + "Could not find the GhostNonGhostCoulombForce in the system" + ) + if not has_gng_lj: + raise ValueError( + "Could not find the GhostNonGhostLJForce in the system" ) - # Get the parameters for the GhostNonGhostNonbondedForce. + # Get the parameters for the GhostNonGhost nonbonded forces. charges = _np.zeros(self._num_atoms, dtype=_np.float32) sigmas = _np.zeros(self._num_atoms, dtype=_np.float32) epsilons = _np.zeros(self._num_atoms, dtype=_np.float32) alphas = _np.zeros(self._num_atoms, dtype=_np.float32) - for i in range(gng_force.getNumParticles()): + for i in range(gng_coul_force.getNumParticles()): # Custom force parameters are returned as floats. - q, half_sigma, two_sqrt_epsilon, alpha, _ = ( - gng_force.getParticleParameters(i) - ) + q, alpha, _ = gng_coul_force.getParticleParameters(i) + half_sigma, two_sqrt_epsilon, _ = gng_lj_force.getParticleParameters(i) # Charge in |e|, sigma in nm, epsilon in kJ/mol. charges[i] = q # Rescale and convert units. @@ -2321,6 +2378,15 @@ def _initialise_gpu_memory(self): (1, self._num_waters), _np.int32 ) + def _init_gcmc_lrc(self, context): + """Detect and cache GCMC LRC parameters from the OpenMM context.""" + try: + self._lrc_w_solute = context.getParameter("lrc_w_solute") + self._lrc_ww_half = context.getParameter("lrc_ww_half") + self._has_gcmc_lrc = True + except Exception: + self._has_gcmc_lrc = False + def _accept_insertion( self, idx, idx_water, positions_openmm, positions_angstrom, context ): @@ -2374,14 +2440,20 @@ def _accept_insertion( ) # Update the custom NonBondedForce parameters. if self._is_fep: - self._custom_nonbonded_force.setParticleParameters( + self._custom_coulomb_force.setParticleParameters( start_idx + i, ( self._water_charge[i], + 0.0, + 0.0, + ), + ) + self._custom_lj_force.setParticleParameters( + start_idx + i, + ( self._water_sigma_custom[i], self._water_epsilon_custom[i], 0.0, - 0.0, ), ) @@ -2393,7 +2465,8 @@ def _accept_insertion( # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_nonbonded_force.updateParametersInContext(context) + self._custom_coulomb_force.updateParametersInContext(context) + self._custom_lj_force.updateParametersInContext(context) # Update the state of the water on the GPU. self._kernels["update_water"]( @@ -2417,6 +2490,10 @@ def _accept_insertion( # Update the number of waters in the sampling volume. self._N += 1 + # Update the GCMC LRC water count in the context. + if self._has_gcmc_lrc: + context.setParameter("n_w", context.getParameter("n_w") + 1.0) + def _accept_deletion(self, idx, context): """ Accept a deletion move. @@ -2445,13 +2522,19 @@ def _accept_deletion(self, idx, context): ) # Update the CustomNonBondedForce parameters. if self._is_fep: - self._custom_nonbonded_force.setParticleParameters( + self._custom_coulomb_force.setParticleParameters( start_idx + i, ( 0.0, - self._water_sigma_custom[i], 0.0, 0.0, + ), + ) + self._custom_lj_force.setParticleParameters( + start_idx + i, + ( + self._water_sigma_custom[i], + 0.0, 0.0, ), ) @@ -2461,7 +2544,8 @@ def _accept_deletion(self, idx, context): # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_nonbonded_force.updateParametersInContext(context) + self._custom_coulomb_force.updateParametersInContext(context) + self._custom_lj_force.updateParametersInContext(context) # Update the state of the water on the GPU. self._kernels["update_water"]( @@ -2487,6 +2571,10 @@ def _accept_deletion(self, idx, context): # Update the number of waters in the sampling volume. self._N -= 1 + # Update the GCMC LRC water count in the context. + if self._has_gcmc_lrc: + context.setParameter("n_w", context.getParameter("n_w") - 1.0) + def _reject_deletion(self, idx, context): """ Reject a deletion move. @@ -2518,14 +2606,20 @@ def _reject_deletion(self, idx, context): ) # Update the CustomNonBondedForce parameters. if self._is_fep: - self._custom_nonbonded_force.setParticleParameters( + self._custom_coulomb_force.setParticleParameters( start_idx + i, ( self._water_charge[i], + 0.0, + 0.0, + ), + ) + self._custom_lj_force.setParticleParameters( + start_idx + i, + ( self._water_sigma_custom[i], self._water_epsilon_custom[i], 0.0, - 0.0, ), ) @@ -2534,7 +2628,8 @@ def _reject_deletion(self, idx, context): # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_nonbonded_force.updateParametersInContext(context) + self._custom_coulomb_force.updateParametersInContext(context) + self._custom_lj_force.updateParametersInContext(context) # Update the state of the water on the GPU. self._kernels["update_water"]( @@ -2560,6 +2655,10 @@ def _reject_deletion(self, idx, context): # Update the number of waters in the sampling volume. self._N += 1 + # Update the GCMC LRC water count in the context. + if self._has_gcmc_lrc: + context.setParameter("n_w", context.getParameter("n_w") + 1.0) + def _set_water_state(self, context, indices=None, states=None, force=False): """ Update the state for a list of waters. This can be used by external @@ -2597,7 +2696,8 @@ def _set_water_state(self, context, indices=None, states=None, force=False): # Assume the context has been recreated, so we need to get the # new forces. self._nonbonded_force = None - self._custom_nonbonded_force = None + self._custom_coulomb_force = None + self._custom_lj_force = None # Update even if the state is unchanged. force = True @@ -2627,13 +2727,19 @@ def _set_water_state(self, context, indices=None, states=None, force=False): ) # Update the CustomNonbondedForce parameters. if self._is_fep: - self._custom_nonbonded_force.setParticleParameters( + self._custom_coulomb_force.setParticleParameters( start_idx + i, ( 0.0, - self._water_sigma_custom[i], 0.0, 0.0, + ), + ) + self._custom_lj_force.setParticleParameters( + start_idx + i, + ( + self._water_sigma_custom[i], + 0.0, 0.0, ), ) @@ -2674,14 +2780,20 @@ def _set_water_state(self, context, indices=None, states=None, force=False): ) # Update the CustomNonBondedForce parameters. if self._is_fep: - self._custom_nonbonded_force.setParticleParameters( + self._custom_coulomb_force.setParticleParameters( start_idx + i, ( self._water_charge[i], + 0.0, + 0.0, + ), + ) + self._custom_lj_force.setParticleParameters( + start_idx + i, + ( self._water_sigma_custom[i], self._water_epsilon_custom[i], 0.0, - 0.0, ), ) @@ -2717,7 +2829,8 @@ def _set_water_state(self, context, indices=None, states=None, force=False): # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_nonbonded_force.updateParametersInContext(context) + self._custom_coulomb_force.updateParametersInContext(context) + self._custom_lj_force.updateParametersInContext(context) def _set_nonbonded_forces(self, context): """ @@ -2730,13 +2843,15 @@ def _set_nonbonded_forces(self, context): The OpenMM context to use. """ if self._nonbonded_force is None or ( - self._is_fep and self._custom_nonbonded_force is None + self._is_fep and self._custom_coulomb_force is None ): for force in context.getSystem().getForces(): if isinstance(force, _openmm.NonbondedForce): self._nonbonded_force = force - elif self._is_fep and force.getName() == "GhostNonGhostNonbondedForce": - self._custom_nonbonded_force = force + elif self._is_fep and force.getName() == "GhostNonGhostCoulombForce": + self._custom_coulomb_force = force + elif self._is_fep and force.getName() == "GhostNonGhostLJForce": + self._custom_lj_force = force elif "Barostat" in force.getName(): msg = ( f"GCMC must be used at constant volume: " @@ -2750,7 +2865,9 @@ def _set_nonbonded_forces(self, context): _logger.error(msg) raise ValueError(msg) - if self._is_fep and self._custom_nonbonded_force is None: + if self._is_fep and ( + self._custom_coulomb_force is None or self._custom_lj_force is None + ): msg = "Could not find a CustomNonbondedForce in the system" _logger.error(msg) raise ValueError(msg) From 23bb1cd70dcd68402a17a2d563729cc5f8bfe532 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 30 Apr 2026 10:15:51 +0100 Subject: [PATCH 2/7] Add Beutler soft-core form for ABFE. --- src/loch/_kernels.py | 30 +++++++++++++++--------------- src/loch/_sampler.py | 33 +++++++++++++++++++++------------ src/loch/_softcore.py | 1 + tests/test_energy.py | 9 ++++----- 4 files changed, 41 insertions(+), 32 deletions(-) diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py index 8165486..d402b89 100644 --- a/src/loch/_kernels.py +++ b/src/loch/_kernels.py @@ -395,10 +395,10 @@ float rf_kappa, float rf_correction, int softcore_form, - float sc_coulomb_power, float sc_shift_coulomb, float sc_shift_delta, - int sc_taylor_power) + int sc_taylor_power, + float sc_beutler_alpha) { // Work out the atom index. const int idx_atom = GET_GLOBAL_ID(0); @@ -562,6 +562,7 @@ // Compute the LJ interaction using the chosen soft-core form. float sig6; + float lj_prefactor = 1.0f; if (softcore_form == 1) { // Taylor soft-core LJ: @@ -573,6 +574,16 @@ : powf(a, (float)sc_taylor_power); sig6 = s6_val / (alpha_m * s6_val + r6); } + else if (softcore_form == 2) + { + // Beutler soft-core LJ: + // sig6 = sigma^6 / (sc_beutler_alpha * sigma^6 * alpha + r^6) + // V_LJ = (1 - alpha) * 4 * epsilon * sig6 * (sig6 - 1) + const float s6_val = powf(s, 6.0f); + const float r6 = r * r * r * r * r * r; + sig6 = s6_val / (sc_beutler_alpha * s6_val * a + r6); + lj_prefactor = 1.0f - a; + } else { // Zacharias soft-core LJ: @@ -581,22 +592,11 @@ const float delta_lj = sc_shift_delta * a; sig6 = powf(s, 6.0f) / powf((s * delta_lj) + (r * r), 3.0f); } - energy_lj[idx] += 4.0f * e * sig6 * (sig6 - 1.0f); - - // Compute the Coulomb power expression. - float cpe; - if (sc_coulomb_power == 0.0f) - { - cpe = 1.0f; - } - else - { - cpe = powf((1.0f - a), sc_coulomb_power); - } + energy_lj[idx] += lj_prefactor * 4.0f * e * sig6 * (sig6 - 1.0f); // Compute the Coulomb interaction. energy_coul[idx] += (q0 * q1) * - ((cpe / sqrtf((sc_shift_coulomb * sc_shift_coulomb * a) + ((1.0f / sqrtf((sc_shift_coulomb * sc_shift_coulomb * a) + (r * r))) + (rf_kappa * r2) - rf_correction); } diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 992b3f9..d2e7ac0 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -79,11 +79,11 @@ def __init__( lambda_value: float = 0.0, rest2_scale: float = 1.0, rest2_selection: _Optional[str] = None, - coulomb_power: float = 0.0, shift_coulomb: str = "1 A", shift_delta: str = "1.5 A", softcore_form: str = "zacharias", taylor_power: int = 1, + beutler_alpha: float = 0.5, swap_end_states: bool = False, restart: bool = False, overwrite: bool = False, @@ -212,8 +212,8 @@ def __init__( softcore_form : str The soft-core potential form to use for alchemical interactions. - This can be either 'zacharias' or 'taylor'. The default is - 'zacharias'. + Valid options are 'zacharias' (default), 'taylor', and 'beutler'. + The Beutler form is recommended for ABFE calculations. taylor_power : int The power to use for the alpha term in the Taylor soft-core LJ @@ -221,6 +221,11 @@ def __init__( Must be between 0 and 4. The default is 1. Only used when softcore_form is 'taylor'. + beutler_alpha : float + The dimensionless scale factor for the r^6 shift in the Beutler + soft-core form. Must be >= 0. The default is 0.5. Only used when + softcore_form is 'beutler'. + swap_end_states: bool Whether to swap the end states of the alchemical systems. @@ -494,12 +499,6 @@ def __init__( ) self._rest2_selection = rest2_selection - try: - coulomb_power = float(coulomb_power) - except Exception: - raise ValueError("'coulomb_power' must be of type 'float'") - self._coulomb_power = float(coulomb_power) - try: self._shift_coulomb = self._validate_sire_unit( "shift_coulomb", shift_coulomb, _sr.u("A") @@ -534,6 +533,14 @@ def __init__( raise ValueError("'taylor_power' must be between 0 and 4") self._taylor_power = taylor_power + try: + beutler_alpha = float(beutler_alpha) + except Exception: + raise ValueError("'beutler_alpha' must be of type 'float'") + if beutler_alpha < 0.0: + raise ValueError("'beutler_alpha' must be >= 0") + self._beutler_alpha = beutler_alpha + if not isinstance(swap_end_states, bool): raise ValueError("'swap_end_states' must be of type 'bool'") self._swap_end_states = swap_end_states @@ -781,7 +788,6 @@ def __str__(self) -> str: f"restart={self._restart}, " f"rest2_scale={self._rest2_scale}, " f"rest2_selection={self._rest2_selection}, " - f"coulomb_power={self._coulomb_power}, " f"shift_coulomb={self._shift_coulomb}, " f"shift_delta={self._shift_delta}, " f"overwrite={self._overwrite}, " @@ -1458,10 +1464,10 @@ def move(self, context: _openmm.Context) -> list[int]: self._rf_kappa, self._rf_correction, self._sc_softcore_form, - self._sc_coulomb_power, self._sc_shift_coulomb, self._sc_shift_delta, self._sc_taylor_power, + self._sc_beutler_alpha, block=(self._num_threads, 1, 1), grid=(self._atom_blocks, self._batch_size, 1), ) @@ -2157,6 +2163,9 @@ def _initialise_gpu_memory(self): if self._softcore_form == _SoftcoreForm.TAYLOR: _map["use_taylor_softening"] = True _map["taylor_power"] = self._taylor_power + elif self._softcore_form == _SoftcoreForm.BEUTLER: + _map["use_beutler_softening"] = True + _map["beutler_alpha"] = self._beutler_alpha # Create a dynamics object. d = mols.dynamics( @@ -2325,10 +2334,10 @@ def _initialise_gpu_memory(self): # Store soft-core parameters as scalars. self._sc_softcore_form = _np.int32(int(self._softcore_form)) - self._sc_coulomb_power = _np.float32(self._coulomb_power) self._sc_shift_coulomb = _np.float32(self._shift_coulomb.value()) self._sc_shift_delta = _np.float32(self._shift_delta.value()) self._sc_taylor_power = _np.int32(self._taylor_power) + self._sc_beutler_alpha = _np.float32(self._beutler_alpha) # Store immutable per-atom buffers on GPU. self._gpu_sigma = sigmas diff --git a/src/loch/_softcore.py b/src/loch/_softcore.py index 6d165af..40af129 100644 --- a/src/loch/_softcore.py +++ b/src/loch/_softcore.py @@ -29,3 +29,4 @@ class SoftcoreForm(IntEnum): ZACHARIAS = 0 TAYLOR = 1 + BEUTLER = 2 diff --git a/tests/test_energy.py b/tests/test_energy.py index 9235cca..5493069 100644 --- a/tests/test_energy.py +++ b/tests/test_energy.py @@ -20,6 +20,7 @@ ("bpti", "zacharias"), ("sd12", "zacharias"), ("sd12", "taylor"), + ("sd12", "beutler"), ], ) def test_energy(fixture, softcore_form, platform, request): @@ -56,6 +57,9 @@ def test_energy(fixture, softcore_form, platform, request): dyn_map = {} if sampler._softcore_form == SoftcoreForm.TAYLOR: dyn_map["use_taylor_softening"] = True + elif sampler._softcore_form == SoftcoreForm.BEUTLER: + dyn_map["use_beutler_softening"] = True + dyn_map["beutler_alpha"] = sampler._beutler_alpha # Create a dynamics object using the modified GCMC system. d = sampler.system().dynamics( @@ -67,7 +71,6 @@ def test_energy(fixture, softcore_form, platform, request): timestep="2 fs", schedule=schedule, lambda_value=lambda_value, - coulomb_power=sampler._coulomb_power, shift_coulomb=str(sampler._shift_coulomb), shift_delta=str(sampler._shift_delta), platform=platform, @@ -227,7 +230,6 @@ def test_platform_consistency(fixture, request): timestep="2 fs", schedule=schedule, lambda_value=lambda_value, - coulomb_power=cuda_sampler._coulomb_power, shift_coulomb=str(cuda_sampler._shift_coulomb), shift_delta=str(cuda_sampler._shift_delta), platform="cuda", @@ -242,7 +244,6 @@ def test_platform_consistency(fixture, request): timestep="2 fs", schedule=schedule, lambda_value=lambda_value, - coulomb_power=opencl_sampler._coulomb_power, shift_coulomb=str(opencl_sampler._shift_coulomb), shift_delta=str(opencl_sampler._shift_delta), platform="opencl", @@ -347,7 +348,6 @@ def test_energy_regression(fixture, platform, request): timestep="2 fs", schedule=schedule, lambda_value=lambda_value, - coulomb_power=sampler._coulomb_power, shift_coulomb=str(sampler._shift_coulomb), shift_delta=str(sampler._shift_delta), platform=platform, @@ -414,7 +414,6 @@ def _create_and_run(seed): timestep="2 fs", schedule=schedule, lambda_value=0.5, - coulomb_power=sampler._coulomb_power, shift_coulomb=str(sampler._shift_coulomb), shift_delta=str(sampler._shift_delta), platform=platform, From c16f401c87af1161db538f66ff6819629cf4deb4 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 1 May 2026 09:56:06 +0100 Subject: [PATCH 3/7] Optimise softcore. --- src/loch/_kernels.py | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py index d402b89..146accc 100644 --- a/src/loch/_kernels.py +++ b/src/loch/_kernels.py @@ -551,14 +551,13 @@ const float e = sqrtf(e0 * e1); const float a = alpha[idx_atom]; - // Compute the distance between the atoms. - float r = sqrtf(r2); + // Clamp r2 to avoid singularities. + const float r2_sc = (r2 < 1e-6f) ? 1e-6f : r2; - // Truncate the distance. - if (r < 0.001f) - { - r = 0.001f; - } + // Precompute r^6 and sigma^6 using r2 directly (avoids sqrtf and powf). + const float r6 = r2_sc * r2_sc * r2_sc; + const float s2 = s * s; + const float s6_val = s2 * s2 * s2; // Compute the LJ interaction using the chosen soft-core form. float sig6; @@ -567,8 +566,6 @@ { // Taylor soft-core LJ: // sig6 = sigma^6 / (alpha^m * sigma^6 + r^6) - const float s6_val = powf(s, 6.0f); - const float r6 = r * r * r * r * r * r; const float alpha_m = (sc_taylor_power == 1) ? a : (sc_taylor_power == 0) ? 1.0f : powf(a, (float)sc_taylor_power); @@ -579,8 +576,6 @@ // Beutler soft-core LJ: // sig6 = sigma^6 / (sc_beutler_alpha * sigma^6 * alpha + r^6) // V_LJ = (1 - alpha) * 4 * epsilon * sig6 * (sig6 - 1) - const float s6_val = powf(s, 6.0f); - const float r6 = r * r * r * r * r * r; sig6 = s6_val / (sc_beutler_alpha * s6_val * a + r6); lj_prefactor = 1.0f - a; } @@ -590,14 +585,15 @@ // sig6 = sigma^6 / (sigma*delta + r^2)^3 // delta = shift_delta * alpha const float delta_lj = sc_shift_delta * a; - sig6 = powf(s, 6.0f) / powf((s * delta_lj) + (r * r), 3.0f); + const float denom = (s * delta_lj) + r2_sc; + sig6 = s6_val / (denom * denom * denom); } energy_lj[idx] += lj_prefactor * 4.0f * e * sig6 * (sig6 - 1.0f); // Compute the Coulomb interaction. energy_coul[idx] += (q0 * q1) * ((1.0f / sqrtf((sc_shift_coulomb * sc_shift_coulomb * a) - + (r * r))) + (rf_kappa * r2) - rf_correction); + + r2_sc)) + (rf_kappa * r2) - rf_correction); } } From 0845bff4e29b4ca663e2cf8a9941b4f438fb7cb7 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 6 May 2026 15:02:51 +0100 Subject: [PATCH 4/7] Recombine ghost nonbonded forces for performance. --- src/loch/_sampler.py | 115 +++++++++++++------------------------------ 1 file changed, 34 insertions(+), 81 deletions(-) diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index d2e7ac0..d367194 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -719,8 +719,7 @@ def __init__( # Null the nonbonded forces. self._nonbonded_force = None - self._custom_coulomb_force = None - self._custom_lj_force = None + self._custom_nonbonded_force = None # Flag for whether the last move was a bulk sampling move. self._is_bulk = False @@ -1160,8 +1159,7 @@ def reset(self) -> None: # Clear the forces. self._nonbonded_force = None - self._custom_coulomb_force = None - self._custom_lj_force = None + self._custom_nonbonded_force = None # Clear the OpenMM context. self._openmm_context = None @@ -2184,38 +2182,32 @@ def _initialise_gpu_memory(self): map=_map, ) - # Flags for the required force. - has_gng_coul = False - has_gng_lj = False + # Flag for the required force. + has_gng = False # Find the required forces. for force in d.context().getSystem().getForces(): - if force.getName() == "GhostNonGhostCoulombForce": - gng_coul_force = force - has_gng_coul = True - elif force.getName() == "GhostNonGhostLJForce": - gng_lj_force = force - has_gng_lj = True + if force.getName() == "GhostNonGhostNonbondedForce": + gng_force = force + has_gng = True + break # Make sure the force was found. - if not has_gng_coul: - raise ValueError( - "Could not find the GhostNonGhostCoulombForce in the system" - ) - if not has_gng_lj: + if not has_gng: raise ValueError( - "Could not find the GhostNonGhostLJForce in the system" + "Could not find the GhostNonGhostNonbondedForce in the system" ) - # Get the parameters for the GhostNonGhost nonbonded forces. + # Get the parameters for the GhostNonGhostNonbondedForce. charges = _np.zeros(self._num_atoms, dtype=_np.float32) sigmas = _np.zeros(self._num_atoms, dtype=_np.float32) epsilons = _np.zeros(self._num_atoms, dtype=_np.float32) alphas = _np.zeros(self._num_atoms, dtype=_np.float32) - for i in range(gng_coul_force.getNumParticles()): + for i in range(gng_force.getNumParticles()): # Custom force parameters are returned as floats. - q, alpha, _ = gng_coul_force.getParticleParameters(i) - half_sigma, two_sqrt_epsilon, _ = gng_lj_force.getParticleParameters(i) + q, half_sigma, two_sqrt_epsilon, alpha, _ = ( + gng_force.getParticleParameters(i) + ) # Charge in |e|, sigma in nm, epsilon in kJ/mol. charges[i] = q # Rescale and convert units. @@ -2449,20 +2441,14 @@ def _accept_insertion( ) # Update the custom NonBondedForce parameters. if self._is_fep: - self._custom_coulomb_force.setParticleParameters( + self._custom_nonbonded_force.setParticleParameters( start_idx + i, ( self._water_charge[i], - 0.0, - 0.0, - ), - ) - self._custom_lj_force.setParticleParameters( - start_idx + i, - ( self._water_sigma_custom[i], self._water_epsilon_custom[i], 0.0, + 0.0, ), ) @@ -2474,8 +2460,7 @@ def _accept_insertion( # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_coulomb_force.updateParametersInContext(context) - self._custom_lj_force.updateParametersInContext(context) + self._custom_nonbonded_force.updateParametersInContext(context) # Update the state of the water on the GPU. self._kernels["update_water"]( @@ -2531,20 +2516,14 @@ def _accept_deletion(self, idx, context): ) # Update the CustomNonBondedForce parameters. if self._is_fep: - self._custom_coulomb_force.setParticleParameters( + self._custom_nonbonded_force.setParticleParameters( start_idx + i, ( 0.0, - 0.0, - 0.0, - ), - ) - self._custom_lj_force.setParticleParameters( - start_idx + i, - ( self._water_sigma_custom[i], 0.0, 0.0, + 0.0, ), ) @@ -2553,8 +2532,7 @@ def _accept_deletion(self, idx, context): # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_coulomb_force.updateParametersInContext(context) - self._custom_lj_force.updateParametersInContext(context) + self._custom_nonbonded_force.updateParametersInContext(context) # Update the state of the water on the GPU. self._kernels["update_water"]( @@ -2615,20 +2593,14 @@ def _reject_deletion(self, idx, context): ) # Update the CustomNonBondedForce parameters. if self._is_fep: - self._custom_coulomb_force.setParticleParameters( + self._custom_nonbonded_force.setParticleParameters( start_idx + i, ( self._water_charge[i], - 0.0, - 0.0, - ), - ) - self._custom_lj_force.setParticleParameters( - start_idx + i, - ( self._water_sigma_custom[i], self._water_epsilon_custom[i], 0.0, + 0.0, ), ) @@ -2637,8 +2609,7 @@ def _reject_deletion(self, idx, context): # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_coulomb_force.updateParametersInContext(context) - self._custom_lj_force.updateParametersInContext(context) + self._custom_nonbonded_force.updateParametersInContext(context) # Update the state of the water on the GPU. self._kernels["update_water"]( @@ -2705,8 +2676,7 @@ def _set_water_state(self, context, indices=None, states=None, force=False): # Assume the context has been recreated, so we need to get the # new forces. self._nonbonded_force = None - self._custom_coulomb_force = None - self._custom_lj_force = None + self._custom_nonbonded_force = None # Update even if the state is unchanged. force = True @@ -2736,20 +2706,14 @@ def _set_water_state(self, context, indices=None, states=None, force=False): ) # Update the CustomNonbondedForce parameters. if self._is_fep: - self._custom_coulomb_force.setParticleParameters( + self._custom_nonbonded_force.setParticleParameters( start_idx + i, ( 0.0, - 0.0, - 0.0, - ), - ) - self._custom_lj_force.setParticleParameters( - start_idx + i, - ( self._water_sigma_custom[i], 0.0, 0.0, + 0.0, ), ) @@ -2789,20 +2753,14 @@ def _set_water_state(self, context, indices=None, states=None, force=False): ) # Update the CustomNonBondedForce parameters. if self._is_fep: - self._custom_coulomb_force.setParticleParameters( + self._custom_nonbonded_force.setParticleParameters( start_idx + i, ( self._water_charge[i], - 0.0, - 0.0, - ), - ) - self._custom_lj_force.setParticleParameters( - start_idx + i, - ( self._water_sigma_custom[i], self._water_epsilon_custom[i], 0.0, + 0.0, ), ) @@ -2838,8 +2796,7 @@ def _set_water_state(self, context, indices=None, states=None, force=False): # Update the CustomNonbondedForce parameters in the context. if self._is_fep: - self._custom_coulomb_force.updateParametersInContext(context) - self._custom_lj_force.updateParametersInContext(context) + self._custom_nonbonded_force.updateParametersInContext(context) def _set_nonbonded_forces(self, context): """ @@ -2852,15 +2809,13 @@ def _set_nonbonded_forces(self, context): The OpenMM context to use. """ if self._nonbonded_force is None or ( - self._is_fep and self._custom_coulomb_force is None + 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 - elif self._is_fep and force.getName() == "GhostNonGhostCoulombForce": - self._custom_coulomb_force = force - elif self._is_fep and force.getName() == "GhostNonGhostLJForce": - self._custom_lj_force = force + elif self._is_fep and force.getName() == "GhostNonGhostNonbondedForce": + self._custom_nonbonded_force = force elif "Barostat" in force.getName(): msg = ( f"GCMC must be used at constant volume: " @@ -2874,9 +2829,7 @@ def _set_nonbonded_forces(self, context): _logger.error(msg) raise ValueError(msg) - if self._is_fep and ( - self._custom_coulomb_force is None or self._custom_lj_force is None - ): + if self._is_fep and self._custom_nonbonded_force is None: msg = "Could not find a CustomNonbondedForce in the system" _logger.error(msg) raise ValueError(msg) From 12a8c94a955cb909a59eb083cc872d7203f803db Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Wed, 6 May 2026 16:04:06 +0100 Subject: [PATCH 5/7] Use Beutler LJ soft-core for decoupling and add support for LRC. --- src/loch/_utils.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/src/loch/_utils.py b/src/loch/_utils.py index c84b1c7..2da9cbd 100644 --- a/src/loch/_utils.py +++ b/src/loch/_utils.py @@ -36,6 +36,7 @@ def excess_chemical_potential( runtime: str = "5 ns", num_lambda: int = 24, replica_exchange: bool = False, + use_dispersion_correction: bool = False, work_dir: _Optional[str] = None, ) -> float: """ @@ -66,6 +67,9 @@ def excess_chemical_potential( replica_exchange: bool, optional Whether to use replica exchange during the calculation (default is False). + use_dispersion_correction: bool, optional + Whether to include the long-range dispersion correction (default is False). + work_dir: str, optional Working directory for the decoupling simulation (default is None, which uses a temporary directory). @@ -120,6 +124,9 @@ def excess_chemical_potential( if not isinstance(replica_exchange, bool): raise TypeError("'replica_exchange' must be a of type 'bool'.") + if not isinstance(use_dispersion_correction, bool): + raise TypeError("'use_dispersion_correction' must be a of type 'bool'.") + if work_dir is not None: if not isinstance(work_dir, str): raise TypeError("'work_dir' must be a of type 'str'.") @@ -154,7 +161,8 @@ def excess_chemical_potential( runtime=runtime, timestep="2 fs", h_mass_factor=1, - shift_delta="2.25 A", + soft_core_form="beutler", + use_dispersion_correction=use_dispersion_correction, output_directory=work_dir, ) except Exception as e: @@ -229,6 +237,7 @@ def standard_volume( cutoff: str = "10 A", num_samples: int = 5000, sample_interval: str = "1 ps", + use_dispersion_correction: bool = False, ) -> float: """ Calculate the standard volume of water at the given temperature and pressure. @@ -254,6 +263,9 @@ def standard_volume( sample_interval: str, optional Interval at which to sample the volume (default is "1 ps"). + use_dispersion_correction: bool, optional + Whether to include the long-range dispersion correction (default is False). + Returns ------- @@ -304,6 +316,9 @@ def standard_volume( if not u.has_same_units(sr.units.picosecond): raise ValueError("'sample_interval' has incorrect units.") + if not isinstance(use_dispersion_correction, bool): + raise TypeError("'use_dispersion_correction' must be a of type 'bool'.") + # Disable the dynamics progress bar. sr.base.ProgressBar.set_silent() @@ -319,7 +334,12 @@ def standard_volume( # Set up the NPT simulation. try: - d = system.dynamics(temperature=temperature, pressure=pressure, timestep="2 fs") + d = system.dynamics( + temperature=temperature, + pressure=pressure, + timestep="2 fs", + map={"use_dispersion_correction": use_dispersion_correction}, + ) except Exception as e: raise ValueError(f"Unable to set up NPT dynamics: {e}") From 80636a3d96a7b1cf8054c150c213b1d40afadc66 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 7 May 2026 15:08:34 +0100 Subject: [PATCH 6/7] Reverse lambda schedule when swap_end_states=True. --- src/loch/_sampler.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index d367194..297bf9a 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -545,6 +545,9 @@ def __init__( raise ValueError("'swap_end_states' must be of type 'bool'") self._swap_end_states = swap_end_states + if swap_end_states and self._lambda_schedule is not None: + self._lambda_schedule = self._lambda_schedule.reverse() + # Check for waters and validate the template. try: self._water_template = system["water and not property is_perturbable"][0] From fcd8e96ad5cd545024d169fcde4357bc72f19786 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 7 May 2026 15:14:08 +0100 Subject: [PATCH 7/7] Update CHANGELOG. --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d5cee1e..319433c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ Changelog * Add support for getting and restoring sampling statistics. * Handle XED force field virtual sites. +* Add support for long-range Lennard-Jones dispersion correction. +* Add support for Beutler soft-core Lennard-Jones form. [2025.2.0](https://github.com/openbiosim/loch/compare/2025.1.0...2025.2.0) - Feb 2026 -------------------------------------------------------------------------------------