Source code for sympy.stats.drv_types

from __future__ import print_function, division

from sympy.stats.drv import SingleDiscreteDistribution, SingleDiscretePSpace
from sympy import factorial, exp, S, sympify, And
from sympy.stats.rv import _value_check
from sympy.sets.sets import Interval
import random

__all__ = ['Geometric', 'Poisson']


def rv(symbol, cls, *args):
    args = list(map(sympify, args))
    dist = cls(*args)
    dist.check(*args)
    return SingleDiscretePSpace(symbol, dist).value


class PoissonDistribution(SingleDiscreteDistribution):
    _argnames = ('lamda',)

    set = S.Naturals0

    @staticmethod
    def check(lamda):
        _value_check(lamda > 0, "Lambda must be positive")

    def pdf(self, k):
        return self.lamda**k / factorial(k) * exp(-self.lamda)

    def sample(self):
        def search(x, y, u):
            while x < y:
                mid = (x + y)//2
                if u <= self.cdf(mid):
                    y = mid
                else:
                    x = mid + 1
            return x

        u = random.uniform(0, 1)
        if u <= self.cdf(S.Zero):
            return S.Zero
        n = S.One
        while True:
            if u > self.cdf(2*n):
                n *= 2
            else:
                return search(n, 2*n, u)


[docs]def Poisson(name, lamda): r""" Create a discrete random variable with a Poisson distribution. The density of the Poisson distribution is given by .. math:: f(k) := \frac{\lambda^{k} e^{- \lambda}}{k!} Parameters ========== lamda: Positive number, a rate Returns ======= A RandomSymbol. Examples ======== >>> from sympy.stats import Poisson, density, E, variance >>> from sympy import Symbol, simplify >>> rate = Symbol("lambda", positive=True) >>> z = Symbol("z") >>> X = Poisson("x", rate) >>> density(X)(z) lambda**z*exp(-lambda)/factorial(z) >>> E(X) lambda >>> simplify(variance(X)) lambda References ========== [1] http://en.wikipedia.org/wiki/Poisson_distribution [2] http://mathworld.wolfram.com/PoissonDistribution.html """ return rv(name, PoissonDistribution, lamda)
class GeometricDistribution(SingleDiscreteDistribution): _argnames = ('p',) set = S.Naturals @staticmethod def check(p): _value_check(And(0 < p, p<=1), "p must be between 0 and 1") def pdf(self, k): return (1 - self.p)**(k - 1) * self.p
[docs]def Geometric(name, p): r""" Create a discrete random variable with a Geometric distribution. The density of the Geometric distribution is given by .. math:: f(k) := p (1 - p)^{k - 1} Parameters ========== p: A probability between 0 and 1 Returns ======= A RandomSymbol. Examples ======== >>> from sympy.stats import Geometric, density, E, variance >>> from sympy import Symbol, S >>> p = S.One / 5 >>> z = Symbol("z") >>> X = Geometric("x", p) >>> density(X)(z) (4/5)**(z - 1)/5 >>> E(X) 5 >>> variance(X) 20 References ========== [1] http://en.wikipedia.org/wiki/Geometric_distribution [2] http://mathworld.wolfram.com/GeometricDistribution.html """ return rv(name, GeometricDistribution, p)