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_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)