# 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, I, zeta, polylog, log, beta, hyper, binomial,
Piecewise, floor)
from sympy.stats.rv import _value_check, RandomSymbol
from sympy.stats.joint_rv_types import JointRV
from sympy.stats.joint_rv import MarginalDistribution, JointPSpace, CompoundDistribution
from sympy.stats import density
import random

__all__ = ['Geometric', 'Logarithmic', 'NegativeBinomial', 'Poisson', 'YuleSimon', 'Zeta']

def rv(symbol, cls, *args):
args = list(map(sympify, args))
dist = cls(*args)
dist.check(*args)
pspace = SingleDiscretePSpace(symbol, dist)
if any(isinstance(arg, RandomSymbol) for arg in args):
pspace = JointPSpace(symbol, CompoundDistribution(dist))
return pspace.value

#-------------------------------------------------------------------------------
# Geometric distribution ------------------------------------------------------------

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

def _characteristic_function(self, t):
p = self.p
return p * exp(I*t) / (1 - (1 - p)*exp(I*t))

def _moment_generating_function(self, t):
p = self.p
return p * exp(t) / (1 - (1 - p) * exp(t))

[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)

#-------------------------------------------------------------------------------
# Logarithmic distribution ------------------------------------------------------------

class LogarithmicDistribution(SingleDiscreteDistribution):
_argnames = ('p',)

set = S.Naturals

@staticmethod
def check(p):
_value_check(And(p > 0, p < 1), "p should be between 0 and 1")

def pdf(self, k):
p = self.p
return (-1) * p**k / (k * log(1 - p))

def _characteristic_function(self, t):
p = self.p
return log(1 - p * exp(I*t)) / log(1 - p)

def _moment_generating_function(self, t):
p = self.p
return log(1 - p * exp(t)) / log(1 - p)

def sample(self):
### TODO
raise NotImplementedError("Sampling of %s is not implemented" % density(self))

def Logarithmic(name, p):
r"""
Create a discrete random variable with a Logarithmic distribution.

The density of the Logarithmic distribution is given by

.. math::
f(k) := \frac{-p^k}{k \ln{(1 - p)}}

Parameters
==========

p: A value between 0 and 1

Returns
=======

A RandomSymbol.

Examples
========

>>> from sympy.stats import Logarithmic, density, E, variance
>>> from sympy import Symbol, S

>>> p = S.One / 5
>>> z = Symbol("z")

>>> X = Logarithmic("x", p)

>>> density(X)(z)
-5**(-z)/(z*log(4/5))

>>> E(X)
-1/(-4*log(5) + 8*log(2))

>>> variance(X)
-1/((-4*log(5) + 8*log(2))*(-2*log(5) + 4*log(2))) + 1/(-64*log(2)*log(5) + 64*log(2)**2 + 16*log(5)**2) - 10/(-32*log(5) + 64*log(2))
"""
return rv(name, LogarithmicDistribution, p)

#-------------------------------------------------------------------------------
# Negative binomial distribution ------------------------------------------------------------

class NegativeBinomialDistribution(SingleDiscreteDistribution):
_argnames = ('r', 'p')
set = S.Naturals0

@staticmethod
def check(r, p):
_value_check(r > 0, 'r should be positive')
_value_check(And(p > 0, p < 1), 'p should be between 0 and 1')

def pdf(self, k):
r = self.r
p = self.p

return binomial(k + r - 1, k) * (1 - p)**r * p**k

def _characteristic_function(self, t):
r = self.r
p = self.p

return ((1 - p) / (1 - p * exp(I*t)))**r

def _moment_generating_function(self, t):
r = self.r
p = self.p

return ((1 - p) / (1 - p * exp(t)))**r

def sample(self):
### TODO
raise NotImplementedError("Sampling of %s is not implemented" % density(self))

def NegativeBinomial(name, r, p):
r"""
Create a discrete random variable with a Negative Binomial distribution.

The density of the Negative Binomial distribution is given by

.. math::
f(k) := \binom{k + r - 1}{k} (1 - p)^r p^k

Parameters
==========

r: A positive value
p: A value between 0 and 1

Returns
=======

A RandomSymbol.

Examples
========

>>> from sympy.stats import NegativeBinomial, density, E, variance
>>> from sympy import Symbol, S

>>> r = 5
>>> p = S.One / 5
>>> z = Symbol("z")

>>> X = NegativeBinomial("x", r, p)

>>> density(X)(z)
1024*5**(-z)*binomial(z + 4, z)/3125

>>> E(X)
5/4

>>> variance(X)
25/16
"""
return rv(name, NegativeBinomialDistribution, r, p)

#-------------------------------------------------------------------------------
# Poisson distribution ------------------------------------------------------------

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)

def _characteristic_function(self, t):
return exp(self.lamda * (exp(I*t) - 1))

def _moment_generating_function(self, t):
return exp(self.lamda * (exp(t) - 1))

[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)

#-------------------------------------------------------------------------------
# Yule-Simon distribution ------------------------------------------------------------

class YuleSimonDistribution(SingleDiscreteDistribution):
_argnames = ('rho',)
set = S.Naturals

@staticmethod
def check(rho):
_value_check(rho > 0, 'rho should be positive')

def pdf(self, k):
rho = self.rho
return rho * beta(k, rho + 1)

def _cdf(self, x):
return Piecewise((1 - floor(x) * beta(floor(x), self.rho + 1), x >= 1), (0, True))

def _characteristic_function(self, t):
rho = self.rho
return rho * hyper((1, 1), (rho + 2,), exp(I*t)) * exp(I*t) / (rho + 1)

def _moment_generating_function(self, t):
rho = self.rho
return rho * hyper((1, 1), (rho + 2,), exp(t)) * exp(t) / (rho + 1)

def sample(self):
### TODO
raise NotImplementedError("Sampling of %s is not implemented" % density(self))

def YuleSimon(name, rho):
r"""
Create a discrete random variable with a Yule-Simon distribution.

The density of the Yule-Simon distribution is given by

.. math::
f(k) := \rho B(k, \rho + 1)

Parameters
==========

rho: A positive value

Returns
=======

A RandomSymbol.

Examples
========

>>> from sympy.stats import YuleSimon, density, E, variance
>>> from sympy import Symbol, simplify

>>> p = 5
>>> z = Symbol("z")

>>> X = YuleSimon("x", p)

>>> density(X)(z)
5*beta(z, 6)

>>> simplify(E(X))
5/4

>>> simplify(variance(X))
25/48
"""
return rv(name, YuleSimonDistribution, rho)

#-------------------------------------------------------------------------------
# Zeta distribution ------------------------------------------------------------

_argnames = ('s',)
set = S.Naturals

@staticmethod
def check(s):
_value_check(s > 1, 's should be greater than 1')

def pdf(self, k):
s = self.s
return 1 / (k**s * zeta(s))

def _characteristic_function(self, t):
return polylog(self.s, exp(I*t)) / zeta(self.s)

def _moment_generating_function(self, t):
return polylog(self.s, exp(t)) / zeta(self.s)

def sample(self):
### TODO
raise NotImplementedError("Sampling of %s is not implemented" % density(self))

def Zeta(name, s):
r"""
Create a discrete random variable with a Zeta distribution.

The density of the Zeta distribution is given by

.. math::
f(k) := \frac{1}{k^s \zeta{(s)}}

Parameters
==========

s: A value greater than 1

Returns
=======

A RandomSymbol.

Examples
========

>>> from sympy.stats import Zeta, density, E, variance
>>> from sympy import Symbol

>>> s = 5
>>> z = Symbol("z")

>>> X = Zeta("x", s)

>>> density(X)(z)
1/(z**5*zeta(5))

>>> E(X)
pi**4/(90*zeta(5))

>>> variance(X)
-pi**8/(8100*zeta(5)**2) + zeta(3)/zeta(5)
"""