diff --git a/doc/source/changelog.rst b/doc/source/changelog.rst index 19e3897b3..82d41c920 100644 --- a/doc/source/changelog.rst +++ b/doc/source/changelog.rst @@ -37,6 +37,8 @@ organisation on `GitHub `__. * 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. diff --git a/src/sire/mol/_dynamics.py b/src/sire/mol/_dynamics.py index 994ca36ab..e07863ae6 100644 --- a/src/sire/mol/_dynamics.py +++ b/src/sire/mol/_dynamics.py @@ -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] @@ -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) @@ -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!") @@ -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 = {} @@ -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"] @@ -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, @@ -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 @@ -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 @@ -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: @@ -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 @@ -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. @@ -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: @@ -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