Last active
June 12, 2020 17:13
-
-
Save synapticarbors/7908b78afd807db0204f90fc94c1cc6d to your computer and use it in GitHub Desktop.
WESTPA plugin example
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment