From 065ef75579fc04735745e1190d28ab08b4a5fa13 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Sun, 29 Mar 2026 10:03:34 +0100 Subject: [PATCH 1/2] Add support for different alchemical softcore forms. --- src/loch/__init__.py | 1 + src/loch/_kernels.py | 31 +++++++++++++++++++++++------ src/loch/_sampler.py | 45 +++++++++++++++++++++++++++++++++++++++++++ src/loch/_softcore.py | 31 +++++++++++++++++++++++++++++ tests/test_energy.py | 21 +++++++++++++++++--- 5 files changed, 120 insertions(+), 9 deletions(-) create mode 100644 src/loch/_softcore.py diff --git a/src/loch/__init__.py b/src/loch/__init__.py index 45a2c0c..2c8d0e7 100644 --- a/src/loch/__init__.py +++ b/src/loch/__init__.py @@ -20,5 +20,6 @@ ##################################################################### from ._sampler import * +from ._softcore import * from ._utils import * from ._version import __version__ diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py index ba20ebb..83c7114 100644 --- a/src/loch/_kernels.py +++ b/src/loch/_kernels.py @@ -394,9 +394,11 @@ float rf_cutoff, float rf_kappa, float rf_correction, + int softcore_form, float sc_coulomb_power, float sc_shift_coulomb, - float sc_shift_delta) + float sc_shift_delta, + int sc_taylor_power) { // Work out the atom index. const int idx_atom = GET_GLOBAL_ID(0); @@ -538,7 +540,7 @@ energy_coul[idx] += (q0 * q1) * (r_inv + (rf_kappa * r2) - rf_correction); } - // Zacharias soft-core potential. + // Soft-core potential for ghost atoms. else { // Store required parameters. @@ -558,10 +560,27 @@ r = 0.001f; } - // Compute the Lennard-Jones interaction. - const float delta_lj = sc_shift_delta * a; - const float s6 = powf(s, 6.0f) / powf((s * delta_lj) + (r * r), 3.0f); - energy_lj[idx] += 4.0f * e * s6 * (s6 - 1.0f); + // Compute the LJ interaction using the chosen soft-core form. + float sig6; + if (softcore_form == 1) + { + // 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 : powf(a, (float)sc_taylor_power); + sig6 = s6_val / (alpha_m * s6_val + r6); + } + else + { + // Zacharias soft-core LJ: + // 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); + } + energy_lj[idx] += 4.0f * e * sig6 * (sig6 - 1.0f); // Compute the Coulomb power expression. float cpe; diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index 2bd98a0..2fe278f 100644 --- a/src/loch/_sampler.py +++ b/src/loch/_sampler.py @@ -37,6 +37,7 @@ from ._platforms import create_backend as _create_backend from ._platforms._rng import RNGManager as _RNGManager +from ._softcore import SoftcoreForm as _SoftcoreForm def _as_float32(arr: _np.ndarray) -> _np.ndarray: @@ -81,6 +82,8 @@ def __init__( coulomb_power: float = 0.0, shift_coulomb: str = "1 A", shift_delta: str = "1.5 A", + softcore_form: str = "zacharias", + taylor_power: int = 1, swap_end_states: bool = False, restart: bool = False, overwrite: bool = False, @@ -207,6 +210,17 @@ def __init__( The soft-core shift-delta parameter. This is used to soften the Lennard-Jones interaction. + softcore_form : str + The soft-core potential form to use for alchemical interactions. + This can be either 'zacharias' or 'taylor'. The default is + 'zacharias'. + + taylor_power : int + The power to use for the alpha term in the Taylor soft-core LJ + expression, i.e. sig6 = sigma^6 / (alpha^m * sigma^6 + r^6). + Must be between 0 and 4. The default is 1. Only used when + softcore_form is 'taylor'. + swap_end_states: bool Whether to swap the end states of the alchemical systems. @@ -495,6 +509,26 @@ def __init__( except Exception as e: raise ValueError(f"Could not validate the 'shift_delta': {e}") + if not isinstance(softcore_form, str): + raise TypeError("'softcore_form' must be of type 'str'") + softcore_form = softcore_form.lower().replace(" ", "") + _valid_softcore_forms = {m.name.lower(): m for m in _SoftcoreForm} + if softcore_form not in _valid_softcore_forms: + raise ValueError( + f"'softcore_form' not recognised. Valid forms are: " + f"{', '.join(_valid_softcore_forms)}" + ) + self._softcore_form = _valid_softcore_forms[softcore_form] + + if not isinstance(taylor_power, int): + try: + taylor_power = int(taylor_power) + except Exception: + raise ValueError("'taylor_power' must be of type 'int'") + if not 0 <= taylor_power <= 4: + raise ValueError("'taylor_power' must be between 0 and 4") + self._taylor_power = taylor_power + if not isinstance(swap_end_states, bool): raise ValueError("'swap_end_states' must be of type 'bool'") self._swap_end_states = swap_end_states @@ -1323,9 +1357,11 @@ def move(self, context: _openmm.Context) -> list[int]: self._rf_cutoff, 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, block=(self._num_threads, 1, 1), grid=(self._atom_blocks, self._batch_size, 1), ) @@ -1902,6 +1938,12 @@ def _initialise_gpu_memory(self): # Link to the reference state. mols = _sr.morph.link_to_reference(self._system) + # Build map of extra options for the dynamics object. + _map = {} + if self._softcore_form == _SoftcoreForm.TAYLOR: + _map["use_taylor_softening"] = True + _map["taylor_power"] = self._taylor_power + # Create a dynamics object. d = mols.dynamics( cutoff_type=self._cutoff, @@ -1916,6 +1958,7 @@ def _initialise_gpu_memory(self): rest2_selection=self._rest2_selection, swap_end_states=self._swap_end_states, platform="cpu", + map=_map, ) # Flags for the required force. @@ -2059,9 +2102,11 @@ 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) # Store immutable per-atom buffers on GPU. self._gpu_sigma = sigmas diff --git a/src/loch/_softcore.py b/src/loch/_softcore.py new file mode 100644 index 0000000..6d165af --- /dev/null +++ b/src/loch/_softcore.py @@ -0,0 +1,31 @@ +###################################################################### +# Loch: GPU accelerated GCMC water sampling engine. +# +# Copyright: 2025-2026 +# +# Authors: The OpenBioSim Team +# +# Loch is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# Loch is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with Loch. If not, see . +##################################################################### + +__all__ = ["SoftcoreForm"] + +from enum import IntEnum + + +class SoftcoreForm(IntEnum): + """Enum for the soft-core potential form used for alchemical interactions.""" + + ZACHARIAS = 0 + TAYLOR = 1 diff --git a/tests/test_energy.py b/tests/test_energy.py index 50e115b..9235cca 100644 --- a/tests/test_energy.py +++ b/tests/test_energy.py @@ -5,7 +5,7 @@ import sire as sr -from loch import GCMCSampler +from loch import GCMCSampler, SoftcoreForm @pytest.mark.skipif( @@ -13,8 +13,16 @@ reason="Requires CUDA enabled GPU.", ) @pytest.mark.parametrize("platform", ["cuda", "opencl"]) -@pytest.mark.parametrize("fixture", ["water_box", "bpti", "sd12"]) -def test_energy(fixture, platform, request): +@pytest.mark.parametrize( + "fixture,softcore_form", + [ + ("water_box", "zacharias"), + ("bpti", "zacharias"), + ("sd12", "zacharias"), + ("sd12", "taylor"), + ], +) +def test_energy(fixture, softcore_form, platform, request): """ Test that the RF energy difference agrees with OpenMM. """ @@ -36,6 +44,7 @@ def test_energy(fixture, platform, request): reference=reference, lambda_schedule=schedule, lambda_value=lambda_value, + softcore_form=softcore_form, log_level="debug", ghost_file=None, log_file=None, @@ -43,6 +52,11 @@ def test_energy(fixture, platform, request): platform=platform, ) + # Build map of extra options for the dynamics object. + dyn_map = {} + if sampler._softcore_form == SoftcoreForm.TAYLOR: + dyn_map["use_taylor_softening"] = True + # Create a dynamics object using the modified GCMC system. d = sampler.system().dynamics( cutoff_type="rf", @@ -57,6 +71,7 @@ def test_energy(fixture, platform, request): shift_coulomb=str(sampler._shift_coulomb), shift_delta=str(sampler._shift_delta), platform=platform, + map=dyn_map, ) # Loop until we accept an insertion move. From a6992fccc9093638bb09702406b11ddc1eecac2e Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Sun, 29 Mar 2026 10:21:30 +0100 Subject: [PATCH 2/2] Optimise taylor_power branch. --- src/loch/_kernels.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/loch/_kernels.py b/src/loch/_kernels.py index 83c7114..8165486 100644 --- a/src/loch/_kernels.py +++ b/src/loch/_kernels.py @@ -568,8 +568,9 @@ // 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 : powf(a, (float)sc_taylor_power); + const float alpha_m = (sc_taylor_power == 1) ? a + : (sc_taylor_power == 0) ? 1.0f + : powf(a, (float)sc_taylor_power); sig6 = s6_val / (alpha_m * s6_val + r6); } else