dice.py

"""
Simple python dice module.
By sixthgear <[email protected]>
"""

import random
import fractions
import itertools

class Dice(object):
    
    def __init__(self, expression):
        """
        The dice object takes a standard dice notation expression
        ---        
        Example inputs:
        '3d6+12'  : roll 3 six-sided dice (summing the results) and add 12.
        '4*d12+3' : roll 1 twelve-sided die, multiply the result by 4 and add 3.
        'd100'    : roll 1 100-sided die.
        '4d6k3'   : roll 4 six-sided dice, keep the best 3 rolls.
        """
        self.expression = expression
        self.terms = self.parse(expression)
                
    def parse(self, expression):
        """
        Parse a given dice notation expression into a list of values
        """        
        expression = expression.replace('-','+-')
        terms = []        
        for t1 in expression.split('+'):        
            mterms = []
            for t2 in t1.split('*'):
                # check for dice expressions            
                dice, d, _roll = t2.partition('d')
                faces, k, keep = _roll.partition('k')
                if not d:
                    # no dice terms here, bail out
                    mterms.append(int(t2))
                    continue            
                if d and not _roll:
                    raise SyntaxError('faces term is required in dice expressions: %s' % t2)
                if k and not keep:
                    raise SyntaxError('keep term is required if you use k: %s' % t2)
                if not dice:
                    # default to a single die
                    dice = 1
                roll_terms = {}                            
                roll_terms['dice'] = int(dice)
                roll_terms['faces'] = int(faces)
                if keep:
                    if int(keep) > int(dice):                    
                        raise SyntaxError('keep term can not exceed total number of dice: %s' % t2)
                    roll_terms['keep'] = int(keep)                
                mterms.append(roll_terms)
            terms.append(mterms)
        return terms
            
    def generate(self, n):
        """
        returns a generator object that will generate a new roll on every iteration
        up to a maximum of n
        """
        while n !=0:
            yield self.roll()
            n -=1

    def roll(self):
        """
        Perform a roll, and return the value
        """
        result = 0 
        for t1 in self.terms:
            # add 
            if type(t1) is list:
                #multiply
                multiply_result = 1
                for t2 in t1:                        
                    if type(t2) is dict:
                        t2 = self._do_roll(**t2)
                    multiply_result *= t2
                t1 = multiply_result
            result += t1
        return result
            
    def _do_roll(self, faces, dice=1, keep=0):
        """
        Perform a simple dice roll, summing the totals.
        ---
        faces: the number of faces on each die
        dice: the number of dice to roll
        keep: how many dice to keep (discarding the lowest dice first)
              a value of 0 will keep all dice (default)
        """
        rolls = [random.randrange(faces)+1 for d in range(dice)]
        if keep:
            rolls.sort(reverse=True)
            rolls = rolls[:keep]
        return sum(rolls)
                
    def analyze(self):
        """
        Perform an analysis of the given dice object, determining every possible
        output, and the probabily of rolling each.        
        """
        result = {}
        
        for t1 in self.terms:
            # add 
            if type(t1) is list:
                #multiply
                multiply_result = {1: fractions.Fraction(1,1)}
                for t2 in t1:                        
                    if type(t2) is dict:
                        # calculate prob
                        t2 = self._do_analyze(**t2)                        
                    else:
                        # prob is 100%
                        t2 = {t2: fractions.Fraction(1,1)}
                    multiply_result = self._multiply_probs(multiply_result, t2)
                t1 = multiply_result
            else:
                t1 = {t1: fractions.Fraction(1,1)}
                
            if result:    
                result = self._add_probs(result, t1)
            else:
                result = t1
        return result
        
    def _multiply_probs(self,d1, d2):
        """
        Take two dictionaries of probability results, and multiply them together
        """
        result = {}
        for r1,r2 in itertools.product(d1.keys(),d2.keys()):
            if not result.has_key(r1*r2):
                result[r1*r2] = d1[r1] * d2[r2]
            else:
                result[r1*r2] += d1[r1] * d2[r2]
        return result

    def _add_probs(self,d1, d2):
        """
        Take two dictionaries of probability results, and add them together
        """
        result = {}
        for r1,r2 in itertools.product(d1.keys(),d2.keys()):
            if not result.has_key(r1+r2):
                result[r1+r2] = d1[r1] * d2[r2]
            else:
                result[r1+r2] += d1[r1] * d2[r2]
        return result
                
    def _do_analyze(self, faces, dice=1, keep=0):
        """
        Perform a simple dice analysis
        ---
        faces: the number of faces on each die
        dice: the number of dice to roll
        keep: how many dice to keep (discarding the lowest dice first)
              a value of 0 will keep all dice (default)
        """    
        def fact(x): return (1 if x==0 else x * fact(x-1))
        def combin(n,r): return (fact(n)/(fact(n-r)*fact(r)))
        result = {}
        if not keep:
            minimum = dice
        else:
            minimum = keep
        maximum = minimum * faces
        average = (minimum + maximum) / 2.0
        denominator = faces ** minimum
        
        # zip together pairs of results counting up from the minimum, and down from the maximum
        # for instance, this will zip together 2,12 3,11 4,10 5,9 6,8 and 7,7 from 2d6
        # this is because those results should have the same probability        
        for count, a in itertools.izip(itertools.count(), xrange(minimum, int(average)+1)):
            b = maximum - count
            numerator = combin(minimum+count-1,minimum-1)
            result[a] = result[b] = fractions.Fraction(numerator, denominator)            
        return result        
                

d = Dice("2d6")

for roll in d.generate(10):
    print '%s => %d' % (d.expression, roll)

print '\nProbability Analysis of %s' % d.expression
print '---'

for n, f in d.analyze().iteritems():
    line =  ('%d: ' % n).ljust(6)
    line += ('%d in %d' % (f.numerator, f.denominator)).ljust(16)
    line += '(%.3g%%)' % (float(f)*100)
    print line
 

gistfile1.txt

2d6 => 5
2d6 => 6
2d6 => 5
2d6 => 5
2d6 => 8
2d6 => 8
2d6 => 9
2d6 => 7
2d6 => 7
2d6 => 7

Probability Analysis of 2d6
---
2:    1 in 36         (2.78%)
3:    1 in 18         (5.56%)
4:    1 in 12         (8.33%)
5:    1 in 9          (11.1%)
6:    5 in 36         (13.9%)
7:    1 in 6          (16.7%)
8:    5 in 36         (13.9%)
9:    1 in 9          (11.1%)
10:   1 in 12         (8.33%)
11:   1 in 18         (5.56%)
12:   1 in 36         (2.78%)