Table of Contents Writing Plugins Introduction My First Plugin Some Examples of WESTPA Plugins

The natural progress of a WESTPA simulation iteration flows through a series of functions, as defined in the run function of the sim_manager.py script. The sim manager handles the sequential execution of the following functions (there are more, but these are the core ones):

• prepare_iteration(): this steps gets everything ready
• pre_propagation()
• propagate(): this is the step in which dynamics are propagated
• post_propagation()
• pre_we()
• run_we(): this is the step in which splitting and merging (as well as recycling) occur
• post_we()
• finalize_iteration(): this step double checks that everything completed successfully

Pre/post functions should be self-explanatory in that they (may) occur before or after the main propagate() or run_we() functions.

You may develop a plugin to overwrite any of the above functions. Some of the functions are actually empty by default (such as the post/pre) and so a plugin targeting these steps would essentially be adding a new function either before or after propagation or we.

The first step in developing a plugin is determining where you need to insert your code and which is the best entry point for that code. For example, if I want to develop a plugin to modify all segment pcoords before splitting/merging, one could conceivably develop a plugin to insert code at either the post_propagation or pre_we stage. Both would likely be equivalent in this scenario, but double check before you begin development to see what information about each walker is present at each stage. For instance, if you want to edit the pcoord values before splitting and merging, you want to make sure they are available at the stage you wish to replace (they should be for this example, but just as a check).

Let's start by writing our first plugin, which simply inserts itself at the pre_we() stage and prints an encouraging message to the west.log file. To do this, create a directory called almost_there, which is where the plugin-related scripts will be placed. Eventually, this plugin directory will be placed in src/westpa/westext, which is where all WESTPA plugins are located. Enter the newly created directory and create an __init__.py that contains the following:

  from .voting_driver import VotingDriver
  __all__ = ['VotingDriver']

Now, create the main plugin script, which we will call almost_there_driver.py. In this script, import any needed Python modules (including westpa) and then define your plugin's main class.

      def __init__(self, sim_manager, plugin_config):
  
          if not sim_manager.work_manager.is_master:
              return
  
          self.sim_manager = sim_manager
          self.data_manager = sim_manager.data_manager
          self.system = sim_manager.system
  

This function is where user-defined parameters from the west.cfg are read in and translated to variables. Below is an example.

          # Parameters from config file
          # enable the plugin
          self.encourage = plugin_config.get('encourage', False)
          # which encouraging message to print
          self.message_id = plugin_config.get('message_id', 1)
          # the plugin priority
          self.priority = plugin_config.get('priority', 0)

          # Register callback
          if self.encourage:
              sim_manager.register_callback(sim_manager.pre_we, self.pre_we, self.priority)

This reads in three variables from the west.cfg file, encourage, message_id and priority. The encourage parameter is really just an on/off switch for the plugin. The second, message_id, we will have to define later on in our function and the third is the priority of the plugin. I'm not sure what that is, so set it to 0 for now.

Next, add a pre_we() function that will be used by the sim manager to run our custom code.

      def pre_we(self):
          messages = ["you got this!", "almost there!", "don't give up!"]
          print(messages[self.message_id-1])

This will print our message just before splitting and merging of trajectories and will print the message id that will be specified in our west.cfg file.

Finally, to actually use your first plugin, add the following to the west.cfg file to enable the plugin and specify plugin-related parameters.

    plugins:
      - plugin: westpa.westext.almost_there.almost_there_driver.AlmostThereDriver
        encourage: True
        message_id: 3

Now if you run your WESTPA simulation, you should see an encouraging message printed in your west.log file. You won't be able to tell, but it was printed right before splitting and merging occurred. And with that, you have successfully created your first plugin. Way to go!

If you are interested in developing a WESTPA plugin, please see the following examples:

• History-augmented Markov state model (haMSM) restarting package
• WE-based string method. Instructions and examples for using this plug-in are provided here.
• Adaptive Voronoi binning scheme
• WExplore method

Consider using the MolSSI Python-centric cookiecutter template.

Wishlist for documentation on this page:

• A recommended skeleton for designing a WESTPA plugin
• Getting config parameters from YAML file
• Available callback hooks
• Registering callbacks

Figure 1 | Key steps of running a weighted ensemble simulation within the WESTPA framework. Upon initialization of the WESTPA environment and trajectories, the simulation manager controls the execution of a weighted ensemble simulation, including dynamics propagation using a selected dynamics engine via the work manager. Trajectories and progress coordinates are consolidated by the segment manager before subsequent uses by the simulation manager and HDF5 data manager. The Plugin manager can be used to initiate callbacks at certain time points during the simulation.