def check_populate_initial(self, prob, target_counts):
istate = InitialState(0, 0, 0, pcoord=[0.0])
self.system.bin_target_counts = numpy.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 numpy.allclose(
[seg.weight for seg in self.we_driver.next_iter_binning[0]],
[0.25 for _i in xrange(4)])
assert segments[0].endpoint_type == Segment.SEG_ENDPOINT_RECYCLED
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 = long(indexrow['parent_id'])
except IndexError:
# old HDF5 version
parent_id = long(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 = numpy.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 = numpy.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)
if (weights < 0.5 * ideal_weight).any():
print(' need to merge')
we.exact_counts = True
we.input_weights = weights
we.input_history_ids = ids
we.run_we(8)
for (struct_id, weight) in we.output_walkers:
try:
istate = initial_states[struct_id]
except KeyError:
istate = InitialState(
state_id=istate_id,
basis_state_id=0,
iter_created=1,
iter_used=1,
istate_type=InitialState.ISTATE_TYPE_GENERATED,
istate_status=InitialState.ISTATE_STATUS_PREPARED,
pcoord=all_pcoord[struct_id])
initial_states[struct_id] = istate
istate_filename = 'istates/struct_{:06d}.gro'.format(istate_id)
uv.selectAtoms('all').positions = coord_ds[struct_id]
uv.selectAtoms('all').write(istate_filename)
print(
' wrote {} containing initial state {} from structure {} with weight {}'
.format(istate_filename, istate_id, struct_id, weight))
istate_id += 1
segment = Segment(seg_id=seg_id,
n_iter=1,
weight=weight,