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 ------------------------------------------------------------------------------------- diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py index 60abc96..146accc 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); @@ -561,6 +561,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: @@ -570,6 +571,14 @@ : 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) + sig6 = s6_val / (sc_beutler_alpha * s6_val * a + r6); + lj_prefactor = 1.0f - a; + } else { // Zacharias soft-core LJ: @@ -579,22 +588,11 @@ const float denom = (s * delta_lj) + r2_sc; sig6 = s6_val / (denom * denom * denom); } - 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) + r2_sc)) + (rf_kappa * r2) - rf_correction); } diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 0b9e04e..297bf9a 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. @@ -287,6 +292,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: @@ -489,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") @@ -529,10 +533,21 @@ 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 + 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] @@ -775,7 +790,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}, " @@ -1273,6 +1287,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 +1304,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( @@ -1442,10 +1465,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), ) @@ -1521,6 +1544,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 +1575,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 +1648,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 +1677,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( @@ -2106,6 +2164,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( @@ -2124,7 +2185,7 @@ def _initialise_gpu_memory(self): map=_map, ) - # Flags for the required force. + # Flag for the required force. has_gng = False # Find the required forces. @@ -2268,10 +2329,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 @@ -2321,6 +2382,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 ): @@ -2417,6 +2487,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. @@ -2487,6 +2561,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. @@ -2560,6 +2638,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 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/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}") 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,