def test_make_count_models(lagtime):
dtrajs = [
np.asarray([1, 1, 2, 3, 1, 1, 1, 2, 0, 0, 1, 3, 1, 4, 2, 2, 2, 2])
]
ttrajs = [
np.asarray([0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2])
]
bias_matrices = make_matching_bias_matrix(dtrajs, 3)
dataset = TRAMDataset(dtrajs=dtrajs,
ttrajs=ttrajs,
bias_matrices=bias_matrices,
lagtime=lagtime)
np.testing.assert_equal(len(dataset.count_models), dataset.n_therm_states)
np.testing.assert_equal(dataset.state_counts.shape,
(dataset.n_therm_states, dataset.n_markov_states))
np.testing.assert_equal(dataset.transition_counts.shape,
(dataset.n_therm_states, dataset.n_markov_states,
dataset.n_markov_states))
np.testing.assert_equal(dataset.count_models[0].state_symbols,
[0, 1, 2, 3])
np.testing.assert_equal(dataset.count_models[1].state_symbols,
[0, 1, 2, 3, 4])
np.testing.assert_equal(dataset.count_models[2].state_symbols, [0, 1, 2])
for k in range(dataset.n_therm_states):
np.testing.assert_equal(
dataset.transition_counts[k].sum(),
len(dataset._find_trajectory_fragments()[k][0]) - lagtime)
np.testing.assert_equal(
dataset.state_counts[k].sum(),
len(dataset._find_trajectory_fragments()[k][0]))
def test_restrict_to_submodel_with_indices_input(test_input, submodel,
expected):
dtrajs = [np.asarray(i) for i in test_input]
bias_matrices = make_matching_bias_matrix(dtrajs)
tram_data = TRAMDataset(dtrajs=dtrajs, bias_matrices=bias_matrices)
tram_data.restrict_to_submodel(submodel)
np.testing.assert_equal(tram_data.dtrajs, expected)
def test_restrict_to_submodel_with_submodel_input(test_input, expected):
tram_input = [np.asarray(i) for i in test_input]
_, bias_matrices = make_random_input_data(2,
7,
n_samples=len(tram_input[0]),
make_ttrajs=False)
tram_data = TRAMDataset(dtrajs=test_input, bias_matrices=bias_matrices)
counts_model = TransitionCountEstimator(1, 'sliding').fit_fetch(tram_input)
cset = counts_model.submodel([1, 2, 3])
tram_data.restrict_to_submodel(cset)
np.testing.assert_equal(tram_data.dtrajs, expected)
def test_transposed_count_matrices_bug(dtrajs):
dtrajs = [np.asarray(traj) for traj in dtrajs]
bias_matrices = make_matching_bias_matrix(dtrajs)
dataset = TRAMDataset(dtrajs=dtrajs, bias_matrices=bias_matrices)
dataset.restrict_to_largest_connected_set(
connectivity='summed_count_matrix')
np.testing.assert_equal(dataset.state_counts,
[[10, 0], [9, 1], [4, 6], [3, 7], [1, 9]])
np.testing.assert_equal(
dataset.transition_counts,
[[[9, 0], [0, 0]], [[7, 1], [1, 0]], [[2, 2], [1, 4]],
[[1, 1], [2, 5]], [[0, 1], [1, 7]]])
def test_property_caching_transition_counts():
dtrajs, bias_matrices, ttrajs = make_random_input_data(2, 5)
# make sure at least one count will be deleted after restricting to submodel, and that the submodel still contains
# some more counts
dtrajs[0][:5] = [0, 1, 2, 3, 4]
dataset = TRAMDataset(dtrajs, bias_matrices, ttrajs)
transition_counts_1 = dataset.transition_counts
dataset.restrict_to_submodel([0, 1, 2, 3])
transition_counts_2 = dataset.transition_counts
with np.testing.assert_raises(AssertionError):
np.testing.assert_array_equal(transition_counts_1, transition_counts_2)
def get_connected_set_from_dtrajs_input(dtrajs,
connectivity,
has_ttrajs=True,
connectivity_factor=1):
dtrajs = [np.asarray(traj) for traj in dtrajs]
if has_ttrajs:
ttrajs = [np.asarray([i] * len(traj)) for i, traj in enumerate(dtrajs)]
else:
ttrajs = None
bias_matrices = [np.ones((len(traj), len(dtrajs))) for traj in dtrajs]
tramdata = TRAMDataset(dtrajs=dtrajs,
ttrajs=ttrajs,
bias_matrices=bias_matrices,
lagtime=1,
count_mode='sliding')
return tramdata._find_largest_connected_set(
connectivity=connectivity, connectivity_factor=connectivity_factor)
def test_dataset_counts_no_ttrajs():
dtrajs, bias_matrices = make_random_input_data(5, 8, make_ttrajs=False)
dataset = TRAMDataset(dtrajs=dtrajs, bias_matrices=bias_matrices)
np.testing.assert_equal(dtrajs, dataset.dtrajs)
np.testing.assert_equal(bias_matrices, dataset.bias_matrices)
np.testing.assert_equal(dataset.state_counts.sum(),
np.sum([len(traj) for traj in dtrajs]))
np.testing.assert_equal(
dataset.transition_counts.sum(),
np.sum([len(traj) - dataset.lagtime for traj in dtrajs]))
def test_get_trajectory_fragments(dtrajs, ttrajs):
dtrajs = [np.asarray(d) for d in dtrajs]
ttrajs = [np.asarray(t) for t in ttrajs]
bias_matrices = make_matching_bias_matrix(dtrajs)
dataset = TRAMDataset(dtrajs=dtrajs,
ttrajs=ttrajs,
bias_matrices=bias_matrices)
# dtraj should be split into fragments [[[1, 3], [5, 6], [8, 9]], [[10, 11], [12, 13]]] due to replica exchanges
# found in ttrajs. This should lead having only 5 transitions in transition counts:
np.testing.assert_equal(dataset.state_counts.sum(), 10)
np.testing.assert_equal(dataset.transition_counts.sum(), 5)
def make_random_dataset(n_therm_states,
n_markov_states,
n_samples=10,
make_ttrajs=True):
if make_ttrajs:
dtrajs, bias_matrices, ttrajs = make_random_input_data(
n_therm_states, n_markov_states, n_samples)
else:
ttrajs = None
dtrajs, bias_matrices = make_random_input_data(n_therm_states,
n_markov_states,
n_samples,
make_ttrajs=False)
return TRAMDataset(dtrajs=dtrajs,
bias_matrices=bias_matrices,
ttrajs=ttrajs)