Merge pull request #16 from OpenBioSim/feature_softcore

Add support for different alchemical softcore forms

Author: lohedges

src/loch/__init__.py

@@ -20,5 +20,6 @@
 #####################################################################
 
 from ._sampler import *
+from ._softcore import *
 from ._utils import *
 from ._version import __version__

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,28 @@
                                     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
+                                        : (sc_taylor_power == 0) ? 1.0f
+                                        : 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;

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

src/loch/_softcore.py

@@ -0,0 +1,31 @@
+######################################################################
+# Loch: GPU accelerated GCMC water sampling engine.
+#
+# Copyright: 2025-2026
+#
+# Authors: The OpenBioSim Team <team@openbiosim.org>
+#
+# 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 <http://www.gnu.org/licenses/>.
+#####################################################################
+
+__all__ = ["SoftcoreForm"]
+
+from enum import IntEnum
+
+
+class SoftcoreForm(IntEnum):
+    """Enum for the soft-core potential form used for alchemical interactions."""
+
+    ZACHARIAS = 0
+    TAYLOR = 1

tests/test_energy.py

@@ -5,16 +5,24 @@
 
 import sire as sr
 
-from loch import GCMCSampler
+from loch import GCMCSampler, SoftcoreForm
 
 
 @pytest.mark.skipif(
     "CUDA_VISIBLE_DEVICES" not in os.environ,
     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,13 +44,19 @@ 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,
         test=True,
         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.