Add support for long-range dispersion correction, Beutler soft-core, and LambdaSchedule.reverse()

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
 -------------------------------------------------------------------------------------

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);
 
                             }

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,15 +212,20 @@ 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
             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'.
 
+        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

src/loch/_softcore.py

@@ -29,3 +29,4 @@ class SoftcoreForm(IntEnum):
 
     ZACHARIAS = 0
     TAYLOR = 1
+    BEUTLER = 2