def check_populate_initial(self, prob, target_counts):
istate = InitialState(0, 0, 0, pcoord=[0.0])
self.system.bin_target_counts = np.array(
[target_counts, target_counts])
self.we_driver.populate_initial([istate], [prob], system=self.system)
assert len(self.we_driver.next_iter_binning[0]) == target_counts
def test_recycle(self):
segments = [
self.segment(0.0, 1.5, weight=0.5),
self.segment(0.0, 0.5, weight=0.5)
]
tstate = TargetState('recycle', [1.5], 0)
istate = InitialState(0, 0, 0, pcoord=[0.0])
self.we_driver.new_iteration(initial_states=[istate],
target_states=[tstate])
n_needed = self.we_driver.assign(segments)
assert n_needed == 0
self.we_driver.construct_next()
n_recycled = len(list(self.we_driver.recycling_segments))
assert n_recycled == 1
assert len(self.we_driver.next_iter_binning[0]) == 4
assert len(self.we_driver.next_iter_binning[1]) == 0
assert abs(
sum(seg.weight for seg in self.we_driver.next_iter_binning[0]) -
1.0) < 4 * EPS
assert np.allclose(
[seg.weight for seg in self.we_driver.next_iter_binning[0]],
[0.25 for _i in range(4)])
assert segments[0].endpoint_type == Segment.SEG_ENDPOINT_RECYCLED
def parent(self):
"""Walker or InitialState: The parent of the walker."""
parent_id = self.iteration.h5group['seg_index']['parent_id'][
self.index]
if parent_id >= 0:
return Walker(parent_id, self.iteration.prev)
istate_id = -(parent_id + 1)
row = self.iteration.h5group['ibstates']['istate_index'][istate_id]
# Initial states may or may not be generated from a basis state.
bstate_id = row['basis_state_id']
try:
bstate = self.iteration.basis_state(bstate_id)
except IndexError:
bstate = None
bstate_id = None
return InitialState(
istate_id,
bstate_id,
row['iter_created'],
iter_used=row['iter_used'],
istate_type=row['istate_type'],
istate_status=row['istate_status'],
pcoord=self.iteration.h5group['ibstates']['istate_pcoord']
[istate_id],
basis_state=bstate,
)
def initial_states(self):
"""list[InitialState]: Initial states."""
return [
InitialState(
row.Index,
row.basis_state_id,
row.iter_created,
iter_used=row.iter_used,
istate_type=row.istate_type,
istate_status=row.istate_status,
pcoord=pcoord,
) for row, pcoord in zip(self.initial_state_info.itertuples(),
self.initial_state_pcoords)
]
def from_data_manager(cls, n_iter, seg_id, data_manager=None):
'''Construct and return a trajectory trace whose last segment is identified
by ``seg_id`` in the iteration number ``n_iter``.'''
data_manager = data_manager or westpa.rc.get_data_manager()
# These values are used later on
endpoint_type = None
pcoord_dtype = None
pcoord_pt_shape = None
seginfo = []
parent_id = seg_id
while n_iter > 0 and parent_id >= 0:
seg_id = parent_id
iter_group = data_manager.get_iter_group(n_iter)
pcoord_ds = iter_group['pcoord']
seg_index = iter_group['seg_index']
n_segs = pcoord_ds.shape[0]
pcoord_len = pcoord_ds.shape[1]
assert seg_id < n_segs
indexrow = seg_index[seg_id]
final_pcoord = pcoord_ds[seg_id, pcoord_len - 1]
weight = indexrow['weight']
cputime = indexrow['cputime']
walltime = indexrow['walltime']
try:
parent_id = int(indexrow['parent_id'])
except IndexError:
# old HDF5 version
parent_id = int(
iter_group['parents'][indexrow['parents_offset']])
if endpoint_type is None:
endpoint_type = indexrow['endpoint_type']
pcoord_pt_shape = pcoord_ds.shape[2:]
pcoord_dtype = pcoord_ds.dtype
seginfo.append(
(n_iter, seg_id, weight, walltime, cputime, final_pcoord))
del iter_group, pcoord_ds, seg_index
n_iter -= 1
# loop terminates with parent_id set to the identifier of the initial state,
# seg_id set to the identifier of the first segment in the trajectory, and
# n_iter set to one less than the iteration of the first segment
first_iter = n_iter + 1
first_seg_id = seg_id
first_parent_id = parent_id
# Initial segment (for fetching initial state)
first_segment = Segment(n_iter=first_iter,
seg_id=first_seg_id,
parent_id=first_parent_id)
seginfo.reverse()
summary_dtype = np.dtype([
('n_iter', n_iter_dtype),
('seg_id', seg_id_dtype),
('weight', weight_dtype),
('walltime', utime_dtype),
('cputime', utime_dtype),
('final_pcoord', pcoord_dtype, pcoord_pt_shape),
])
summary = np.array(seginfo, dtype=summary_dtype)
try:
initial_state = data_manager.get_segment_initial_states(
[first_segment], first_iter)[0]
except KeyError:
# old HDF5 version
assert parent_id < 0
istate_pcoord = data_manager.get_iter_group(first_iter)['pcoord'][
first_seg_id, 0]
istate_id = -(first_parent_id + 1)
basis_state = None
initial_state = InitialState(istate_id,
None,
iter_created=0,
pcoord=istate_pcoord)
else:
basis_state = data_manager.get_basis_states(first_iter)[
initial_state.basis_state_id]
return cls(summary, endpoint_type, basis_state, initial_state,
data_manager)