WESTPA plugin example

replicasperbin_plugin.py

from __future__ import division
import logging
import numpy as np
log = logging.getLogger(__name__)


class ReplicasPerBinDriver(object):
    def __init__(self, sim_manager, plugin_config):
        super(ReplicasPerBinDriver, self).__init__()

        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
        self.priority = plugin_config.get('priority', 0)

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

    def get_bin_indices(self, region, prev):
        '''Return bin indices corresponding to region and previous state'''
        x_range = np.linspace(-1.0, 30.0, 100)
        y_range = np.linspace(-1.0, 30.0, 100)

        if region == 0:
            z_range = np.linspace(0.0, 15.0, 100)
        elif region == 1:
            z_range = np.linspace(15.0, 40.0, 100)
        elif region == 2:
            z_range = np.linspace(40.0, 56.0, 100)

        X, Y, Z = np.meshgrid(x_range, y_range, z_range)
        tmp = np.column_stack(map(np.ravel, [X, Y, Z]))

        pcoords = np.zeros((tmp.shape[0], 5), dtype=np.float32)
        pcoords[:,:3] = tmp
        pcoords[:,3] = region
        pcoords[:,4] = prev

        assignments = self.system.bin_mapper.assign(pcoords)

        return np.unique(assignments)

    def pre_we(self):

        n_iter = self.sim_manager.n_iter

        if n_iter < 10:
            return
        elif n_iter >= 10 and n_iter < 20:
            rpb = 2
        elif n_iter >= 20:
            rpb = 1

        log.info('Updating the bins with undefined prev region to {} replicas per bin.'.format(rpb))

        ii = self.get_bin_indices(0, -1)
        self.system.bin_target_counts[ii] = rpb

        ii = self.get_bin_indices(1, -1)
        self.system.bin_target_counts[ii] = rpb

        ii = self.get_bin_indices(2, -1)
        self.system.bin_target_counts[ii] = rpb

west.cfg

# The master WEST configuration file for a simulation.
# vi: set filetype=yaml :
---
west: 
  system:
    driver: system.System
    module_path: $WEST_SIM_ROOT
  propagation:
    max_total_iterations: 700
    max_run_wallclock: 47:50:00
    propagator: executable
    gen_istates: false
  plugins:
    - plugin: replicasperbin_plugin.ReplicasPerBinDriver
      enabled: true
  data:
    west_data_file: west.h5
    datasets: # dataset storage options
      #- name: pcoord   # store progress coordinates with HDF5 scale/offset
      #  scaleoffset: 4 # with 4 decimal places of precision
      - name: na_coords
        dtype: float32
        scaleoffset: 3
      - name: cl_coords
        dtype: float32
        scaleoffset: 3

    data_refs: # how to convert segments and states to paths, etc
      segment:       $WEST_SIM_ROOT/traj_segs/{segment.n_iter:06d}/{segment.seg_id:06d}
      basis_state:   $WEST_SIM_ROOT/bstates/{basis_state.auxref}
      initial_state: $WEST_SIM_ROOT/istates/struct_{initial_state.state_id:06d}.gro

  executable:
    environ: # environment variables for all executables
      SEG_DEBUG: 1

    datasets: # input/output for datasets
      - name:  pcoord
        loader: system.pcoord_loader
        filename: $WEST_SIM_ROOT/traj_segs/{segment.n_iter:06d}/{segment.seg_id:06d}/seg.xtc
      - name:    coord
        enabled: true
        loader:  system.coord_loader
        filename: $WEST_SIM_ROOT/traj_segs/{segment.n_iter:06d}/{segment.seg_id:06d}/seg.xtc

    propagator:
      executable: $WEST_SIM_ROOT/runseg.sh
      #stdout:     $WEST_SIM_ROOT/seg_logs/{segment.n_iter:06d}-{segment.seg_id:06d}.log
      #stdout:     /dev/null
      #no entry: do not redirect; remains in per-node logs
      stderr:     stdout # append stderr to stdout
      stdin:      null # do not redirect standard input
      cwd:        null # do not change directory

    get_pcoord:
      executable: $WEST_SIM_ROOT/get_pcoord.sh
      #stdout: /dev/null
      stderr: stdout
    
    gen_istate:
      executable: $WEST_SIM_ROOT/gen_istate.sh
      stdout: /dev/null
      stderr: stdout

    post_iteration:
      enabled: true
      executable: $WEST_SIM_ROOT/post_iter.sh
      stdout: $WEST_SIM_ROOT/post_logs/post_{n_iter:06d}.log
      stderr: stdout

    pre_iteration:
      enabled: false
      executable: $WEST_SIM_ROOT/pre_iter.sh
      stderr: stdout