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)