Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
allendowney
GitHub Repository: allendowney/thinkbayes2
 Path: blob/master/scripts/dungeons.py
1901 views
1
"""This file contains code for use with "Think Bayes",

2
by Allen B. Downey, available from greenteapress.com

3


4
Copyright 2012 Allen B. Downey

5
License: GNU GPLv3 http://www.gnu.org/licenses/gpl.html

6
"""

7


8
from __future__ import print_function, division

9


10
import random

11


12
import thinkbayes2

13
import thinkplot

14


15
FORMATS = ['pdf', 'eps', 'png']

16


17


18
class Die(thinkbayes2.Pmf):

19
    """Represents the PMF of outcomes for a die."""

20


21
    def __init__(self, sides, label=''):

22
        """Initializes the die.

23


24
        sides: int number of sides

25
        label: string

26
        """

27
        hypos = range(1, sides+1)

28
        thinkbayes2.Pmf.__init__(self, hypos, label=label)

29


30


31
def PmfMax(pmf1, pmf2):

32
    """Computes the distribution of the max of values drawn from two Pmfs.

33


34
    pmf1, pmf2: Pmf objects

35


36
    returns: new Pmf

37
    """

38
    res = thinkbayes2.Pmf()

39
    for v1, p1 in pmf1.Items():

40
        for v2, p2 in pmf2.Items():

41
            res.Incr(max(v1, v2), p1*p2)

42
    return res

43
    

44


45
def main():

46
    pmf_dice = thinkbayes2.Pmf()

47
    pmf_dice.Set(Die(4), 5)

48
    pmf_dice.Set(Die(6), 4)

49
    pmf_dice.Set(Die(8), 3)

50
    pmf_dice.Set(Die(12), 2)

51
    pmf_dice.Set(Die(20), 1)

52
    pmf_dice.Normalize()

53


54
    mix = thinkbayes2.Pmf()

55
    for die, weight in pmf_dice.Items():

56
        for outcome, prob in die.Items():

57
            mix.Incr(outcome, weight*prob)

58


59
    mix = thinkbayes2.MakeMixture(pmf_dice)

60


61
    thinkplot.Hist(mix, width=0.9)

62
    thinkplot.Save(root='dungeons3',

63
                xlabel='Outcome',

64
                ylabel='Probability',

65
                formats=FORMATS)

66


67
    random.seed(17)

68


69
    d6 = Die(6, 'd6')

70


71
    dice = [d6] * 3

72
    three = thinkbayes2.SampleSum(dice, 1000)

73
    three.label = 'sample'

74
    three.Print()

75


76
    three_exact = d6 + d6 + d6

77
    three_exact.label = 'exact'

78
    three_exact.Print()

79


80
    thinkplot.PrePlot(num=2)

81
    thinkplot.Pmf(three)

82
    thinkplot.Pmf(three_exact, linestyle='dashed')

83
    thinkplot.Save(root='dungeons1',

84
                xlabel='Sum of three d6',

85
                ylabel='Probability',

86
                axis=[2, 19, 0, 0.15],

87
                formats=FORMATS)

88


89
    thinkplot.Clf()

90
    thinkplot.PrePlot(num=1)

91
    

92
    # compute the distribution of the best attribute the hard way

93
    best_attr2 = PmfMax(three_exact, three_exact)

94
    best_attr4 = PmfMax(best_attr2, best_attr2)

95
    best_attr6 = PmfMax(best_attr4, best_attr2)

96
    # thinkplot.Pmf(best_attr6)

97


98
    # and the easy way

99
    best_attr_cdf = three_exact.Max(6)

100
    best_attr_cdf.label = ''

101
    best_attr_pmf = best_attr_cdf.MakePmf()

102
    best_attr_pmf.Print()

103


104
    thinkplot.Pmf(best_attr_pmf)

105
    thinkplot.Save(root='dungeons2',

106
                   xlabel='Best of three d6',

107
                   ylabel='Probability',

108
                   axis=[2, 19, 0, 0.23],

109
                   formats=FORMATS,

110
                   legend=False)

111
    

112


113


114
if __name__ == '__main__':

115
    main()

116


117