Skip to content

Commit

Permalink
Merge branch 'devel' into bugfix-force-specific-equation [ci skip]
Browse files Browse the repository at this point in the history
  • Loading branch information
@fjclark
fjclark committed Jan 30, 2025
2 parents 5c2dacf + 961e15e commit 3cbe190
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 16 deletions.
2 changes: 2 additions & 0 deletions doc/source/changelog.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ organisation on `GitHub <https://github.com/openbiosim/sire>`__.

* Added support for angle and dihedral restraints which can be used in alchemical and standard simulations.

* Allow user to compute energy trajectory over a subset of the lambda windows for each lambda.

* Fix the ``hasForceSpecificEquation`` function in the ``LambdaLever`` class so that it returns true if
there is a default equation for the force.

Expand Down
113 changes: 97 additions & 16 deletions src/sire/mol/_dynamics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -208,12 +208,15 @@ def __init__(self, mols=None, map=None, **kwargs):
)

# Store all the perturbable molecules associated with the selection
# and remove perturbable atoms from the selection.
# and remove perturbable atoms from the selection. Remove alchemical ions
# from the selection.
pert_mols = {}
non_pert_atoms = atoms.to_list()
for atom in atoms:
mol = atom.molecule()
if mol.has_property("is_perturbable"):
if mol.has_property("is_alchemical_ion"):
non_pert_atoms.remove(atom)
elif mol.has_property("is_perturbable"):
non_pert_atoms.remove(atom)
if mol.number() not in pert_mols:
pert_mols[mol.number()] = [atom]
Expand All @@ -240,6 +243,20 @@ def __init__(self, mols=None, map=None, **kwargs):
# Update the system.
self._sire_mols.update(mol)

# Search for alchemical ions and exclude them via a REST2 mask.
try:
for mol in self._sire_mols.molecules("property is_alchemical_ion"):
is_rest2 = [False] * mol.num_atoms()
mol = (
mol.edit()
.set_property("is_rest2", is_rest2)
.molecule()
.commit()
)
self._sire_mols.update(mol)
except:
pass

from ..convert import to

self._omm_mols = to(self._sire_mols, "openmm", map=self._map)
Expand Down Expand Up @@ -341,6 +358,8 @@ def _exit_dynamics_block(
rest2_scale_factors=[],
save_velocities: bool = False,
delta_lambda: float = None,
num_energy_neighbours: int = None,
null_energy: float = None,
):
if not self._is_running:
raise SystemError("Cannot stop dynamics that is not running!")
Expand Down Expand Up @@ -370,6 +389,14 @@ def _exit_dynamics_block(
self._elapsed_time = current_time
self._current_time += delta

# store the number of lambda windows
if lambda_windows is not None:
num_lambda_windows = len(lambda_windows)

# compute energies for all windows
if num_energy_neighbours is None:
num_energy_neighbours = num_lambda_windows

if save_energy:
# should save energy here
nrgs = {}
Expand All @@ -391,7 +418,7 @@ def _exit_dynamics_block(
sim_lambda_value = self._omm_mols.get_lambda()
sim_rest2_scale = self._omm_mols.get_rest2_scale()

# Store the potential energy and accumulated non-equilibrium work.
# store the potential energy and accumulated non-equilibrium work
if self._is_interpolate:
nrg = nrgs["potential"]

Expand All @@ -406,21 +433,32 @@ def _exit_dynamics_block(
nrgs[str(sim_lambda_value)] = nrgs["potential"]

if lambda_windows is not None:
for lambda_value, rest2_scale in zip(
lambda_windows, rest2_scale_factors
# get the index of the simulation lambda value in the
# lambda windows list
try:
lambda_index = lambda_windows.index(sim_lambda_value)
except:
num_energy_neighbours = num_lambda_windows
lambda_index = i

for i, (lambda_value, rest2_scale) in enumerate(
zip(lambda_windows, rest2_scale_factors)
):
if lambda_value != sim_lambda_value:
self._omm_mols.set_lambda(
lambda_value,
rest2_scale=rest2_scale,
update_constraints=False,
)
nrgs[str(lambda_value)] = (
self._omm_mols.get_potential_energy(
to_sire_units=False
).value_in_unit(openmm.unit.kilocalorie_per_mole)
* kcal_per_mol
)
if abs(lambda_index - i) <= num_energy_neighbours:
self._omm_mols.set_lambda(
lambda_value,
rest2_scale=rest2_scale,
update_constraints=False,
)
nrgs[str(lambda_value)] = (
self._omm_mols.get_potential_energy(
to_sire_units=False
).value_in_unit(openmm.unit.kilocalorie_per_mole)
* kcal_per_mol
)
else:
nrgs[str(lambda_value)] = null_energy * kcal_per_mol

self._omm_mols.set_lambda(
sim_lambda_value,
Expand Down Expand Up @@ -893,6 +931,8 @@ def run(
save_frame_on_exit: bool = False,
save_energy_on_exit: bool = False,
auto_fix_minimise: bool = True,
num_energy_neighbours: int = None,
null_energy: str = None,
):
if self.is_null():
return
Expand All @@ -908,6 +948,8 @@ def run(
"save_frame_on_exit": save_frame_on_exit,
"save_energy_on_exit": save_energy_on_exit,
"auto_fix_minimise": auto_fix_minimise,
"num_energy_neighbours": num_energy_neighbours,
"null_energy": null_energy,
}

from concurrent.futures import ThreadPoolExecutor
Expand All @@ -927,6 +969,17 @@ def run(
if energy_frequency is not None:
energy_frequency = u(energy_frequency)

if null_energy is not None:
null_energy = u(null_energy)
else:
null_energy = u("10000 kcal/mol")

if num_energy_neighbours is not None:
try:
num_energy_neighbours = int(num_energy_neighbours)
except:
num_energy_neighbours = len(lambda_windows)

try:
steps_to_run = int(time.to(picosecond) / self.timestep().to(picosecond))
except Exception:
Expand Down Expand Up @@ -1187,6 +1240,8 @@ class NeedsMinimiseError(Exception):
rest2_scale_factors=rest2_scale_factors,
save_velocities=save_velocities,
delta_lambda=delta_lambda,
num_energy_neighbours=num_energy_neighbours,
null_energy=null_energy.value(),
)

saved_last_frame = False
Expand Down Expand Up @@ -1459,6 +1514,8 @@ def run(
save_frame_on_exit: bool = False,
save_energy_on_exit: bool = False,
auto_fix_minimise: bool = True,
num_energy_neighbours: int = None,
null_energy: str = None,
):
"""
Perform dynamics on the molecules.
Expand Down Expand Up @@ -1538,6 +1595,28 @@ def run(
Such failures often indicate that the system needs
minimsing. This automatically runs the minimisation
in these cases, and then runs the requested dynamics.
num_energy_neighbours: int
The number of neighbouring windows to use when computing
the energy trajectory for the simulation lambda value.
This can be used to compute energies over a subset of the
values in 'lambda_windows', hence reducing the cost of
computing the energy trajectory. Note that the simulation
lambda value must be contained in 'lambda_windows', so it
is recommended that the values are rounded. A value of
'null_energy' will be added to the energy trajectory for the
lambda windows that are omitted. Note that a similar result
can be achieved by simply removing any lambda values from
'lambda_windows' that you don't want, but this will result
in an energy trajectory that only contains results for
the specified lambda values. If None, then all lambda windows
will be used.
null_energy: str
The energy value to use for lambda windows that are not
being computed as part of the energy trajectory, i.e. when
'num_energy_neighbours' is less than len(lambda_windows).
By default, a value of '10000 kcal mol-1' is used.
"""
if not self._d.is_null():
if save_velocities is None:
Expand All @@ -1557,6 +1636,8 @@ def run(
save_frame_on_exit=save_frame_on_exit,
save_energy_on_exit=save_energy_on_exit,
auto_fix_minimise=auto_fix_minimise,
num_energy_neighbours=num_energy_neighbours,
null_energy=null_energy,
)

return self
Expand Down

0 comments on commit 3cbe190

Please sign in to comment.