Calculates damage probabilities

aggdmg.py

#!/usr/bin/env python2
import fractions
import itertools
import operator
import numbers
import collections

class OutcomeProb(object):
    '''combines probabilistic outcomes with operator.add'''
    def __init__(self, *args, **kwargs):
        args = (lambda: fractions.Fraction(0),)+args
        self._outcomes = collections.defaultdict(*args)
        self.default_outcome = kwargs.get("default_outcome")
        self._normalize()

    def avg(self):
        return sum(p*o for o,p in self._outcomes.iteritems()).limit_denominator(120)

    def _normalize(self):
        rem = fractions.Fraction(1)-sum(self._outcomes.itervalues())
        if rem != 0:
            self._outcomes[self.default_outcome] = rem

    def _agg(self, other):
        outcomes = collections.defaultdict(lambda: fractions.Fraction(0))
        for o1,o2 in itertools.product(self._outcomes,other._outcomes):
            outcomes[o1+o2] += self[o1]*other[o2]
        return OutcomeProb(outcomes, default_outcome=self.default_outcome)

    def _mul(self, other):
        other = fractions.Fraction(other)
        outcomes = self._outcomes.copy()
        for o in outcomes.iterkeys():
            outcomes[o] *= other
        return OutcomeProb(outcomes, default_outcome=self.default_outcome)

    def __lt__(self, outcome):
        return sum(p for o,p in self._outcomes.iteritems() if o<outcome).limit_denominator(120)

    def __le__(self, outcome):
        return sum(p for o,p in self._outcomes.iteritems() if o<=outcome).limit_denominator(120)

    def __eq__(self, outcome):
        return sum(p for o,p in self._outcomes.iteritems() if o==outcome).limit_denominator(120)

    def __ge__(self, outcome):
        return sum(p for o,p in self._outcomes.iteritems() if o>=outcome).limit_denominator(120)

    def __gt__(self, outcome):
        return sum(p for o,p in self._outcomes.iteritems() if o>outcome).limit_denominator(120)

    def __getitem__(self, outcome):
        return self._outcomes[outcome]

    def __setitem__(self, outcome, prob):
        if not 0<=prob<=1:
            raise ValueError("Probability cannot be outside [0,1]")
        self._outcomes[outcome] = prob
        self._normalize()

    def __mul__(self, other):
        if isinstance(other, numbers.Number):
            if 0<=other<=1: #probability
                return self._mul(other)
            else:
                return reduce(operator.add, [self]*int(other))
        else:
            raise TypeError("OutcomeProb multiplication requires numbers.Number type")

    def __add__(self, other):
        if isinstance(other, OutcomeProb):
            return self._agg(other)
        else:
            raise TypeError("OutcomeProb addition requires OutcomeProb type")

    def __radd__(self, other):
        return self.__add__(other)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __repr__(self):
        return "\n".join(["Outcome "+str(o)+":\t{:.0%}".format(float(p)) for o,p in self._outcomes.iteritems()]) +"\nAverage: "+str(self.avg())

if __name__=="__main__":
    att = OutcomeProb({8:fractions.Fraction(1,3),9:fractions.Fraction(1,3),10:fractions.Fraction(1,3)}, default_outcome=0)
    sui = OutcomeProb({15: fractions.Fraction(1)}, default_outcome=0)

    print "0.5*att + 0.5*att"
    print 0.5*att + 0.5*att

    print "\n0.5*(att + att)"
    print 0.5*(att+att)

    print "\natt"
    print att

    print "\n2*att"
    print 2*att

    print "\n2*att <= 18"
    print 2*att <= 18

    print "\n0.75*att + 0.5*att + 0.75*att < 8"
    print 0.75*att + 0.5*att + 0.75*att < 8

    print "\n0.8*sui + 0.5*att"
    print 0.8*sui + 0.5*att

    print "\n0.8*sui + 0.5*att < 20"
    print 0.8*sui + 0.5*att < 20

'''Example output:
att = OutcomeProb({8:fractions.Fraction(1,3),9:fractions.Fraction(1,3),10:fractions.Fraction(1,3)}, default_outcome=0)
sui = OutcomeProb({15: fractions.Fraction(1)}, default_outcome=0)

0.5*att + 0.5*att
Outcome 0:      25%
Outcome 8:      17%
Outcome 9:      17%
Outcome 10:     17%
Outcome 16:     3%
Outcome 17:     6%
Outcome 18:     8%
Outcome 19:     6%
Outcome 20:     3%
Average: 9

0.5*(att + att)
Outcome 0:      50%
Outcome 16:     6%
Outcome 17:     11%
Outcome 18:     17%
Outcome 19:     11%
Outcome 20:     6%
Average: 9

att
Outcome 8:      33%
Outcome 9:      33%
Outcome 10:     33%
Average: 9

2*att
Outcome 16:     11%
Outcome 17:     22%
Outcome 18:     33%
Outcome 19:     22%
Outcome 20:     11%
Average: 18

2*att <= 18
2/3

0.75*att + 0.5*att + 0.75*att < 8
1/32

0.8*sui + 0.5*att
Outcome 0:      10%
Outcome 8:      3%
Outcome 9:      3%
Outcome 10:     3%
Outcome 15:     40%
Outcome 23:     13%
Outcome 24:     13%
Outcome 25:     13%
Average: 33/2

0.8*sui + 0.5*att < 20
3/5
'''