Skip to content

Allow a sampler to be re-used across lambda values #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lohedges merged 1 commit into from
Jul 27, 2026
Merged

Allow a sampler to be re-used across lambda values #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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 CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ Changelog

* Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
* Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
* Add `set_lambda` and `precompute_lambda` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context.

[2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
-------------------------------------------------------------------------------------
Expand Down
335 changes: 268 additions & 67 deletions src/loch/_sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,8 @@ def __init__(
lambda_schedule: _Optional[_Any] = None,
lambda_value: float = 0.0,
rest2_scale: float = 1.0,
lambda_values: _Optional[list[float]] = None,
rest2_scales: _Optional[list[float]] = None,
rest2_selection: _Optional[str] = None,
shift_coulomb: str = "1 A",
shift_delta: str = "1.5 A",
Expand Down Expand Up @@ -205,6 +207,16 @@ def __init__(
(REST2) for alchemical systems. This should specify the temperature
of the REST2 system relative to the rest of the system.

lambda_values: [float]
The lambda values that the sampler will be switched between via
set_lambda(). The non-bonded parameters for each are cached at
construction, so that switching later doesn't need to build an
OpenMM context. If None, only 'lambda_value' is cached.

rest2_scales: [float]
The REST2 scaling factor for each entry of 'lambda_values'. If
None, 'rest2_scale' is used for all of them.

rest2_selection: str
A selection string for atoms to include in the REST2 region in
addition to any perturbable molecules. For example, "molidx 0 and
Expand Down Expand Up @@ -506,6 +518,37 @@ def __init__(
raise ValueError("'rest2_scale' must be greater than or equal to 1.0")
self._rest2_scale = rest2_scale

# Cache of the lambda dependent non-bonded parameters, keyed by
# (lambda_value, rest2_scale). Populated below and consumed by
# set_lambda().
self._lambda_params = {}

if lambda_values is None:
self._lambda_values = None
else:
try:
self._lambda_values = [float(x) for x in lambda_values]
except:
raise ValueError("'lambda_values' must be a list of 'float'")
if not all(0.0 <= x <= 1.0 for x in self._lambda_values):
raise ValueError("'lambda_values' must be between 0 and 1")

if rest2_scales is None:
self._rest2_scales = [rest2_scale] * len(self._lambda_values)
else:
try:
self._rest2_scales = [float(x) for x in rest2_scales]
except:
raise ValueError("'rest2_scales' must be a list of 'float'")
if len(self._rest2_scales) != len(self._lambda_values):
raise ValueError(
"'rest2_scales' must be the same length as 'lambda_values'"
)
if any(x < 1.0 for x in self._rest2_scales):
raise ValueError(
"'rest2_scales' must be greater than or equal to 1.0"
)

if rest2_selection is not None:
if not isinstance(rest2_selection, str):
raise ValueError("'rest2_selection' must be of type 'str'")
Expand Down Expand Up @@ -694,6 +737,17 @@ def __init__(
self._water_blocks = self._num_waters // self._num_threads + 1

# Initialise the GPU memory.
# Cache the non-bonded parameters for every lambda value that the
# sampler will be switched to, so that set_lambda() never has to build
# an OpenMM context mid-simulation. This is done before the GPU memory
# is initialised, so that the current lambda value is part of the same
# scan rather than needing a context of its own.
if self._lambda_values is not None:
self._precompute_lambdas(
[self._lambda_value] + self._lambda_values,
[self._rest2_scale] + self._rest2_scales,
)

self._initialise_gpu_memory()

# Set the box information.
Expand Down Expand Up @@ -2296,6 +2350,210 @@ def _prepare_system(system, water_template, rng, num_ghost_waters):
_np.array(water_residues),
)

def _precompute_lambdas(
self, lambda_values: list[float], rest2_scales: list[float]
) -> None:
"""
Cache the lambda dependent non-bonded parameters for a set of lambda
values, so that set_lambda() can switch between them without touching
OpenMM.

The system only holds the end-state properties, so the parameters are
obtained from an OpenMM context built using the specified lambda
schedule. A single context is created and scanned over the requested
lambda values, re-reading the GhostNonGhostNonbondedForce at each one,
since building the context is by far the expensive part.

This is a no-op for non-alchemical systems, and for lambda values that
have already been cached.

Parameters
----------

lambda_values: [float]
The lambda values to cache.

rest2_scales: [float]
The REST2 scaling factor for each lambda value.
"""

if not self._is_fep:
return

# Only build a context if something is actually missing.
wanted = []
for lambda_value, rest2_scale in zip(lambda_values, rest2_scales):
key = (float(lambda_value), float(rest2_scale))
if key not in self._lambda_params and key not in wanted:
wanted.append(key)

if not wanted:
return

# 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
elif self._softcore_form == _SoftcoreForm.BEUTLER:
_map["use_beutler_softening"] = True
_map["beutler_alpha"] = self._beutler_alpha

# Create a single dynamics object, which is then scanned over the
# requested lambda values.
d = mols.dynamics(
cutoff_type=self._cutoff,
cutoff=self._cutoff,
lambda_value=wanted[0][0],
schedule=self._lambda_schedule,
pressure=None,
timestep="2fs",
constraint="h_bonds",
perturbable_constraint="h_bonds_not_heavy_perturbed",
rest2_scale=wanted[0][1],
rest2_selection=self._rest2_selection,
swap_end_states=self._swap_end_states,
platform="cpu",
map=_map,
)

# Find the required force. set_lambda() updates it in place, so it
# only needs to be located once.
gng_force = None
for force in d.context().getSystem().getForces():
if force.getName() == "GhostNonGhostNonbondedForce":
gng_force = force
break

if gng_force is None:
raise ValueError(
"Could not find the GhostNonGhostNonbondedForce in the system"
)

num_particles = gng_force.getNumParticles()

for lambda_value, rest2_scale in wanted:
d.set_lambda(lambda_value, rest2_scale=rest2_scale)

# Get the parameters for the GhostNonGhostNonbondedForce.
charges = _np.zeros(self._num_atoms, dtype=_np.float32)
sigmas = _np.zeros(self._num_atoms, dtype=_np.float32)
epsilons = _np.zeros(self._num_atoms, dtype=_np.float32)
alphas = _np.zeros(self._num_atoms, dtype=_np.float32)
for i in range(num_particles):
# Custom force parameters are returned as floats.
q, half_sigma, two_sqrt_epsilon, alpha, _ = (
gng_force.getParticleParameters(i)
)
# Charge in |e|, sigma in nm, epsilon in kJ/mol.
charges[i] = q
# Rescale and convert units.
sigmas[i] = _sr.u(f"{2.0 * half_sigma} nm").to("angstrom")
epsilons[i] = _sr.u(f"{(0.5 * two_sqrt_epsilon) ** 2} kJ/mol").to(
"kcal/mol"
)
# Store the softening parameter.
alphas[i] = alpha

self._lambda_params[(lambda_value, rest2_scale)] = (
charges,
sigmas,
epsilons,
alphas,
)

def set_lambda(
self, lambda_value: float, rest2_scale: _Optional[float] = None
) -> None:
"""
Set the lambda value for the sampler, updating the non-bonded
parameters used to evaluate insertion and deletion energies.

Parameters for lambda values passed to the constructor are already
cached. Any other value is computed here, which requires building an
OpenMM context and is slow, though the result is then cached too.

This must be kept consistent with the lambda value of the OpenMM
context that the sampler is used with.

Parameters
----------

lambda_value: float
The lambda value.

rest2_scale: float
The REST2 scaling factor. If None, the current value is retained.
"""

try:
lambda_value = float(lambda_value)
except:
raise ValueError("'lambda_value' must be of type 'float'")
if not 0.0 <= lambda_value <= 1.0:
raise ValueError("'lambda_value' must be between 0 and 1")

if rest2_scale is None:
rest2_scale = self._rest2_scale
else:
try:
rest2_scale = float(rest2_scale)
except:
raise ValueError("'rest2_scale' must be of type 'float'")
if rest2_scale < 1.0:
raise ValueError("'rest2_scale' must be greater than or equal to 1.0")

# Nothing to do.
if lambda_value == self._lambda_value and rest2_scale == self._rest2_scale:
return

# There are no lambda dependent parameters for a non-alchemical system.
if not self._is_fep:
self._lambda_value = lambda_value
self._rest2_scale = rest2_scale
return

# Make sure the parameters are cached. This builds an OpenMM context,
# so is slow; pass 'lambda_values' to the constructor to avoid it.
self._precompute_lambdas([lambda_value], [rest2_scale])

charges, sigmas, epsilons, alphas = self._lambda_params[
(lambda_value, rest2_scale)
]

# Upload the new parameters to the GPU.
self._gpu_charge = self._backend.to_gpu(charges)
self._gpu_sigma = self._backend.to_gpu(sigmas)
self._gpu_epsilon = self._backend.to_gpu(epsilons)
self._gpu_alpha = self._backend.to_gpu(alphas)

self._lambda_value = lambda_value
self._rest2_scale = rest2_scale

def set_ghost_file(self, ghost_file: _Optional[str]) -> None:
"""
Set the file that write_ghost_residues() appends to.

Unlike the constructor, this does not create or truncate the file,
since it is intended for switching between files that are already
being written to.

Parameters
----------

ghost_file: str
The path to the ghost residue file. If None, ghost residues
cannot be written.
"""

if ghost_file is not None and not isinstance(ghost_file, str):
raise TypeError("'ghost_file' must be of type 'str'")

self._ghost_file = ghost_file

def _initialise_gpu_memory(self):
"""
Initialise the GPU memory.
Expand Down Expand Up @@ -2360,76 +2618,19 @@ def _initialise_gpu_memory(self):
# schedule and value, then extract the required properties from the forces
# within the context. (The system just contains the end-state properties.)
else:
# 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
elif self._softcore_form == _SoftcoreForm.BEUTLER:
_map["use_beutler_softening"] = True
_map["beutler_alpha"] = self._beutler_alpha

# Create a dynamics object.
d = mols.dynamics(
cutoff_type=self._cutoff,
cutoff=self._cutoff,
lambda_value=self._lambda_value,
schedule=self._lambda_schedule,
pressure=None,
timestep="2fs",
constraint="h_bonds",
perturbable_constraint="h_bonds_not_heavy_perturbed",
rest2_scale=self._rest2_scale,
rest2_selection=self._rest2_selection,
swap_end_states=self._swap_end_states,
platform="cpu",
map=_map,
)

# Flag for the required force.
has_gng = False

# Find the required forces.
for force in d.context().getSystem().getForces():
if force.getName() == "GhostNonGhostNonbondedForce":
gng_force = force
has_gng = True
break
# Cache the host arrays so that set_lambda() can switch back to
# this lambda value without rebuilding an OpenMM context.
self._precompute_lambdas([self._lambda_value], [self._rest2_scale])

# Make sure the force was found.
if not has_gng:
raise ValueError(
"Could not find the GhostNonGhostNonbondedForce in the system"
)

# Get the parameters for the GhostNonGhostNonbondedForce.
charges = _np.zeros(self._num_atoms, dtype=_np.float32)
sigmas = _np.zeros(self._num_atoms, dtype=_np.float32)
epsilons = _np.zeros(self._num_atoms, dtype=_np.float32)
alphas = _np.zeros(self._num_atoms, dtype=_np.float32)
for i in range(gng_force.getNumParticles()):
# Custom force parameters are returned as floats.
q, half_sigma, two_sqrt_epsilon, alpha, _ = (
gng_force.getParticleParameters(i)
)
# Charge in |e|, sigma in nm, epsilon in kJ/mol.
charges[i] = q
# Rescale and convert units.
sigmas[i] = _sr.u(f"{2.0 * half_sigma} nm").to("angstrom")
epsilons[i] = _sr.u(f"{(0.5 * two_sqrt_epsilon) ** 2} kJ/mol").to(
"kcal/mol"
)
# Store the softening parameter.
alphas[i] = alpha
charges, sigmas, epsilons, alphas = self._lambda_params[
(self._lambda_value, self._rest2_scale)
]

# Convert to GPU arrays.
charges = self._backend.to_gpu(charges.astype(_np.float32))
sigmas = self._backend.to_gpu(sigmas.astype(_np.float32))
epsilons = self._backend.to_gpu(epsilons.astype(_np.float32))
alphas = self._backend.to_gpu(alphas.astype(_np.float32))
charges = self._backend.to_gpu(charges)
sigmas = self._backend.to_gpu(sigmas)
epsilons = self._backend.to_gpu(epsilons)
alphas = self._backend.to_gpu(alphas)

# Create the ghost atom array.
is_ghost_fep = _np.zeros(self._num_atoms, dtype=_np.int32)
Expand Down
Loading