wepy.runners.openmm module

OpenMM molecular dynamics runner with accessory classes.

OpenMM is a library with support for running molecular dynamics simulations with specific support for fast GPU calculations. The component based architecture of OpenMM makes it a perfect fit with wepy.

In addition to the principle OpenMMRunner class there are a few classes here that make using OpenMM runner more efficient.

First is a WalkerState class (OpenMMState) that wraps the openmm state object directly, itself is a wrapper around the C++ datastructures. This gives better performance by not performing copies to a WalkerState dictionary.

Second, is the OpenMMWalker which is identical to the Walker class except that it enforces the state is an actual instantiation of OpenMMState. Use of this is optional.

Finally, is the OpenMMGPUWorker class. This is to be used as the worker type for the WorkerMapper work mapper. This is necessary to allow passing of the device index to OpenMM for which GPU device to use.

wepy.runners.openmm.KEYS = ('positions', 'velocities', 'forces', 'kinetic_energy', 'potential_energy', 'time', 'box_vectors', 'box_volume', 'parameters', 'parameter_derivatives')

Names of the fields of the OpenMMState.

wepy.runners.openmm.GET_STATE_KWARG_DEFAULTS = (('getPositions', True), ('getVelocities', True), ('getForces', True), ('getEnergy', True), ('getParameters', True), ('getParameterDerivatives', False), ('enforcePeriodicBox', True))

Mapping of key word arguments to the simulation.context.getState method for retrieving data for a simulation state. By default we set each as True to retrieve all information. The presence or absence of them is handled by the OpenMMState.

wepy.runners.openmm.resolve_state_data_type_enum_values()

wepy.runners.openmm.get_state_fields_present(sim_state)

For a state returns a set of the field data types present in it.

wepy.runners.openmm.UNIT_NAMES = (('positions_unit', 'nanometer'), ('time_unit', 'picosecond'), ('box_vectors_unit', 'nanometer'), ('velocities_unit', 'nanometer/picosecond'), ('forces_unit', 'kilojoule/(nanometer*mole)'), ('box_volume_unit', 'nanometer'), ('kinetic_energy_unit', 'kilojoule/mole'), ('potential_energy_unit', 'kilojoule/mole'))

Mapping of unit identifier strings to the serialized string spec of the unit.

class wepy.runners.openmm.OpenMMRunner(system, topology, integrator, platform=None, platform_kwargs=None, enforce_box=False, get_state_kwargs=None)

Bases: Runner

Runner for OpenMM simulations.

Constructor for OpenMMRunner.

Parameters:

• system (simtk.openmm.System object) – The system (forcefields) for the simulation.

• topology (simtk.openmm.app.Topology object) – The topology for you system.

• integrator (subclass simtk.openmm.Integrator object) – Integrator for propagating dynamics.

• platform (str) – The specification for the default computational platform to use. Platform can also be set when run_segment is called. If None uses OpenMM default platform, see OpenMM documentation for all value but typical ones are: Reference, CUDA, OpenCL. If value is None the automatic platform determining mechanism in OpenMM will be used.

• platform_kwargs (dict of str : bool, optional) – key-values to set for a platform with platform.setPropertyDefaultValue as the default for this runner.

• enforce_box (bool) –

  Calls ‘context.getState’ with ‘enforcePeriodicBox’ if True.

  (Default value = False)

• get_state_kwargs (dict of str : bool, optional) – key-values to set for getting the state from the OpenMM context. keys not included will use the values in GET_STATE_KWARG_DEFAULTS. Will override the enforce_box flag.

Warning

Regarding the enforce_box option.

When retrieving states from an OpenMM simulation Context, you have the option to enforce periodic boundary conditions in the resulting atomic positions in a topology aware way that doesn’t break bonds through boundaries. This is convenient for post-processing as this can be a complex task and is not readily exposed in the OpenMM API as a standalone function.

However, in some types of simulations the periodic box vectors are ignored (such as implicit solvent ones) despite there being no option to not have periodic boundaries in the context itself. Likely if you are running one of these kinds of simulations you will not pay attention to the box vectors at all and the random defaults that exist will be very wrong but this incorrectness will not show in a non-wepy simulation with openmm unless you are handling the context states yourself. Then when you run in wepy the default of True to enforce the boxes will be applied and confusingly wrong answers will result that are difficult to find root cause of.

pre_cycle(platform=None, platform_kwargs=None, **kwargs)

Perform pre-cycle behavior. run_segment will be called for each walker so this allows you to perform changes of state on a per-cycle basis.

Parameters:

kwargs (key-word arguments) – Key-value pairs to be interpreted by each runner implementation.

post_cycle(**kwargs)

Perform post-cycle behavior. run_segment will be called for each walker so this allows you to perform changes of state on a per-cycle basis.

Parameters:

kwargs (key-word arguments) – Key-value pairs to be interpreted by each runner implementation.

_resolve_platform(platform, platform_kwargs)

run_segment(walker, segment_length, getState_kwargs=None, platform=None, platform_kwargs=None, **kwargs)

Run dynamics for the walker.

Parameters:

• walker (object implementing the Walker interface) – The walker for which dynamics will be propagated.

• segment_length (int or float) – The numerical value that specifies how much dynamics are to be run.

• getState_kwargs (dict of str : bool, optional) – Specify the key-word arguments to pass to simulation.context.getState when getting simulation states. If None defaults object values.

• platform (str or None or Ellipsis) – The specification for the computational platform to use. If None will use the default for the runner and ignore platform_kwargs. If Ellipsis forces the use of the OpenMM default or environmentally defined platform. See OpenMM documentation for all value but typical ones are: Reference, CUDA, OpenCL. If value is None the automatic platform determining mechanism in OpenMM will be used.

• platform_kwargs (dict of str : bool, optional) – key-values to set for a platform with platform.setPropertyDefaultValue for this segment only.

Returns:

new_walker – Walker after dynamics was run, only the state should be modified.

Return type:

object implementing the Walker interface

generate_state(simulation, segment_length, starting_walker, getState_kwargs)

Method for generating a wepy compliant state from an OpenMM simulation object and data about the last segment of dynamics run.

Parameters:

• simulation (simtk.openmm.app.Simulation object) – A complete simulation object from which the state will be extracted.

• segment_length (int) – The number of integration steps run in a segment of simulation.

• starting_walker (wepy.walker.Walker subclass object) – The walker that was the beginning of this segment of simyulation.

• getState_kwargs (dict of str : bool) – Specify the key-word arguments to pass to simulation.context.getState when getting simulation states.

Returns:

• new_state (wepy.runners.openmm.OpenMMState object) – A new state from the simulation state.

• This method is meant to be called from within the

• run_segment method during a simulation. It can be customized

• in subclasses to allow for the addition of custom attributes

• for a state, in addition to the base ones implemented in the

• interface to the openmm simulation state in OpenMMState.

• The extra arguments to this function are data that would allow

• for the calculation of integral values over the duration of

• the segment, such as time elapsed and differences from the

• starting state.

class wepy.runners.openmm.OpenMMState(sim_state, **kwargs)

Bases: WalkerState

Walker state that wraps an simtk.openmm.State object.

The keys for which values in the state are available are given by the KEYS module constant (accessible through the class constant of the same name as well).

Additional fields can be added to these states through passing extra kwargs to the constructor. These will be automatically given a suffix of “_OTHER” to avoid name clashes.

Constructor for OpenMMState.

Parameters:

• state (simtk.openmm.State object) – The simulation state retrieved from the simulation constant.

• kwargs (optional) – Additional attributes to set for the state. Will add the

• keys ("_OTHER" suffix to the)

KEYS = ('positions', 'velocities', 'forces', 'kinetic_energy', 'potential_energy', 'time', 'box_vectors', 'box_volume', 'parameters', 'parameter_derivatives')

The provided attribute keys for the state.

OTHER_KEY_TEMPLATE = '{}_OTHER'

String formatting template for attributes not set in KEYS.

property sim_state

The underlying simtk.openmm.State object this is wrapping.

property positions

The positions of the state as a numpy array simtk.units.Quantity object.

property positions_unit

The units (as a simtk.units.Unit object) the positions are in.

positions_values()

The positions of the state as a numpy array in the positions_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property velocities

The velocities of the state as a numpy array simtk.units.Quantity object.

property velocities_unit

The units (as a simtk.units.Unit object) the velocities are in.

velocities_values()

The velocities of the state as a numpy array in the velocities_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property forces

The forces of the state as a numpy array simtk.units.Quantity object.

property forces_unit

The units (as a simtk.units.Unit object) the forces are in.

forces_values()

The forces of the state as a numpy array in the forces_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property box_vectors

The box vectors of the state as a numpy array simtk.units.Quantity object.

property box_vectors_unit

The units (as a simtk.units.Unit object) the box vectors are in.

box_vectors_values()

The box vectors of the state as a numpy array in the box_vectors_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property kinetic_energy

The kinetic energy of the state as a numpy array simtk.units.Quantity object.

property kinetic_energy_unit

The units (as a simtk.units.Unit object) the kinetic energy is in.

kinetic_energy_value()

The kinetic energy of the state as a numpy array in the kinetic_energy_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property potential_energy

The potential energy of the state as a numpy array simtk.units.Quantity object.

property potential_energy_unit

The units (as a simtk.units.Unit object) the potential energy is in.

potential_energy_value()

The potential energy of the state as a numpy array in the potential_energy_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property time

The time of the state as a numpy array simtk.units.Quantity object.

property time_unit

The units (as a simtk.units.Unit object) the time is in.

time_value()

The time of the state as a numpy array in the time_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property box_volume

The box volume of the state as a numpy array simtk.units.Quantity object.

property box_volume_unit

The units (as a simtk.units.Unit object) the box volume is in.

box_volume_value()

The box volume of the state as a numpy array in the box_volume_unit simtk.units.Unit. This is what is returned by the __getitem__ accessor.

property parameters

The parameters of the state as a dictionary mapping the names of the parameters to their values which are numpy array simtk.units.Quantity objects.

property parameters_unit

The units for each parameter as a dictionary mapping parameter names to their corresponding unit as a simtk.units.Unit object.

parameters_values()

The parameters of the state as a dictionary mapping the name of the parameter to a numpy array in the unit for the parameter of the same name in the parameters_unit corresponding simtk.units.Unit object. This is what is returned by the __getitem__ accessor using the compound key syntax with the prefix ‘parameters’, e.g. state[‘parameter/paramA’] for the parameter ‘paramA’.

property parameter_derivatives

The parameter derivatives of the state as a dictionary mapping the names of the parameters to their values which are numpy array simtk.units.Quantity objects.

property parameter_derivatives_unit

The units for each parameter derivative as a dictionary mapping parameter names to their corresponding unit as a simtk.units.Unit object.

parameter_derivatives_values()

The parameter derivatives of the state as a dictionary mapping the name of the parameter to a numpy array in the unit for the parameter of the same name in the parameters_unit corresponding simtk.units.Unit object. This is what is returned by the __getitem__ accessor using the compound key syntax with the prefix ‘parameter_derivatives’, e.g. state[‘parameter_derivatives/paramA’] for the parameter ‘paramA’.

_dict_attr_to_compound_key_dict(root_key, attr_dict)

Transform a dictionary of values within the compound key ‘root_key’ to a dictionary mapping compound keys to values.

For example give the root_key ‘parameters’ and the parameters dictionary {‘paramA’ : 1.234} returns {‘parameters/paramA’ : 1.234}.

Parameters:

• root_key (str) – The compound key prefix

• attr_dict (dict of str : value) – The dictionary with simple keys within the root key namespace.

Returns:

compound_key_dict – The dictionary with the compound keys.

Return type:

dict of str : value

_get_nested_attr_from_compound_key(compound_key, compound_feat_dict)

Get arbitrarily deeply nested compound keys from the full dictionary tree.

Parameters:

• compound_key (str) – Compound key separated by ‘/’ characters

• compound_feat_dict (dict) – Dictionary of arbitrary depth

Returns:

Value requested by the key.

Return type:

value

parameters_features()

Returns a dictionary of the parameters with their appropriate compound keys. This can be used for placing them in the same namespace as the rest of the attributes.

parameter_derivatives_features()

Returns a dictionary of the parameter derivatives with their appropriate compound keys. This can be used for placing them in the same namespace as the rest of the attributes.

omm_state_dict()

Return a dictionary with all of the default keys from the wrapped simtk.openmm.State object

dict()

Return all key-value pairs as a dictionary.

to_mdtraj(topology)

Returns an mdtraj.Trajectory object from this walker’s state.

Parameters:

topology (mdtraj.Topology object) – Topology for the state.

Returns:

state_traj

Return type:

mdtraj.Trajectory object

wepy.runners.openmm.gen_sim_state(positions, system, integrator, getState_kwargs=None)

Convenience function for generating an omm.State object.

Parameters:

• positions (arraylike of float) – The positions for the system you want to set

• system (openmm.app.System object)

• integrator (openmm.Integrator object)

Returns:

sim_state

Return type:

openmm.State object

wepy.runners.openmm.gen_walker_state(positions, system, integrator, getState_kwargs=None)

Convenience function for generating a wepy walker State object for an openmm simulation state.

Parameters:

• positions (arraylike of float) – The positions for the system you want to set

• system (openmm.app.System object)

• integrator (openmm.Integrator object)

Returns:

walker_state

Return type:

wepy.runners.openmm.OpenMMState object

class wepy.runners.openmm.OpenMMWalker(state, weight)

Bases: Walker

Walker for OpenMMRunner simulations.

This simply enforces the use of an OpenMMState object for the walker state attribute.

Constructor for Walker.

Parameters:

• state (object implementing the WalkerState interface)

• weight (float)

clone(number=1)

Clone this walker by making a copy with the same state and split the probability uniformly between clones.

The number is the increase in the number of walkers.

e.g. number=1 will return 2 walkers with the same state as this object but with probability split 50/50 between them

Parameters:

number (int) –

Number of extra clones to make

(Default value = 1)

Returns:

cloned_walkers

Return type:

list of objects implementing the Walker interface

merge(other_walkers)

Merge a set of other walkers into this one using the merge function.

Parameters:

other_walkers (list of objects implementing the Walker interface) – The walkers that will be merged together

Returns:

merged_walker

Return type:

object implementing the Walker interface

squash(merge_target)

Add the weight of this walker to another.

Parameters:

merge_target (object implementing the Walker interface) – The walker to add this one’s weight to.

Returns:

merged_walker

Return type:

object implementing the Walker interface

class wepy.runners.openmm.OpenMMCPUWorker(*args, **kwargs)

Bases: Worker

Worker for OpenMM GPU simulations (CUDA or OpenCL platforms).

This is intended to be used with the wepy.work_mapper.WorkerMapper work mapper class.

This class must be used in order to ensure OpenMM runs jobs on the appropriate GPU device.

Constructor for the Worker class.

Parameters:

• worker_idx (int) – The index of the worker. Should be unique.

• task_queue (multiprocessing.JoinableQueue) – The shared task queue the worker will watch for new tasks to complete.

• result_queue (multiprocessing.Queue) – The shared queue that completed task results will be placed on.

• interrupt_connection (multiprocessing.Connection) – One end of a pipe to listen for messages specific to this worker.

• mapper_attributes (None or dict) – A dictionary of the attributes of the mapper for reference in workers.

• kwargs – The worker specific attributes

NAME_TEMPLATE = 'OpenMMCPUWorker-{}'

The name template the worker processes are named to substituting in the process number.

DEFAULT_NUM_THREADS = 1

run_task(task)

Actually executes the task.

This default runner simply executes the task thunk.

This can be customized by subclasses in order to allow for injection of worker specific data.

Parameters:

task (Task object) – The partially evaluated task; function plus arguments

Returns:

Results of running the task.

Return type:

task_result

static _Popen(process_obj)

static _after_fork()

_bootstrap(parent_sentinel=None)

_check_closed()

_run_task(task)

Runs the given task and returns the results.

This manages handling exceptions and tracebacks from the actual run_task function which is intended to be specialized by different workers to inject worker specific arguments to tasks. Such as node and device identification.

Parameters:

task (Task object) – The partially evaluated task; function plus arguments

Returns:

Results of running the task.

Return type:

task_result

_run_worker()

_shutdown()

_sigterm_shutdown(signum, frame)

_start_method = None

property attributes

Dictionary of attributes of the worker.

property authkey

close()

Close the Process object.

This method releases resources held by the Process object. It is an error to call this method if the child process is still running.

property daemon

Return whether process is a daemon

property exitcode

Return exit code of process or None if it has yet to stop

property ident

Return identifier (PID) of process or None if it has yet to start

is_alive()

Return whether process is alive

join(timeout=None)

Wait until child process terminates

kill()

Terminate process; sends SIGKILL signal or uses TerminateProcess()

property mapper_attributes

Dictionary of attributes of the worker.

property name

property pid

Return identifier (PID) of process or None if it has yet to start

run()

Method to be run in sub-process; can be overridden in sub-class

property sentinel

Return a file descriptor (Unix) or handle (Windows) suitable for waiting for process termination.

start()

Start child process

terminate()

Terminate process; sends SIGTERM signal or uses TerminateProcess()

property worker_idx

Dictionary of attributes of the worker.

class wepy.runners.openmm.OpenMMGPUWorker(worker_idx, task_queue, result_queue, exception_queue, interrupt_connection, mapper_attributes=None, log_level='INFO', **kwargs)

Bases: Worker

Worker for OpenMM GPU simulations (CUDA or OpenCL platforms).

This is intended to be used with the wepy.work_mapper.WorkerMapper work mapper class.

This class must be used in order to ensure OpenMM runs jobs on the appropriate GPU device.

Constructor for the Worker class.

Parameters:

• worker_idx (int) – The index of the worker. Should be unique.

• task_queue (multiprocessing.JoinableQueue) – The shared task queue the worker will watch for new tasks to complete.

• result_queue (multiprocessing.Queue) – The shared queue that completed task results will be placed on.

• interrupt_connection (multiprocessing.Connection) – One end of a pipe to listen for messages specific to this worker.

• mapper_attributes (None or dict) – A dictionary of the attributes of the mapper for reference in workers.

• kwargs – The worker specific attributes

NAME_TEMPLATE = 'OpenMMGPUWorker-{}'

The name template the worker processes are named to substituting in the process number.

static _Popen(process_obj)

static _after_fork()

_bootstrap(parent_sentinel=None)

_check_closed()

_run_task(task)

Runs the given task and returns the results.

This manages handling exceptions and tracebacks from the actual run_task function which is intended to be specialized by different workers to inject worker specific arguments to tasks. Such as node and device identification.

Parameters:

task (Task object) – The partially evaluated task; function plus arguments

Returns:

Results of running the task.

Return type:

task_result

_run_worker()

_shutdown()

_sigterm_shutdown(signum, frame)

_start_method = None

property attributes

Dictionary of attributes of the worker.

property authkey

close()

Close the Process object.

This method releases resources held by the Process object. It is an error to call this method if the child process is still running.

property daemon

Return whether process is a daemon

property exitcode

Return exit code of process or None if it has yet to stop

property ident

Return identifier (PID) of process or None if it has yet to start

is_alive()

Return whether process is alive

join(timeout=None)

Wait until child process terminates

kill()

Terminate process; sends SIGKILL signal or uses TerminateProcess()

property mapper_attributes

Dictionary of attributes of the worker.

property name

property pid

Return identifier (PID) of process or None if it has yet to start

run()

Method to be run in sub-process; can be overridden in sub-class

property sentinel

Return a file descriptor (Unix) or handle (Windows) suitable for waiting for process termination.

start()

Start child process

terminate()

Terminate process; sends SIGTERM signal or uses TerminateProcess()

property worker_idx

Dictionary of attributes of the worker.

run_task(task)

Actually executes the task.

This default runner simply executes the task thunk.

This can be customized by subclasses in order to allow for injection of worker specific data.

Parameters:

task (Task object) – The partially evaluated task; function plus arguments

Returns:

Results of running the task.

Return type:

task_result

class wepy.runners.openmm.OpenMMCPUWalkerTaskProcess(walker_idx, mapper_attributes, func, task_args, task_kwargs, worker_queue, results_list, worker_segment_times, interrupt_connection, **kwargs)

Bases: WalkerTaskProcess

static _Popen(process_obj)

static _after_fork()

_bootstrap(parent_sentinel=None)

_check_closed()

_external_sigterm_shutdown(signum, frame)

_run_task(task)

_run_walker()

_shutdown()

The normal shutdown which can be ordered by the work mapper.

_start_method = None

property attributes

property authkey

close()

Close the Process object.

This method releases resources held by the Process object. It is an error to call this method if the child process is still running.

property daemon

Return whether process is a daemon

property exitcode

Return exit code of process or None if it has yet to stop

property ident

Return identifier (PID) of process or None if it has yet to start

is_alive()

Return whether process is alive

join(timeout=None)

Wait until child process terminates

kill()

Terminate process; sends SIGKILL signal or uses TerminateProcess()

property name

property pid

Return identifier (PID) of process or None if it has yet to start

run()

Method to be run in sub-process; can be overridden in sub-class

property sentinel

Return a file descriptor (Unix) or handle (Windows) suitable for waiting for process termination.

start()

Start child process

terminate()

Terminate process; sends SIGTERM signal or uses TerminateProcess()

NAME_TEMPLATE = 'OpenMM_CPU_Walker_Task-{}'

run_task(task)

class wepy.runners.openmm.OpenMMGPUWalkerTaskProcess(walker_idx, mapper_attributes, func, task_args, task_kwargs, worker_queue, results_list, worker_segment_times, interrupt_connection, **kwargs)

Bases: WalkerTaskProcess

static _Popen(process_obj)

static _after_fork()

_bootstrap(parent_sentinel=None)

_check_closed()

_external_sigterm_shutdown(signum, frame)

_run_task(task)

_run_walker()

_shutdown()

The normal shutdown which can be ordered by the work mapper.

_start_method = None

property attributes

property authkey

close()

Close the Process object.

This method releases resources held by the Process object. It is an error to call this method if the child process is still running.

property daemon

Return whether process is a daemon

property exitcode

Return exit code of process or None if it has yet to stop

property ident

Return identifier (PID) of process or None if it has yet to start

is_alive()

Return whether process is alive

join(timeout=None)

Wait until child process terminates

kill()

Terminate process; sends SIGKILL signal or uses TerminateProcess()

property name

property pid

Return identifier (PID) of process or None if it has yet to start

run()

Method to be run in sub-process; can be overridden in sub-class

property sentinel

Return a file descriptor (Unix) or handle (Windows) suitable for waiting for process termination.

start()

Start child process

terminate()

Terminate process; sends SIGTERM signal or uses TerminateProcess()

NAME_TEMPLATE = 'OpenMM_GPU_Walker_Task-{}'

run_task(task)