def test_can_pad_with_zero(self): # probability sums up to 0.7 p3 = ProbabilityDistribution({ 1: 0.5, 2: 0.2 }, normalize=False) # sanity check self.assertAlmostEqual(p3.p(0), 0) self.assertEqual(p3.mean, None) p3.pad() self.assertAlmostEqual(p3.p(0), 0.3) self.assertEqual(p3.mean, 0.9)
def test_can_initiate_distribution_with_list_of_values(self): prob_dist = [ (4, 0.5), (5, 0.2), ((5, 6), 0.3) ] try: p = ProbabilityDistribution(prob_dist) except AttributeError as e: self.fail(f'Should not have raised error "{e}"') self.assertAlmostEqual(p.p(4), 0.5) self.assertAlmostEqual(p.p((5, 6)), 0.5) self.assertAlmostEqual(p.p((5.5, 6)), 0.15)
def test_can_normalize_distribution(self): p3 = ProbabilityDistribution({1: 1, (2, 3): 1}, normalize=False) self.assertFalse(p3.is_normalized) p3.normalize() self.assertTrue(p3.is_normalized) self.assertAlmostEqual(p3.p(1), 0.5) self.assertAlmostEqual(p3.mean, 1.75)
def test_can_initiate_distribution_with_duplicate_entries(self): prob_dist = [ (4, 0.5), (4, 0.2), (5, 0.3) ] p = ProbabilityDistribution(prob_dist) self.assertAlmostEqual(p.p(4), 0.7)
def test_can_initiate_distribution_with_intersecting_ranges(self): prob_dist = [ ((4, 7), 0.3), ((5, 8), 0.3), ((5, 6), 0.4), ] p = ProbabilityDistribution(prob_dist) self.assertAlmostEqual(p.p((5, 6)), 0.6) self.assertAlmostEqual(p.interval(0.6)[0], 5 + 1/6) self.assertAlmostEqual(p.interval(0.6)[1], 6.5)
class TestProbabilityUtils(unittest.TestCase): def setUp(self): self.p1_dict = { 0: 0.1, 1: 0.2, 2: 0.3, 3: 0.3, 4: 0.1 # mean: 2.1 } self.p1 = ProbabilityDistribution(self.p1_dict) self.p2_dict = { (1, 2): 0.5, 4: 0.2, (4, 5): 0.3 # mean: 2.9 } self.p2 = ProbabilityDistribution(self.p2_dict) def test_can_get_probability(self): prob = self.p1.p(0) self.assertAlmostEqual(prob, 0.1) def test_can_get_probability_outside_definition(self): prob = self.p1.p(10) self.assertEqual(prob, 0) def test_can_get_probability_for_range(self): prob = self.p1.p([0, 2.5]) self.assertAlmostEqual(prob, 0.6) def test_raises_valueerror_for_badly_defined_ranges(self): # contains only one value self.assertRaisesRegex(ValueError, 'Expected length of interval to be 2, but length is', self.p1.p, [0]) # contains three values value self.assertRaisesRegex(ValueError, 'Expected length of interval to be 2, but length is', self.p1.p, [0, 1, 3]) def test_can_get_mean(self): self.assertAlmostEqual(self.p1.mean, 2.1) def test_can_get_interval(self): interval = self.p1.interval(0.95) self.assertEqual(interval, (0, 4)) interval = self.p1.interval(0.5) self.assertEqual(interval, (1, 3)) # bordering interval = self.p1.interval(0.8) self.assertEqual(interval, (1, 3)) # complete interval = self.p1.interval(1) self.assertEqual(interval, (0, 4)) def test_raises_when_interval_value_wrong(self): self.assertRaisesRegex(ValueError, 'Confidence value for interval should be', self.p1.interval, 0) self.assertRaisesRegex(ValueError, 'Confidence value for interval should be', self.p1.interval, 1.2) def test_supports_ranges_for_probabilities(self): # only in a range self.assertEqual(self.p2.p(1), 0) # specified both in range and as point self.assertEqual(self.p2.p(4), 0.2) self.assertEqual(self.p2.p((4.5, 5)), 0.15) def test_supports_ranges_for_mean(self): self.assertEqual(self.p2.mean, 2.9) def test_supports_ranges_for_interval(self): self.assertEqual(self.p2.interval(0.95), (1.05, 4.916666666666667)) def test_can_add_distributions(self): p3 = self.p1 + self.p2 self.assertEqual(type(p3), ProbabilityDistribution) p3.normalize() self.assertEqual(p3.mean, 2.5) self.assertAlmostEqual(p3.p(4), 0.15) def test_can_add_distributions_with_intersection_ranges(self): p4 = ProbabilityDistribution({ (1, 3): 0.5, (4, 5): 0.5 # mean: 3.25 }) p5 = ProbabilityDistribution({ (2, 4): 0.5, (5, 6): 0.5 # mean: 4.25 }) p6 = p4 + p5 p6.normalize() self.assertEqual(p6.mean, 3.75) self.assertEqual(p6.p((2, 3)), 0.25) def test_can_add_distributions_with_multiple_intersection_ranges(self): p4 = ProbabilityDistribution({ (1, 4): 0.3, (4, 6): 0.6, (7, 8): 0.1 }) p5 = ProbabilityDistribution({ (0, 2): 0.2, (2, 3): 0.3, (3, 9): 0.5 # intersects all ranges }) p6 = p4 + p5 p6.normalize() p6_prob_ranges_expected = { (0, 1): 0.05, (1, 2): 0.1, (2, 3): 0.2, (3, 4): (0.3/3 + 0.5/6)/2, (4, 6): (0.6 + 0.5/6*2)/2, (6, 7): (0.5/6)/2, (7, 8): (0.1 + 0.5/6)/2, (8, 9): (0.5/6)/2 } # custom asserting since the dicts contain floats self.assertEqual(set(p6.prob_ranges.keys()), set(p6_prob_ranges_expected.keys())) for key in p6.prob_ranges.keys(): # need AlmostEqual for floats self.assertAlmostEqual(p6.prob_ranges[key], p6_prob_ranges_expected[key]) def test_can_sum_distributions(self): p3 = ProbabilityDistribution({ 1: 0.5, 2: 0.5 # mean: 1.5 }) p4 = sum([self.p1, self.p2, p3]) p4.normalize() # mean: (2.1 + 2.9 + 1.5) / 3 = 2 1/6 self.assertAlmostEqual(p4.mean, 2 + 1/6) def test_raises_when_adding_not_distribution(self): try: self.p1 + 1 except TypeError: pass except Exception as e: self.fail(f'Should have raised TypeError, but raised {type(e)}') else: self.fail('Should have raised') def test_can_initiate_distribution_with_list_of_values(self): prob_dist = [ (4, 0.5), (5, 0.2), ((5, 6), 0.3) ] try: p = ProbabilityDistribution(prob_dist) except AttributeError as e: self.fail(f'Should not have raised error "{e}"') self.assertAlmostEqual(p.p(4), 0.5) self.assertAlmostEqual(p.p((5, 6)), 0.5) self.assertAlmostEqual(p.p((5.5, 6)), 0.15) def test_can_initiate_distribution_with_duplicate_entries(self): prob_dist = [ (4, 0.5), (4, 0.2), (5, 0.3) ] p = ProbabilityDistribution(prob_dist) self.assertAlmostEqual(p.p(4), 0.7) def test_can_initiate_distribution_with_intersecting_ranges(self): prob_dist = [ ((4, 7), 0.3), ((5, 8), 0.3), ((5, 6), 0.4), ] p = ProbabilityDistribution(prob_dist) self.assertAlmostEqual(p.p((5, 6)), 0.6) self.assertAlmostEqual(p.interval(0.6)[0], 5 + 1/6) self.assertAlmostEqual(p.interval(0.6)[1], 6.5) def test_can_multiply_distribution_by_number(self): p7 = self.p1 * 0.5 self.assertEqual(type(p7), ProbabilityDistribution) # this is not a true ProbabilityDistribution, probabilities sum to 0.5 self.assertAlmostEqual(p7.p(4), 0.05) def test_can_multiply_distributions_with_ranges(self): p3 = ProbabilityDistribution({(0.5, 0.625): 1.0}) p4 = p3 * 3 self.assertAlmostEqual(p4.p((0.5, 0.625)), 3) def test_can_multiply_distribution_left_hand(self): p7 = 0.5 * self.p1 self.assertEqual(type(p7), ProbabilityDistribution) # this is not a true ProbabilityDistribution, probabilities sum to 0.5 self.assertAlmostEqual(p7.p(4), 0.05) def test_can_normalize_distribution(self): p3 = ProbabilityDistribution({1: 1, (2, 3): 1}, normalize=False) self.assertFalse(p3.is_normalized) p3.normalize() self.assertTrue(p3.is_normalized) self.assertAlmostEqual(p3.p(1), 0.5) self.assertAlmostEqual(p3.mean, 1.75) def test_automatically_normalized_distribution_on_creation(self): p3 = ProbabilityDistribution({1: 1, (2, 3): 1}) # No need to call .normalize() manually. self.assertAlmostEqual(p3.p(1), 0.5) self.assertAlmostEqual(p3.mean, 1.75) def test_raises_ValueError_when_calculating_intervals_on_not_normalized_distribution(self): p = ProbabilityDistribution({1: 0.5}, normalize=False) self.assertRaisesRegex(ValueError, 'Cannot compute interval\(\) for the ProbabilityDistribution since it is not normalized', p.interval, 0.5) def test_mean_is_None_for_not_normalized_distribution(self): p = ProbabilityDistribution({1: 0.5}, normalize=False) self.assertEqual(p.mean, None) def test_add_range_to_ranges_raises_on_wrong_input(self): # Get the function from an instance. add_range_to_ranges = self.p1.add_range_to_ranges # ranges is a list self.assertRaises(ValueError, add_range_to_ranges, [], {(0, 1): 1}) # new_range_value is a list self.assertRaises(ValueError, add_range_to_ranges, {}, []) # empty new range self.assertRaises(ValueError, add_range_to_ranges, {}, {}) def test_add_range_to_ranges_simple(self): add_range_to_ranges = self.p1.add_range_to_ranges ranges = { (0, 1): 0.5, (1, 2): 0.5 } new_range_value = { (2, 3): 1 } expected = { (0, 1): 0.5, (1, 2): 0.5, (2, 3): 1 } add_range_to_ranges(ranges, new_range_value) self.assertEqual(ranges, expected) def test_add_range_to_ranges_one_intersection(self): add_range_to_ranges = self.p1.add_range_to_ranges ranges = { (0, 1): 0.5, (1, 2): 0.5 } new_range_value = { (1, 3): 1 } expected = { (0, 1): 0.5, (1, 2): 1, (2, 3): 0.5 } add_range_to_ranges(ranges, new_range_value) self.assertEqual(ranges, expected) def test_add_range_to_ranges(self): # Extra test since this went wrong, although the other tests worked. add_range_to_ranges = self.p1.add_range_to_ranges prob_ranges = {(1, 3): 0.5, (4, 5): 0.5} prob_range = {(2, 4): 0.5} add_range_to_ranges(prob_ranges, prob_range) expected_ranges = { (1, 2): 0.25, (2, 3): 0.5, (3, 4): 0.25, (4, 5): 0.5 } self.assertEqual(prob_ranges, expected_ranges) def test_split_range(self): split_range = self.p1.split_range test_range = (1, 4) points = [3, 2] self.assertEqual(split_range(test_range, points), [(1, 2), (2, 3), (3, 4)]) test_range = (1, 4) points = [4, 2] self.assertEqual(split_range(test_range, points), [(1, 2), (2, 4)]) def test_split_prob_range(self): # Extra test of values that didn't work out at first. split_prob_range = self.p1.split_prob_range prob_range = { (1, 4): 0.5 } points = [3, 2] self.assertEqual(split_prob_range(prob_range, points), { (1, 2): 1/6, (2, 3): 1/6, (3, 4): 1/6 }) points = [4, 2] self.assertEqual(split_prob_range(prob_range, points), { (1, 2): 1/6, (2, 4): 2/6 }) def test_split_prob_range_extra(self): split_prob_range = self.p1.split_prob_range prob_range = {(1, 3): 0.5} points = [1, 2, 3, 4] expected_split_prob_range = { (1, 2): 0.25, (2, 3): 0.25 } self.assertEqual(split_prob_range(prob_range, points), expected_split_prob_range) def test_split_range_extra(self): split_range = self.p1.split_range split = split_range((1,3), [1, 2, 3, 4]) expected_split = [(1, 2), (2, 3)] self.assertEqual(split, expected_split) def test_merge_prob_ranges(self): merge_prob_ranges = self.p1.merge_prob_ranges prob_ranges_1 = { (1, 3): 0.5, (4, 5): 0.5 } prob_ranges_2 = { (2, 4): 0.5, (5, 6): 0.5 } merged_prob_ranges = merge_prob_ranges(prob_ranges_1, prob_ranges_2) merged_prob_ranges_expected = { (1, 2): 0.25, (2, 3): 0.5, (3, 4): 0.25, (4, 5): 0.5, (5, 6): 0.5 } self.assertEqual(merged_prob_ranges, merged_prob_ranges_expected) def test_can_add_values(self): # This is different from adding the distributions, # but if # X ~ pd_x, Y ~ pd_y # and Z ~ pd_z, # then pd_x & pd_y ~ pd_z. # TODO: consider whether switching '+' and '&' # makes more sense? # flipping a coin coin = ProbabilityDistribution({ 0: 0.5, 1: 0.5 }) # flipping two p4 = coin & coin self.assertEqual(type(p4), ProbabilityDistribution) self.assertEqual(p4.mean, 1) self.assertEqual(p4.p(0), 0.25) self.assertEqual(p4.p(0), 0.25) def test_can_add_values_with_at_most_one_range(self): # This is different from adding the distributions, # but if # X ~ pd_x, Y ~ pd_y # and Z ~ pd_z, # then pd_x & pd_y ~ pd_z. # TODO: consider whether switching '+' and '&' # makes more sense? # flipping a coin coin = ProbabilityDistribution({ 0: 0.5, 1: 0.5 }) p3 = ProbabilityDistribution({ (0, 1): 0.5, (2, 3): 0.5 }) p4 = coin & p3 self.assertEqual(type(p4), ProbabilityDistribution) self.assertEqual(p4.mean, coin.mean + p3.mean) self.assertEqual(p4.p(0), 0) self.assertEqual(p4.p((0, 1)), 0.25) def test_raises_NotImplementedError_when_adding_distributions_with_both_ranges(self): p3 = ProbabilityDistribution({ (0, 1): 0.5, (2, 3): 0.5 }) try: p4 = p3 & p3 except NotImplementedError: pass def test_can_add_number(self): p3 = ProbabilityDistribution({ 5: 0.2, (0, 1): 0.3, (2, 3): 0.5 }) p4 = 2 & p3 self.assertEqual(type(p4), ProbabilityDistribution) self.assertEqual(p4.mean, p3.mean + 2) self.assertEqual(p4.p(7), 0.2) self.assertEqual(p4.p((4, 5)), 0.5) p5 = p3 & 2 self.assertEqual(type(p5), ProbabilityDistribution) self.assertEqual(p5.mean, p3.mean + 2) self.assertEqual(p5.p(7), 0.2) self.assertEqual(p5.p((4, 5)), 0.5) def test_can_pad_with_zero(self): # probability sums up to 0.7 p3 = ProbabilityDistribution({ 1: 0.5, 2: 0.2 }, normalize=False) # sanity check self.assertAlmostEqual(p3.p(0), 0) self.assertEqual(p3.mean, None) p3.pad() self.assertAlmostEqual(p3.p(0), 0.3) self.assertEqual(p3.mean, 0.9) def test_can_copy(self): p3 = ProbabilityDistribution({ 1: 0.5, 2: 0.2 }, normalize=False) p4 = p3.copy() self.assertEqual(type(p4), ProbabilityDistribution) self.assertNotEqual(id(p3), id(p4)) # Same as p3 self.assertEqual(p4.mean, None) p4.normalize() self.assertAlmostEqual(p4.mean, 0.9/0.7) # Original object is unchanged. self.assertEqual(p3.mean, None) def test_filters_out_p_0_values(self): p3 = ProbabilityDistribution({ 1: 1, 2: 0 }) self.assertEqual(list(p3.prob_points.keys()), [1]) def test_string_representation_sorts_keys(self): p3 = ProbabilityDistribution({ 2: 0.5, 1: 0.3, (0, 1): 0.2 }) # Note that it should be sorted expected_str = "ProbabilityDistribution({(0, 1): 0.2, 1: 0.3, 2: 0.5})" self.assertEqual(str(p3), expected_str)
def test_automatically_normalized_distribution_on_creation(self): p3 = ProbabilityDistribution({1: 1, (2, 3): 1}) # No need to call .normalize() manually. self.assertAlmostEqual(p3.p(1), 0.5) self.assertAlmostEqual(p3.mean, 1.75)