Source code for sympy.stats.frv_types

"""
Finite Discrete Random Variables - Prebuilt variable types

Contains
========
FiniteRV
DiscreteUniform
Die
Bernoulli
Coin
Binomial
Hypergeometric
"""

from __future__ import print_function, division

from sympy.stats.frv import (SingleFinitePSpace, SingleFiniteDistribution)
from sympy import (S, sympify, Rational, binomial, cacheit, Symbol, Integer,
        Dict, Basic)

__all__ = ['FiniteRV', 'DiscreteUniform', 'Die', 'Bernoulli', 'Coin',
        'Binomial', 'Hypergeometric']

def rv(name, cls, *args):
    density = cls(*args)
    return SingleFinitePSpace(name, density).value

class FiniteDistributionHandmade(SingleFiniteDistribution):
    @property
    def dict(self):
        return self.args[0]

    def __new__(cls, density):
        density = Dict(density)
        return Basic.__new__(cls, density)

[docs]def FiniteRV(name, density): """ Create a Finite Random Variable given a dict representing the density. Returns a RandomSymbol. >>> from sympy.stats import FiniteRV, P, E >>> density = {0: .1, 1: .2, 2: .3, 3: .4} >>> X = FiniteRV('X', density) >>> E(X) 2.00000000000000 >>> P(X>=2) 0.700000000000000 """ return rv(name, FiniteDistributionHandmade, density)
class DiscreteUniformDistribution(SingleFiniteDistribution): @property def p(self): return Rational(1, len(self.args)) @property @cacheit def dict(self): return dict((k, self.p) for k in self.set) @property def set(self): return self.args def pdf(self, x): if x in self.args: return self.p else: return S.Zero
[docs]def DiscreteUniform(name, items): """ Create a Finite Random Variable representing a uniform distribution over the input set. Returns a RandomSymbol. Examples ======== >>> from sympy.stats import DiscreteUniform, density >>> from sympy import symbols >>> X = DiscreteUniform('X', symbols('a b c')) # equally likely over a, b, c >>> density(X).dict {a: 1/3, b: 1/3, c: 1/3} >>> Y = DiscreteUniform('Y', list(range(5))) # distribution over a range >>> density(Y).dict {0: 1/5, 1: 1/5, 2: 1/5, 3: 1/5, 4: 1/5} """ return rv(name, DiscreteUniformDistribution, *items)
class DieDistribution(SingleFiniteDistribution): _argnames = ('sides',) @property def set(self): return list(map(Integer, list(range(1, self.sides+1)))) def pdf(self, x): x = sympify(x) if x.is_Integer and x >= 1 and x <= self.sides: return Rational(1, self.sides) else: return 0
[docs]def Die(name, sides=6): """ Create a Finite Random Variable representing a fair die. Returns a RandomSymbol. >>> from sympy.stats import Die, density >>> D6 = Die('D6', 6) # Six sided Die >>> density(D6).dict {1: 1/6, 2: 1/6, 3: 1/6, 4: 1/6, 5: 1/6, 6: 1/6} >>> D4 = Die('D4', 4) # Four sided Die >>> density(D4).dict {1: 1/4, 2: 1/4, 3: 1/4, 4: 1/4} """ return rv(name, DieDistribution, sides)
class BernoulliDistribution(SingleFiniteDistribution): _argnames = ('p', 'succ', 'fail') @property @cacheit def dict(self): return {self.succ: self.p, self.fail: 1 - self.p}
[docs]def Bernoulli(name, p, succ=1, fail=0): """ Create a Finite Random Variable representing a Bernoulli process. Returns a RandomSymbol >>> from sympy.stats import Bernoulli, density >>> from sympy import S >>> X = Bernoulli('X', S(3)/4) # 1-0 Bernoulli variable, probability = 3/4 >>> density(X).dict {0: 1/4, 1: 3/4} >>> X = Bernoulli('X', S.Half, 'Heads', 'Tails') # A fair coin toss >>> density(X).dict {Heads: 1/2, Tails: 1/2} """ return rv(name, BernoulliDistribution, p, succ, fail)
[docs]def Coin(name, p=S.Half): """ Create a Finite Random Variable representing a Coin toss. Probability p is the chance of gettings "Heads." Half by default Returns a RandomSymbol. >>> from sympy.stats import Coin, density >>> from sympy import Rational >>> C = Coin('C') # A fair coin toss >>> density(C).dict {H: 1/2, T: 1/2} >>> C2 = Coin('C2', Rational(3, 5)) # An unfair coin >>> density(C2).dict {H: 3/5, T: 2/5} """ return rv(name, BernoulliDistribution, p, 'H', 'T')
class BinomialDistribution(SingleFiniteDistribution): _argnames = ('n', 'p', 'succ', 'fail') @property @cacheit def dict(self): n, p, succ, fail = self.n, self.p, self.succ, self.fail return dict((k*succ + (n - k)*fail, binomial(n, k) * p**k * (1 - p)**(n - k)) for k in range(0, n + 1))
[docs]def Binomial(name, n, p, succ=1, fail=0): """ Create a Finite Random Variable representing a binomial distribution. Returns a RandomSymbol. Examples ======== >>> from sympy.stats import Binomial, density >>> from sympy import S >>> X = Binomial('X', 4, S.Half) # Four "coin flips" >>> density(X).dict {0: 1/16, 1: 1/4, 2: 3/8, 3: 1/4, 4: 1/16} """ return rv(name, BinomialDistribution, n, p, succ, fail)
class HypergeometricDistribution(SingleFiniteDistribution): _argnames = ('N', 'm', 'n') @property @cacheit def dict(self): N, m, n = self.N, self.m, self.n N, m, n = list(map(sympify, (N, m, n))) density = dict((sympify(k), Rational(binomial(m, k) * binomial(N - m, n - k), binomial(N, n))) for k in range(max(0, n + m - N), min(m, n) + 1)) return density
[docs]def Hypergeometric(name, N, m, n): """ Create a Finite Random Variable representing a hypergeometric distribution. Returns a RandomSymbol. Examples ======== >>> from sympy.stats import Hypergeometric, density >>> from sympy import S >>> X = Hypergeometric('X', 10, 5, 3) # 10 marbles, 5 white (success), 3 draws >>> density(X).dict {0: 1/12, 1: 5/12, 2: 5/12, 3: 1/12} """ return rv(name, HypergeometricDistribution, N, m, n)