Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/loch/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,5 +20,6 @@
#####################################################################

from ._sampler import *
from ._softcore import *
from ._utils import *
from ._version import __version__
32 changes: 26 additions & 6 deletions src/loch/_kernels.py
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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.
Expand All @@ -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;
Expand Down
45 changes: 45 additions & 0 deletions src/loch/_sampler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
)
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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
Expand Down
31 changes: 31 additions & 0 deletions src/loch/_softcore.py
Original file line number Diff line number Diff line change
@@ -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
21 changes: 18 additions & 3 deletions tests/test_energy.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand All @@ -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",
Expand All @@ -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.
Expand Down
Loading