def minimize(thermodynamic_state: states.ThermodynamicState,
sampler_state: states.SamplerState,
max_iterations: int = 100) -> states.SamplerState:
"""
Minimize the given system and state, up to a maximum number of steps.
This does not return a copy of the samplerstate; it is an update-in-place.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The state at which the system could be minimized
sampler_state : openmmtools.states.SamplerState
The starting state at which to minimize the system.
max_iterations : int, optional, default 100
The maximum number of minimization steps. Default is 100.
Returns
-------
sampler_state : openmmtools.states.SamplerState
The posititions and accompanying state following minimization
"""
if type(cache.global_context_cache) == cache.DummyContextCache:
integrator = openmm.VerletIntegrator(
1.0) #we won't take any steps, so use a simple integrator
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state, integrator)
_logger.debug(f"using dummy context cache")
else:
_logger.debug(f"using global context cache")
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state)
sampler_state.apply_to_context(context, ignore_velocities=True)
openmm.LocalEnergyMinimizer.minimize(context, maxIterations=max_iterations)
sampler_state.update_from_context(context)
class Propagator(OMMBIP):
"""
Propagator pseudocode:
Step 1: initialization--
set iteration = 0, n_iterations = n_iterations, lambda = 0 (i.e. iteration / n_iterations); work_accumulated = 0.0
generate sample x_0 ~ e^(-p(x))
evaluate work_incremental = 0 (i.e. u_mm(x_0) - g(x_0), but we presume that g = u_mm(.))
work_accumulated <- work_accumulated + work_incremental
x' = x_0
Step 2: sampling
for increment in range(n_iterations):
x = x'
ante_perturbation_potential = (1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)
set iteration <- iteration + 1.0; lambda <- iteration / n_iterations
evaluate work_incremental = [(1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)] - ante_perturbation_potential
work_accumulated <- work_accumulated + work_incremental
create a modified force: modified_f = (1 - lambda) * f_mm + lambda * f_ani_mm_mix (where f_. = -grad(u_.) )
x' = V R O R (where V deterministic update is according to modified_f defined above) w.r.t x
NOTE: in this regime, the last x' is propagated w.r.t. a propagator whose invariant distribution respects u_ani_mm_mix;
this should _not_ be the case. There should be an exception in the Step 2 for loop that breaks once the final work_incremental is computed and updated
to the work_accumulated. Regardless, the distribution of accumulated works is unaffected by this 'bug'; only expectations (as a function of x) w.r.t. these
weights may be affected.
See: 3.1.1. of https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf (esp. Remark 1.)
"""
def __init__(self,
openmm_pdf_state,
openmm_pdf_state_subset,
subset_indices_map,
integrator,
ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter=None,
write_trajectory_interval = 1,
**kwargs):
"""
arguments
openmm_pdf_state : openmmtools.states.ThermodynamicState
the pdf state of the propagator
openmm_pdf_state_subset : openmmtools.states.ThermodynamicState
the pdf state of the atom subset
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
integrator : openmm.Integrator
integrator of dynamics
ani_handler : ANI1_force_and_energy
handler for ani forces and potential energy
context_cache : openmmtools.cache.ContextCache, optional default:None
The ContextCache to use for Context creation. If None, the global cache
openmmtools.cache.global_context_cache is used.
reassign_velocities : bool, optional default:False
If True, the velocities will be reassigned from the Maxwell-Boltzmann
distribution at the beginning of the move.
n_restart_attempts : int, optional default:0
When greater than 0, if after the integration there are NaNs in energies,
the move will restart. When the integrator has a random component, this
may help recovering. On the last attempt, the ``Context`` is
re-initialized in a slower process, but better than the simulation
crashing. An IntegratorMoveError is raised after the given number of
attempts if there are still NaNs.
reporter : coddiwomple.openmm.reporter.OpenMMReporter, default None
a reporter object to write trajectories
write_trajectory_interval : int
frequency of writing trajectory
"""
super().__init__(openmm_pdf_state,
integrator,
context_cache,
reassign_velocities,
n_restart_attempts)
#create a pdf state for the subset indices (usually a vacuum system)
self.pdf_state_subset = openmm_pdf_state_subset
assert self.pdf_state_subset.temperature == self.pdf_state.temperature, f"the temperatures of the pdf states do not match"
#create a dictionary for subset indices
self._subset_indices_map = subset_indices_map
#create an ani handler attribute that can be referenced
self.ani_handler = ani_handler
#create a context for the subset atoms that can be referenced
self.context_subset, _ = cache.global_context_cache.get_context(self.pdf_state_subset)
#create a reporter for the accumulated works
self._state_works = {}
self._state_works_counter = 0
#create a reporter
self._write_trajectory = False if reporter is None else True
self.reporter=reporter
if self._write_trajectory:
from coddiwomple.particles import Particle
self.particle = Particle(0)
self.write_trajectory_interval=write_trajectory_interval
else:
self.particle = None
self.write_trajectory_interval=None
def _initialize_state_works(self):
"""
initialize an empty list and add 0.0 to it (state works)
"""
self._current_state_works = [] #define an interim (auxiliary) list that will track the thermodynamic work of the current application
self._current_state_works.append(0.0) #the first incremental work is always 0 since the importance function is identical to the first target distribution (i.e. fully interacting MM)
def _initialize_iterations(self, n_iterations):
"""
initialize the iteration counter
"""
self._iteration = 0.0 #define the first iteration as 0
self._n_iterations = n_iterations #the number of iterations in the protocol is equal to the number of steps in the application
def _update_particle_state_substate(self, particle_state, new_state_subset=False):
"""
update the particle state from the context, create a particle substate and update from context
"""
#update the particle state and the particle state subset
particle_state.update_from_context(self.context, ignore_velocities=True) #update the particle state from the context
if new_state_subset:
self.particle_state_subset = SamplerState(positions = particle_state.positions[list(self._subset_indices_map.keys())]) #create a particle state from the subset context
else:
self.particle_state_subset.positions = particle_state.positions[list(self._subset_indices_map.keys())] #update the particle subset positions appropriately
self.particle_state_subset.apply_to_context(self.context_subset, ignore_velocities=True) #apply the subset particle state to its context
self.particle_state_subset.update_from_context(self.context_subset, ignore_velocities=True) #update the subset particle state from its context to updated the potential energy
def _update_current_state_works(self, particle_state):
"""
update the current state and associated works
"""
#get the reduced potential
reduced_potential = self._compute_hybrid_potential(_lambda = self._iteration / self._n_iterations, particle_state = particle_state)
perturbed_reduced_potential = self._compute_hybrid_potential(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state)
self._current_state_works.append(self._current_state_works[-1] + (perturbed_reduced_potential - reduced_potential))
def _update_force(self, particle_state):
"""
update the force
"""
mm_force_matrix = self._compute_hybrid_forces(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state).value_in_unit_system(unit.md_unit_system)
self.integrator.setPerDofVariableByName('modified_force', mm_force_matrix)
def _before_integration(self, *args, **kwargs):
particle_state = args[0] #define the particle state
n_iterations = args[1] #define the number of iterations
self._initialize_state_works()
self._initialize_iterations(n_iterations)
#update the particle state and the particle state subset
self._update_particle_state_substate(particle_state, new_state_subset=True)
self._update_current_state_works(particle_state)
self._update_force(particle_state)
#report
if self._write_trajectory: # the first state is always saved for processing purposes
self.particle.update_state(particle_state)
self.reporter.record([self.particle])
def _during_integration(self, *args, **kwargs):
particle_state = args[0]
self._iteration += 1.0
self._update_particle_state_substate(particle_state)
#get the reduced potential
if self._iteration < self._n_iterations:
self._update_current_state_works(particle_state)
self._update_force(particle_state)
else:
#we are done
pass
if self._write_trajectory and int(self._iteration) % self.write_trajectory_interval == 0:
self.particle.update_state(particle_state)
if self._iteration == self._n_iterations:
self.reporter.record([self.particle], save_to_disk=True)
else:
self.reporter.record([self.particle], save_to_disk=False)
def _after_integration(self, *args, **kwargs):
self._state_works[self._state_works_counter] = deepcopy(self._current_state_works)
self._state_works_counter += 1
if self._write_trajectory:
self.reporter.reset()
#self._log_context_parameters()
def _compute_hybrid_potential(self,_lambda, particle_state):
"""
function to compute the hybrid reduced potential defined as follows:
U(x_rec, x_lig) = u_mm,rec(x_rec) - lambda*u_mm,lig(x_lig) + lambda*u_ani,lig(x_lig)
"""
reduced_potential = (self.pdf_state.reduced_potential(particle_state)
- _lambda * self.pdf_state_subset.reduced_potential(self.particle_state_subset)
+ _lambda * self.ani_handler.calculate_energy(self.particle_state_subset.positions) * self.pdf_state.beta)
return reduced_potential
def _compute_hybrid_forces(self, _lambda, particle_state):
"""
function to compute a hybrid force matrix of shape num_particles x 3
in the spirit of the _compute_hybrid_potential, we compute the forces in the following way
F(x_rec, x_lig) = F_mm(x_rec, x_lig) - lambda * F_mm(x_lig) + lambda * F_ani(x_lig)
"""
# get the complex mm forces
state = self.context.getState(getForces=True)
mm_force_matrix = state.getForces(asNumpy=True) # returns forces in kJ/(nm mol)
# get the ligand mm forces
subset_state = self.context_subset.getState(getForces=True)
mm_force_matrix_subset = subset_state.getForces(asNumpy=True)
# get the ligand ani forces
coords = self.particle_state_subset.positions
subset_ani_force_matrix, energie = self.ani_handler.calculate_force(coords) # returns force in kJ/(A mol)
#print(f"ani force matrix head: ",subset_ani_force_matrix[0])
# now combine the ligand forces
subset_force_matrix = _lambda * (subset_ani_force_matrix - mm_force_matrix_subset) #we are adding two Quantities with different units, but they are compatible
#print(f"mm subset force matrix head", mm_force_matrix_subset[0])
# and append to the complex forces...
#print(f"mm force matrix head", mm_force_matrix[0])
mm_force_matrix[list(self._subset_indices_map.keys()), :] += subset_force_matrix #and same, here...
#print(f"mm force matrix head (after ani modification)", mm_force_matrix[0])
return mm_force_matrix
def _get_context_subset_parameters(self):
"""
return a dictionary of the self.context_subset's parameters
returns
context_parameters : dict
{parameter name <str> : parameter value value <float>}
"""
swig_parameters = self.context_subset.getParameters()
context_parameters = {q: swig_parameters[q] for q in swig_parameters}
return context_parameters
def _log_context_parameters(self):
"""
log the context and context subset parameters
"""
context_parameters = self._get_context_parameters()
context_subset_parameters = self._get_context_subset_parameters()
_logger.debug(f"\tcontext_parameters during integration:")
for key, val in context_parameters.items():
_logger.debug(f"\t\t{key}: {val}")
_logger.debug(f"\tcontext subset parameters during integration:")
for key, val in context_subset_parameters:
_logger.debug(f"\t\t{key}: {val}")
@property
def state_works(self):
return self._state_works