def update_args_env_initial_state(self, template_args, environ,
initial_state):
new_template_args = {'initial_state': initial_state}
new_env = {
self.ENV_ISTATE_ID:
str(initial_state.state_id if initial_state.
state_id is not None else -1),
self.ENV_ISTATE_DATA_REF:
self.makepath(self.initial_state_ref_template, new_template_args),
}
if initial_state.basis_state is not None:
basis_state = initial_state.basis_state
elif initial_state.istate_type == InitialState.ISTATE_TYPE_START:
basis_state = BasisState(label=f"sstate_{initial_state.state_id}",
pcoord=initial_state.pcoord,
probability=0.0,
auxref="")
else:
basis_state = self.basis_states[initial_state.basis_state_id]
self.update_args_env_basis_state(new_template_args, new_env,
basis_state)
template_args.update(new_template_args)
environ.update(new_env)
return template_args, environ
def basis_states(self):
"""list[BasisState]: Basis states in use for the iteration."""
return [
BasisState(row.label,
row.probability,
pcoord=pcoord,
auxref=row.auxref,
state_id=row.Index) for row, pcoord in
zip(self.basis_state_info.itertuples(), self.basis_state_pcoords)
]
def setUp(self):
parser = argparse.ArgumentParser()
westpa.rc.add_args(parser)
here = os.path.dirname(__file__)
os.environ['WEST_SIM_ROOT'] = os.path.join(here, 'fixtures', 'odld')
config_file_name = os.path.join(here, 'fixtures', 'odld', 'west.cfg')
args = parser.parse_args(['-r={}'.format(config_file_name)])
westpa.rc.process_args(args)
self.sim_manager = westpa.rc.get_sim_manager()
self.test_dir = tempfile.mkdtemp()
self.hdf5 = os.path.join(self.test_dir, "west.h5")
self.basis_states = [BasisState(label="label", probability=1.0)]
self.segments = [self.segment(0.0, 1.5, weight=0.125) for _i in range(4)] + [
self.segment(1.5, 0.5, weight=0.125) for _i in range(4)
]
self.sim_manager.we_driver.new_iteration()
self.sim_manager.we_driver.assign(self.segments)
self.sim_manager.we_driver.construct_next()
self.sim_manager.segments = {segment.seg_id: segment for segment in self.segments}
self.sim_manager.incomplete_segments = self.sim_manager.segments
self.sim_manager.current_iter_istates = self.sim_manager.segments
self.sim_manager.completed_segments = self.sim_manager.segments
self.sim_manager.report_bin_statistics = MagicMock(return_value=True)
data = self.sim_manager.we_driver.rc.get_data_manager()
data.we_h5filename = self.hdf5
data.prepare_backing()
data.create_ibstate_group([])
data.create_initial_states(1)
data.save_target_states([])
data.update_segments = MagicMock(return_value=None)
n_iter = 0
it_name = data.iter_group_name(n_iter)
for group in ["seg_index", "parents", "ibstates", "pcoord"]:
data.we_h5file.create_group(it_name + "/" + group)
data.get_new_weight_data = MagicMock(return_value=None)
data.get_segments = MagicMock(return_value=self.segments)
self.sim_manager.we_driver.rc.get_data_manager = MagicMock(return_value=data)
self.sim_manager.n_iter = n_iter
def basis_state(self, index):
"""Return the basis state with the given index.
Parameters
----------
index : int
Basis state index (0-based).
Returns
-------
BasisState
The basis state indexed by `index`.
"""
row = self.h5group['ibstates']['bstate_index'][index]
pcoord = self.h5group['ibstates']['bstate_pcoord'][index]
return BasisState(
row['label'].decode(),
row['probability'],
pcoord=pcoord,
auxref=row['auxref'].decode(),
state_id=index,
)
def update_args_env_segment(self, template_args, environ, segment):
template_args['segment'] = segment
environ[self.ENV_CURRENT_SEG_INITPOINT] = Segment.initpoint_type_names[
segment.initpoint_type]
if segment.initpoint_type == Segment.SEG_INITPOINT_CONTINUES:
# Could use actual parent object here if the work manager cared to pass that much data
# to us (we'd need at least the subset of parents for all segments sent in the call to propagate)
# that may make a good west.cfg option for future crazy extensibility, but for now,
# just populate the bare minimum
parent = Segment(n_iter=segment.n_iter - 1,
seg_id=segment.parent_id)
parent_template_args = dict(template_args)
parent_template_args['segment'] = parent
environ[self.ENV_PARENT_SEG_ID] = str(
segment.parent_id if segment.parent_id is not None else -1)
environ[self.ENV_PARENT_DATA_REF] = self.makepath(
self.segment_ref_template, parent_template_args)
elif segment.initpoint_type == Segment.SEG_INITPOINT_NEWTRAJ:
# This segment is initiated from a basis state; WEST_PARENT_SEG_ID and WEST_PARENT_DATA_REF are
# set to the basis state ID and data ref
initial_state = self.initial_states[segment.initial_state_id]
if initial_state.istate_type == InitialState.ISTATE_TYPE_START:
basis_state = BasisState(
label=f"sstate_{initial_state.state_id}",
pcoord=initial_state.pcoord,
probability=0.0,
auxref="")
else:
basis_state = self.basis_states[initial_state.basis_state_id]
if self.ENV_BSTATE_ID not in environ:
self.update_args_env_basis_state(template_args, environ,
basis_state)
if self.ENV_ISTATE_ID not in environ:
self.update_args_env_initial_state(template_args, environ,
initial_state)
assert initial_state.istate_type in (
InitialState.ISTATE_TYPE_BASIS,
InitialState.ISTATE_TYPE_GENERATED,
InitialState.ISTATE_TYPE_START,
)
if initial_state.istate_type == InitialState.ISTATE_TYPE_BASIS:
environ[self.ENV_PARENT_DATA_REF] = environ[
self.ENV_BSTATE_DATA_REF]
elif initial_state.istate_type == InitialState.ISTATE_TYPE_START:
# This points to the start-state PDB
environ[self.ENV_PARENT_DATA_REF] = environ[
self.
ENV_BSTATE_DATA_REF] + '/' + initial_state.basis_auxref
else: # initial_state.type == InitialState.ISTATE_TYPE_GENERATED
environ[self.ENV_PARENT_DATA_REF] = environ[
self.ENV_ISTATE_DATA_REF]
environ[self.ENV_CURRENT_SEG_ID] = str(
segment.seg_id if segment.seg_id is not None else -1)
environ[self.ENV_CURRENT_SEG_DATA_REF] = self.makepath(
self.segment_ref_template, template_args)
return template_args, environ
def entry_point():
parser = argparse.ArgumentParser(
'w_states',
description='''\
Display or manipulate basis (initial) or target (recycling) states for a WEST simulation. By default, states are
displayed (or dumped to files). If ``--replace`` is specified, all basis/target states are replaced for the
next iteration. If ``--append`` is specified, the given target state(s) are appended to the list for the
next iteration.
Appending basis states is not permitted, as this would require renormalizing basis state
probabilities in ways that may be error-prone. Instead, use ``w_states --show --bstate-file=bstates.txt``
and then edit the resulting ``bstates.txt`` file to include the new desired basis states, then use
``w_states --replace --bstate-file=bstates.txt`` to update the WEST HDF5 file appropriately.
''',
)
westpa.rc.add_args(parser)
smgroup = parser.add_argument_group('modes of operation')
mode_group = smgroup.add_mutually_exclusive_group()
mode_group.add_argument(
'--show', dest='mode', action='store_const', const='show', help='Display current basis/target states (or dump to files).'
)
mode_group.add_argument(
'--append',
dest='mode',
action='store_const',
const='append',
help='Append the given basis/target states to those currently in use.',
)
mode_group.add_argument(
'--replace',
dest='mode',
action='store_const',
const='replace',
help='Replace current basis/target states with those specified.',
)
parser.add_argument(
'--bstate-file',
metavar='BSTATE_FILE',
help='''Read (--append/--replace) or write (--show) basis state names, probabilities,
and data references from/to BSTATE_FILE.''',
)
parser.add_argument(
'--bstate',
action='append',
dest='bstates',
help='''Add the given basis state (specified as a string 'label,probability[,auxref]')
to the list of basis states (after those specified in --bstate-file, if any). This argument
may be specified more than once, in which case the given states are appended in the order
they are given on the command line.''',
)
parser.add_argument(
'--tstate-file',
metavar='TSTATE_FILE',
help='''Read (--append/--replace) or write (--show) target state names
and representative progress coordinates from/to TSTATE_FILE''',
)
parser.add_argument(
'--tstate',
action='append',
dest='tstates',
help='''Add the given target state (specified as a string 'label,pcoord0[,pcoord1[,...]]') to the
list of target states (after those specified in the file given by --tstates-from, if any).
This argument may be specified more than once, in which case the given states are appended
in the order they appear on the command line.''',
)
parser.set_defaults(mode='show')
work_managers.environment.add_wm_args(parser)
args = parser.parse_args()
westpa.rc.process_args(args)
work_managers.environment.process_wm_args(args)
work_manager = make_work_manager()
system = westpa.rc.get_system_driver()
with work_manager:
if work_manager.is_master:
data_manager = westpa.rc.get_data_manager()
data_manager.open_backing(mode='a')
sim_manager = westpa.rc.get_sim_manager()
n_iter = data_manager.current_iteration
assert args.mode in ('show', 'replace', 'append')
if args.mode == 'show':
basis_states = data_manager.get_basis_states(n_iter)
if basis_states:
bstate_file = sys.stdout if not args.bstate_file else open(args.bstate_file, 'wt')
bstate_file.write('# Basis states for iteration {:d}\n'.format(n_iter))
BasisState.states_to_file(basis_states, bstate_file)
target_states = data_manager.get_target_states(n_iter)
if target_states:
tstate_file = sys.stdout if not args.tstate_file else open(args.tstate_file, 'wt')
tstate_file.write('# Target states for iteration {:d}\n'.format(n_iter))
TargetState.states_to_file(target_states, tstate_file)
elif args.mode == 'replace':
seg_index = data_manager.get_seg_index(n_iter)
if (seg_index['status'] == Segment.SEG_STATUS_COMPLETE).any():
print('Iteration {:d} has completed segments; applying new states to iteration {:d}'.format(n_iter, n_iter + 1))
n_iter += 1
basis_states = []
if args.bstate_file:
basis_states.extend(BasisState.states_from_file(args.bstate_file))
if args.bstates:
for bstate_str in args.bstates:
fields = bstate_str.split(',')
label = fields[0]
probability = float(fields[1])
try:
auxref = fields[2]
except IndexError:
auxref = None
basis_states.append(BasisState(label=label, probability=probability, auxref=auxref))
if basis_states:
# Check that the total probability of basis states adds to one
tprob = sum(bstate.probability for bstate in basis_states)
if abs(1.0 - tprob) > len(basis_states) * EPS:
pscale = 1 / tprob
log.warning('Basis state probabilities do not add to unity; rescaling by {:g}'.format(pscale))
for bstate in basis_states:
bstate.probability *= pscale
# Assign progress coordinates to basis states
sim_manager.get_bstate_pcoords(basis_states, n_iter)
data_manager.create_ibstate_group(basis_states, n_iter)
sim_manager.report_basis_states(basis_states)
# Now handle target states
target_states = []
if args.tstate_file:
target_states.extend(TargetState.states_from_file(args.tstate_file, system.pcoord_dtype))
if args.tstates:
tstates_strio = io.StringIO('\n'.join(args.tstates).replace(',', ' '))
target_states.extend(TargetState.states_from_file(tstates_strio, system.pcoord_dtype))
del tstates_strio
if not target_states:
westpa.rc.pstatus('No target states specified.')
else:
data_manager.save_target_states(target_states, n_iter)
sim_manager.report_target_states(target_states)
data_manager.update_iter_group_links(n_iter)
else: # args.mode == 'append'
if args.bstate_file or args.bstates:
sys.stderr.write('refusing to append basis states; use --show followed by --replace instead\n')
sys.exit(2)
target_states = data_manager.get_target_states(n_iter)
seg_index = data_manager.get_seg_index(n_iter)
if (seg_index['status'] == Segment.SEG_STATUS_COMPLETE).any():
print('Iteration {:d} has completed segments; applying new states to iteration {:d}'.format(n_iter, n_iter + 1))
n_iter += 1
if args.tstate_file:
target_states.extend(TargetState.states_from_file(args.tstate_file, system.pcoord_dtype))
if args.tstates:
tstates_strio = io.StringIO('\n'.join(args.tstates).replace(',', ' '))
target_states.extend(TargetState.states_from_file(tstates_strio, system.pcoord_dtype))
del tstates_strio
if not target_states:
westpa.rc.pstatus('No target states specified.')
else:
data_manager.save_target_states(target_states, n_iter)
sim_manager.report_target_states(target_states)
data_manager.update_iter_group_links(n_iter)
else:
work_manager.run()
def entry_point():
parser = argparse.ArgumentParser('w_init',
description='''\
Initialize a new WEST simulation, creating the WEST HDF5 file and preparing the first
iteration's segments.
Initial states are generated from one or more "basis states" which are specified either in a file specified with --bstates-from, or
by one or more "--bstate" arguments. If neither --bstates-from nor at least one "--bstate" argument is provided, then a default
basis state of probability one identified by the state ID zero and label "basis" will be created (a warning will be printed in this
case, to remind you of this behavior, in case it is not what you wanted).
Target states for (non-equilibrium) steady-state simulations are specified either in a file specified with --tstates-from, or
by one or more --tstate arguments. If neither --tstates-from nor at least one --tstate argument is provided, then an equilibrium
simulation (without any sinks) will be performed.
''')
westpa.rc.add_args(parser)
parser.add_argument('--force',
dest='force',
action='store_true',
help='Overwrite any existing simulation data')
parser.add_argument(
'--bstate-file',
'--bstates-from',
metavar='BSTATE_FILE',
help=
'Read basis state names, probabilities, and (optionally) data references from BSTATE_FILE.'
)
parser.add_argument(
'--bstate',
action='append',
dest='bstates',
help=
'''Add the given basis state (specified as a string 'label,probability[,auxref]')
to the list of basis states (after those specified in --bstates-from, if any). This argument
may be specified more than once, in which case the given states are appended in the order
they are given on the command line.''')
parser.add_argument(
'--tstate-file',
'--tstates-from',
metavar='TSTATE_FILE',
help=
'Read target state names and representative progress coordinates from TSTATE_FILE'
)
parser.add_argument(
'--tstate',
action='append',
dest='tstates',
help=
'''Add the given target state (specified as a string 'label,pcoord0[,pcoord1[,...]]') to the
list of target states (after those specified in the file given by --tstates-from, if any).
This argument may be specified more than once, in which case the given states are appended
in the order they appear on the command line.''')
parser.add_argument(
'--segs-per-state',
type=int,
metavar='N',
default=1,
help=
'''Initialize N segments from each basis state (default: %(default)s).'''
)
parser.add_argument(
'--no-we',
'--shotgun',
dest='shotgun',
action='store_true',
help=
'''Do not run the weighted ensemble bin/split/merge algorithm on newly-created segments.'''
)
work_managers.environment.add_wm_args(parser)
args = parser.parse_args()
westpa.rc.process_args(args)
work_managers.environment.process_wm_args(args)
westpa.rc.work_manager = work_manager = make_work_manager()
system = westpa.rc.get_system_driver()
sim_manager = westpa.rc.get_sim_manager()
propagator = westpa.rc.get_propagator()
data_manager = westpa.rc.get_data_manager()
h5file = data_manager.we_h5filename
data_manager.system = system
we_driver = westpa.rc.get_we_driver()
with work_manager:
if work_manager.is_master:
# Process target states
target_states = []
if args.tstate_file:
target_states.extend(
TargetState.states_from_file(args.tstate_file,
system.pcoord_dtype))
if args.tstates:
tstates_strio = io.StringIO('\n'.join(args.tstates).replace(
',', ' '))
target_states.extend(
TargetState.states_from_file(tstates_strio,
system.pcoord_dtype))
del tstates_strio
# Process basis states
basis_states = []
if args.bstate_file:
basis_states.extend(
BasisState.states_from_file(args.bstate_file))
if args.bstates:
for bstate_str in args.bstates:
fields = bstate_str.split(',')
label = fields[0]
probability = float(fields[1])
try:
auxref = fields[2]
except IndexError:
auxref = None
basis_states.append(
BasisState(label=label,
probability=probability,
auxref=auxref))
if not basis_states:
log.error('At least one basis state is required')
sys.exit(3)
# Check that the total probability of basis states adds to one
tprob = sum(bstate.probability for bstate in basis_states)
if abs(1.0 - tprob) > len(basis_states) * EPS:
pscale = 1 / tprob
log.warning(
'Basis state probabilities do not add to unity; rescaling by {:g}'
.format(pscale))
for bstate in basis_states:
bstate.probability *= pscale
# Prepare simulation
sim_manager.initialize_simulation(
basis_states,
target_states,
segs_per_state=args.segs_per_state,
suppress_we=args.shotgun)
else:
work_manager.run()
def initialize(mode, bstates, _bstate_file, tstates, _tstate_file):
work_manager = make_work_manager()
system = westpa.rc.get_system_driver()
with work_manager:
if work_manager.is_master:
data_manager = westpa.rc.get_data_manager()
data_manager.open_backing(mode='a')
sim_manager = westpa.rc.get_sim_manager()
n_iter = data_manager.current_iteration
assert mode in ('show', 'replace', 'append')
if mode == 'show':
basis_states = data_manager.get_basis_states(n_iter)
if basis_states:
bstate_file = sys.stdout if not _bstate_file else open(
_bstate_file, 'wt')
bstate_file.write(
'# Basis states for iteration {:d}\n'.format(n_iter))
BasisState.states_to_file(basis_states, bstate_file)
target_states = data_manager.get_target_states(n_iter)
if target_states:
tstate_file = sys.stdout if not _tstate_file else open(
_tstate_file, 'wt')
tstate_file.write(
'# Target states for iteration {:d}\n'.format(n_iter))
TargetState.states_to_file(target_states, tstate_file)
elif mode == 'replace':
seg_index = data_manager.get_seg_index(n_iter)
if (seg_index['status'] == Segment.SEG_STATUS_COMPLETE).any():
print(
'Iteration {:d} has completed segments; applying new states to iteration {:d}'
.format(n_iter, n_iter + 1))
n_iter += 1
basis_states = []
if _bstate_file:
basis_states.extend(
BasisState.states_from_file(_bstate_file))
if bstates:
for bstate_str in bstates:
fields = bstate_str.split(',')
label = fields[0]
probability = float(fields[1])
try:
auxref = fields[2]
except IndexError:
auxref = None
basis_states.append(
BasisState(label=label,
probability=probability,
auxref=auxref))
if basis_states:
# Check that the total probability of basis states adds to one
tprob = sum(bstate.probability for bstate in basis_states)
if abs(1.0 - tprob) > len(basis_states) * EPS:
pscale = 1 / tprob
log.warning(
'Basis state probabilities do not add to unity; rescaling by {:g}'
.format(pscale))
for bstate in basis_states:
bstate.probability *= pscale
# Assign progress coordinates to basis states
sim_manager.get_bstate_pcoords(basis_states, n_iter)
data_manager.create_ibstate_group(basis_states, n_iter)
sim_manager.report_basis_states(basis_states)
# Now handle target states
target_states = []
if _tstate_file:
target_states.extend(
TargetState.states_from_file(_tstate_file,
system.pcoord_dtype))
if tstates:
tstates_strio = io.StringIO('\n'.join(tstates).replace(
',', ' '))
target_states.extend(
TargetState.states_from_file(tstates_strio,
system.pcoord_dtype))
del tstates_strio
if not target_states:
westpa.rc.pstatus('No target states specified.')
else:
data_manager.save_target_states(target_states, n_iter)
sim_manager.report_target_states(target_states)
data_manager.update_iter_group_links(n_iter)
else: # args.mode == 'append'
if _bstate_file or bstates:
sys.stderr.write(
'refusing to append basis states; use --show followed by --replace instead\n'
)
sys.exit(2)
target_states = data_manager.get_target_states(n_iter)
seg_index = data_manager.get_seg_index(n_iter)
if (seg_index['status'] == Segment.SEG_STATUS_COMPLETE).any():
print(
'Iteration {:d} has completed segments; applying new states to iteration {:d}'
.format(n_iter, n_iter + 1))
n_iter += 1
if _tstate_file:
target_states.extend(
TargetState.states_from_file(_tstate_file,
system.pcoord_dtype))
if tstates:
tstates_strio = io.StringIO('\n'.join(tstates).replace(
',', ' '))
target_states.extend(
TargetState.states_from_file(tstates_strio,
system.pcoord_dtype))
del tstates_strio
if not target_states:
westpa.rc.pstatus('No target states specified.')
else:
data_manager.save_target_states(target_states, n_iter)
sim_manager.report_target_states(target_states)
data_manager.update_iter_group_links(n_iter)
else:
work_manager.run()
def initialize(tstates, tstate_file, bstates, bstate_file, sstates=None, sstate_file=None, segs_per_state=1, shotgun=False):
"""
Initialize a WESTPA simulation.
tstates : list of str
tstate_file : str
bstates : list of str
bstate_file : str
sstates : list of str
sstate_file : str
segs_per_state : int
shotgun : bool
"""
westpa.rc.work_manager = work_manager = make_work_manager()
system = westpa.rc.get_system_driver()
sim_manager = westpa.rc.get_sim_manager()
data_manager = westpa.rc.get_data_manager()
data_manager.system = system
with work_manager:
if work_manager.is_master:
# Process target states
target_states = []
if tstate_file:
target_states.extend(TargetState.states_from_file(tstate_file, system.pcoord_dtype))
if tstates:
tstates_strio = io.StringIO('\n'.join(tstates).replace(',', ' '))
target_states.extend(TargetState.states_from_file(tstates_strio, system.pcoord_dtype))
del tstates_strio
# Process basis states
basis_states = []
if bstate_file:
basis_states.extend(BasisState.states_from_file(bstate_file))
if bstates:
for bstate_str in bstates:
fields = bstate_str.split(',')
label = fields[0]
probability = float(fields[1])
try:
auxref = fields[2]
except IndexError:
auxref = None
basis_states.append(BasisState(label=label, probability=probability, auxref=auxref))
# Process the list of start states, creating a BasisState from each
start_states = []
if sstate_file:
start_states.extend(BasisState.states_from_file(sstate_file))
if sstates:
for sstate_str in sstates:
fields = sstate_str.split(',')
label = fields[0]
probability = float(fields[1])
try:
auxref = fields[2]
except IndexError:
auxref = None
start_states.append(BasisState(label=label, probability=probability, auxref=auxref))
if not basis_states:
log.error('At least one basis state is required')
sys.exit(3)
# Check that the total probability of basis states adds to one
bstate_prob, sstate_prob = (
sum(bstate.probability for bstate in basis_states),
sum(sstate.probability for sstate in start_states),
)
# tprob = sum(bstate.probability for bstate in basis_states)
tprob = bstate_prob + sstate_prob
if abs(1.0 - tprob) > len(basis_states) * EPS:
pscale = 1 / tprob
log.warning(
'Basis state probabilities do not add to unity (basis: {:.2f}, start states: {:.2f}); rescaling by {:g}. If using start states, some rescaling is normal.'.format(
bstate_prob, sstate_prob, pscale
)
)
for bstate in basis_states:
bstate.probability *= pscale
for sstate in start_states:
sstate.probability *= pscale
# Prepare simulation
sim_manager.initialize_simulation(
basis_states=basis_states,
target_states=target_states,
start_states=start_states,
segs_per_state=segs_per_state,
suppress_we=shotgun,
)
else:
work_manager.run()