odld with restructured westpa

gistfile1.txt

rm -f west.h5
BSTATES="--bstate initial,1.0"
#TSTATES="--tstate drift,10.01 --tstate bound,1.3"
#TSTATES="--tstate bound,1.2"
w_init $BSTATES $TSTATES "$@"

init.sh

rm -f west.h5
BSTATES="--bstate initial,1.0"
#TSTATES="--tstate drift,10.01 --tstate bound,1.3"
#TSTATES="--tstate bound,1.2"
w_init $BSTATES $TSTATES "$@"

odld_system.py

import numpy as np
from numpy.random import normal as random_normal

from westpa.core.binning import RectilinearBinMapper
from westpa.core.propagators import WESTPropagator
from westpa.core.systems import WESTSystem

PI = np.pi

pcoord_len = 21
pcoord_dtype = np.float32


class ODLDPropagator(WESTPropagator):
    def __init__(self, rc=None):
        super(ODLDPropagator, self).__init__(rc)

        self.coord_len = pcoord_len
        self.coord_dtype = pcoord_dtype
        self.coord_ndim = 1

        self.initial_pcoord = np.array([8.0], dtype=self.coord_dtype)

        self.sigma = 0.001**(0.5)

        self.A = 2
        self.B = 10
        self.C = 0.5
        self.x0 = 1

        # Implement a reflecting boundary at this x value
        # (or None, for no reflection)
        self.reflect_at = 10.0

    def get_pcoord(self, state):
        '''Get the progress coordinate of the given basis or initial state.'''
        state.pcoord = self.initial_pcoord.copy()

    def gen_istate(self, basis_state, initial_state):
        initial_state.pcoord = self.initial_pcoord.copy()
        initial_state.istate_status = initial_state.ISTATE_STATUS_PREPARED
        return initial_state

    def propagate(self, segments):

        A, B, C, x0 = self.A, self.B, self.C, self.x0

        n_segs = len(segments)

        coords = np.empty((n_segs, self.coord_len, self.coord_ndim), dtype=self.coord_dtype)

        for iseg, segment in enumerate(segments):
            coords[iseg, 0] = segment.pcoord[0]

        twopi_by_A = 2 * PI / A
        half_B = B / 2
        sigma = self.sigma
        gradfactor = self.sigma * self.sigma / 2
        coord_len = self.coord_len
        reflect_at = self.reflect_at
        all_displacements = np.zeros((n_segs, self.coord_len, self.coord_ndim), dtype=self.coord_dtype)
        for istep in range(1,coord_len):
            x = coords[:,istep-1,0]

            xarg = twopi_by_A*(x - x0)

            eCx = np.exp(C*x)
            eCx_less_one = eCx - 1.0

            all_displacements[:,istep,0] = displacements = random_normal(scale=sigma, size=(n_segs,))
            grad = half_B / (eCx_less_one*eCx_less_one)*(twopi_by_A*eCx_less_one*np.sin(xarg)+C*eCx*np.cos(xarg))

            newx = x - gradfactor*grad + displacements
            if reflect_at is not None:
                # Anything that has moved beyond reflect_at must move back that much

                # boolean array of what to reflect
                to_reflect = newx > reflect_at

                # how far the things to reflect are beyond our boundary
                reflect_by = newx[to_reflect] - reflect_at

                # subtract twice how far they exceed the boundary by
                # puts them the same distance from the boundary, on the other side
                newx[to_reflect] -= 2*reflect_by
            coords[:,istep,0] = newx

        for iseg, segment in enumerate(segments):
            segment.pcoord[...] = coords[iseg,:]
            segment.data['displacement'] = all_displacements[iseg]
            segment.status = segment.SEG_STATUS_COMPLETE

        return segments


class ODLDSystem(WESTSystem):
    def initialize(self):
        self.pcoord_ndim = 1
        self.pcoord_dtype = pcoord_dtype
        self.pcoord_len = pcoord_len

        #self.bin_mapper = RectilinearBinMapper([[0,1.3] + list(np.arange(1.4, 10.1, 0.1)) + [float('inf')]])
        self.bin_mapper = RectilinearBinMapper([list(np.arange(0.0, 10.1, 0.1))])
        self.bin_target_counts = np.empty((self.bin_mapper.nbins,), np.int_)
        self.bin_target_counts[...] = 10

run.sh

#!/bin/bash

rm -f west.log
w_run "$@"  &> west.log

west.cfg

# The master WEST configuration file for a simulation.
# vi: set filetype=yaml :
---
west: 
  system:
    driver: odld_system.ODLDSystem
    module_path: $WEST_SIM_ROOT
  propagation:
    max_total_iterations: 100
    max_run_wallclock: 3:00:00
    propagator: odld_system.ODLDPropagator
    gen_istates: true
    block_size: 10000
  data:
    west_data_file: west.h5
    aux_compression_threshold: 16384 # data sets bigger than this are compressed
                                     # unless overridden by an entry in ``datasets`` below
    datasets: # dataset storage options
      - name: displacement            # name used to refer to this in segment.data/env vars
        #h5path: auxdata/displacement # HDF5 storage path, overrides default
        #store: true                  # store when writing segment data (defaults to true)
        #load:  true                  # load when reading segment data (defaults to false)
        store: false
        load: false
        dtype: float32                # numpy dtype
        compression: false            # whether to store compressed
        scaleoffset: 4                # whether to store with scale/offset filter
        chunks: null                  # custom chunking, or null for auto/no chunking
                                      #  - ignored if necessary for other options
      - name: pcoord                  # you can mess CAREFULLY with pcoord as well
        scaleoffset: 4
        
    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/{initial_state.iter_created}/{initial_state.state_id}.gro
  plugins:
    #- plugin: westext.wess.WESSDriver # must name Python object
    #  enabled: true # optional, implied by presence in plugins list
    #  do_reweighting: true
    #  window_size: 0.5
  analysis:
     # Settings for w_ipa, an interactive analysis program that can also automate analysis.
     directory: ANALYSIS                # specify the directory all analysis files should exist in.
     postanalysis: True                 # should the routines for w_reweight be run?
     kinetics:                          # general options for both kinetics routines.
       # Command line arguments with values should be specified as key: value (see below)
       # Command line arguments that are flags without values should be included as a list value
       # in the extra key (extra: [ 'disable-correl', 'disable-bootstrap' ])
       # These are global options for each scheme; individual schemes can have different values, 
       # set in their respective section.
       step_iter: 10
       evolution: cumulative
       extra: [ 'disable-correl' ]
     analysis_schemes:                  # Analysis schemes.  Required: name (TEST), states, and bins
       TEST:
         enabled: True
         states:
           - label: unbound
             coords: [[8.0]]
           - label: bound
             coords: [[3.99]]
         bins:
           - type: RectilinearBinMapper
             boundaries: [[0.0,4.0,8.00,100000]]