class TestFiniteDistribution(unittest.TestCase):
def setUp(self):
self.die = Choose({1, 2, 3, 4, 5, 6})
self.ragged = Categorical({0: 0.9, 1: 0.05, 2: 0.025, 3: 0.025})
def test_map(self):
plusOne = self.die.map(lambda x: x + 1)
assert_almost_equal(self, plusOne, Choose({2, 3, 4, 5, 6, 7}))
evenOdd = self.die.map(lambda x: x % 2 == 0)
assert_almost_equal(self, evenOdd, Choose({True, False}))
greaterThan4 = self.die.map(lambda x: x > 4)
assert_almost_equal(self, greaterThan4,
Categorical({
True: 1 / 3,
False: 2 / 3
}))
def test_expectation(self):
self.assertAlmostEqual(self.die.expectation(float), 3.5)
even = self.die.map(lambda n: n % 2 == 0)
self.assertAlmostEqual(even.expectation(float), 0.5)
self.assertAlmostEqual(self.ragged.expectation(float), 0.175)
class TestDistribution(unittest.TestCase):
def setUp(self):
self.finite = Choose(range(0, 6))
self.sampled = SampledDistribution(lambda: self.finite.sample(),
100000)
def test_expectation(self):
expected_finite = self.finite.expectation(lambda x: x)
expected_sampled = self.sampled.expectation(lambda x: x)
self.assertLess(abs(expected_finite - expected_sampled), 0.02)
def test_sample_n(self):
samples = self.sampled.sample_n(10)
self.assertEqual(len(samples), 10)
self.assertTrue(all(0 <= s < 6 for s in samples))
