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', True), ('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.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)

Bases: wepy.runners.runner.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)

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)

post_cycle(**kwargs)

_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: wepy.walker.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

• suffix to the keys ("_OTHER") –

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: wepy.walker.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: wepy.work_mapper.mapper.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)

_bootstrap()

_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: wepy.work_mapper.mapper.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)

_bootstrap()

_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

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

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.OpenMMCPUWalkerTaskProcess(walker_idx, mapper_attributes, func, task_args, task_kwargs, worker_queue, results_list, worker_segment_times, interrupt_connection, **kwargs)

Bases: wepy.work_mapper.task_mapper.WalkerTaskProcess

static _Popen(process_obj)

_bootstrap()

_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: wepy.work_mapper.task_mapper.WalkerTaskProcess

static _Popen(process_obj)

_bootstrap()

_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)