else:
self.equilibrium_state_counts_full = _np.zeros((self.nthermo, self.nstates_full), dtype=_np.float64)
pg = _ProgressReporter()
stage = "cset"
with pg.context(stage=stage):
self.csets, pcset = _cset.compute_csets_TRAM(
self.connectivity, state_counts_full, count_matrices_full,
equilibrium_state_counts=self.equilibrium_state_counts_full,
ttrajs=ttrajs+equilibrium_ttrajs, dtrajs=dtrajs_full+equilibrium_dtrajs_full, bias_trajs=self.btrajs+self.equilibrium_btrajs,
nn=self.nn, factor=self.connectivity_factor,
callback=_IterationProgressIndicatorCallBack(pg, "finding connected set", stage=stage))
self.active_set = pcset
// check for empty states
for k in range(self.nthermo):
if len(self.csets[k]) == 0:
_warnings.warn(
"Thermodynamic state %d" % k \
+ " contains no samples after reducing to the connected set.", EmptyState)
// deactivate samples not in the csets, states are *not* relabeled
self.state_counts, self.count_matrices, self.dtrajs, _ = _cset.restrict_to_csets(
self.csets,
state_counts=state_counts_full, count_matrices=count_matrices_full,
ttrajs=ttrajs, dtrajs=dtrajs_full)
if self.equilibrium is not None:
self.equilibrium_state_counts, _, self.equilibrium_dtrajs, _ = _cset.restrict_to_csets(
self.csets,
state_counts=self.equilibrium_state_counts_full, ttrajs=equilibrium_ttrajs, dtrajs=equilibrium_dtrajs_full)
else:
self.equilibrium_state_counts = _np.zeros((self.nthermo, self.nstates_full), dtype=_np.intc) // (remember: no relabeling)
self.equilibrium_dtrajs = []
// self-consistency tests
assert _np.all(self.state_counts >= _np.maximum(self.count_matrices.sum(axis=1), \
self.count_matrices.sum(axis=2)))
assert _np.all(_np.sum(
[_np.bincount(d[d>=0], minlength=self.nstates_full) for d in self.dtrajs],
axis=0) == self.state_counts.sum(axis=0))
assert _np.all(_np.sum(
[_np.bincount(t[d>=0], minlength=self.nthermo) for t, d in zip(ttrajs, self.dtrajs)],
axis=0) == self.state_counts.sum(axis=1))
if self.equilibrium is not None:
assert _np.all(_np.sum(
[_np.bincount(d[d >= 0], minlength=self.nstates_full) for d in self.equilibrium_dtrajs],
axis=0) == self.equilibrium_state_counts.sum(axis=0))
assert _np.all(_np.sum(
[_np.bincount(t[d >= 0], minlength=self.nthermo) for t, d in zip(equilibrium_ttrajs, self.equilibrium_dtrajs)],
axis=0) == self.equilibrium_state_counts.sum(axis=1))
// check for empty states
for k in range(self.state_counts.shape[0]):
if self.count_matrices[k, :, :].sum() == 0 and self.equilibrium_state_counts[k, :].sum()==0:
_warnings.warn(
"Thermodynamic state %d" % k \
+ " contains no transitions and no equilibrium data after reducing to the connected set.", EmptyState)
if self.init == "mbar" and self.biased_conf_energies is None:
if self.direct_space:
mbar = _mbar_direct
else:
mbar = _mbar
stage = "MBAR init."
with pg.context(stage=stage):
self.mbar_therm_energies, self.mbar_unbiased_conf_energies, \
self.mbar_biased_conf_energies, _ = mbar.estimate(
(state_counts_full.sum(axis=1)+self.equilibrium_state_counts_full.sum(axis=1)).astype(_np.intc),
self.btrajs+self.equilibrium_btrajs, dtrajs_full+equilibrium_dtrajs_full,
maxiter=self.init_maxiter, maxerr=self.init_maxerr,
callback=_ConvergenceProgressIndicatorCallBack(
pg, stage, self.init_maxiter, self.init_maxerr),
n_conf_states=self.nstates_full)
self.biased_conf_energies = self.mbar_biased_conf_energies.copy()
// run estimator
if self.direct_space:
