def test_sample(self): model = MC(['a', 'b'], [2, 2]) model.transition_models['a'] = { 0: { 0: 0.1, 1: 0.9 }, 1: { 0: 0.2, 1: 0.8 } } model.transition_models['b'] = { 0: { 0: 0.3, 1: 0.7 }, 1: { 0: 0.4, 1: 0.6 } } sample = model.sample(start_state=[State('a', 0), State('b', 1)], size=2) self.assertEqual(len(sample), 2) self.assertEqual(list(sample.columns), ['a', 'b']) self.assertTrue( list(sample.loc[0]) in [[0, 0], [0, 1], [1, 0], [1, 1]]) self.assertTrue( list(sample.loc[1]) in [[0, 0], [0, 1], [1, 0], [1, 1]])
def test_sample_less_arg(self, random_state): model = MC(['a', 'b'], [2, 2]) random_state.return_value = [State('a', 0), State('b', 1)] sample = model.sample(size=1) random_state.assert_called_once_with(model) self.assertEqual(model.state, random_state.return_value) self.assertEqual(len(sample), 1) self.assertEqual(list(sample.columns), ['a', 'b']) self.assertEqual(list(sample.loc[0]), [0, 1])
def test_sample(self): model = MC(['a', 'b'], [2, 2]) model.transition_models['a'] = {0: {0: 0.1, 1: 0.9}, 1: {0: 0.2, 1: 0.8}} model.transition_models['b'] = {0: {0: 0.3, 1: 0.7}, 1: {0: 0.4, 1: 0.6}} sample = model.sample(start_state=[State('a', 0), State('b', 1)], size=2) self.assertEqual(len(sample), 2) self.assertEqual(list(sample.columns), ['a', 'b']) self.assertTrue(list(sample.loc[0]) in [[0, 0], [0, 1], [1, 0], [1, 1]]) self.assertTrue(list(sample.loc[1]) in [[0, 0], [0, 1], [1, 0], [1, 1]])