Traktor/myenv/Lib/site-packages/sympy/stats/sampling/sample_scipy.py
2024-05-26 05:12:46 +02:00

168 lines
6.2 KiB
Python

from functools import singledispatch
from sympy.core.symbol import Dummy
from sympy.functions.elementary.exponential import exp
from sympy.utilities.lambdify import lambdify
from sympy.external import import_module
from sympy.stats import DiscreteDistributionHandmade
from sympy.stats.crv import SingleContinuousDistribution
from sympy.stats.crv_types import ChiSquaredDistribution, ExponentialDistribution, GammaDistribution, \
LogNormalDistribution, NormalDistribution, ParetoDistribution, UniformDistribution, BetaDistribution, \
StudentTDistribution, CauchyDistribution
from sympy.stats.drv_types import GeometricDistribution, LogarithmicDistribution, NegativeBinomialDistribution, \
PoissonDistribution, SkellamDistribution, YuleSimonDistribution, ZetaDistribution
from sympy.stats.frv import SingleFiniteDistribution
scipy = import_module("scipy", import_kwargs={'fromlist':['stats']})
@singledispatch
def do_sample_scipy(dist, size, seed):
return None
# CRV
@do_sample_scipy.register(SingleContinuousDistribution)
def _(dist: SingleContinuousDistribution, size, seed):
# if we don't need to make a handmade pdf, we won't
import scipy.stats
z = Dummy('z')
handmade_pdf = lambdify(z, dist.pdf(z), ['numpy', 'scipy'])
class scipy_pdf(scipy.stats.rv_continuous):
def _pdf(dist, x):
return handmade_pdf(x)
scipy_rv = scipy_pdf(a=float(dist.set._inf),
b=float(dist.set._sup), name='scipy_pdf')
return scipy_rv.rvs(size=size, random_state=seed)
@do_sample_scipy.register(ChiSquaredDistribution)
def _(dist: ChiSquaredDistribution, size, seed):
# same parametrisation
return scipy.stats.chi2.rvs(df=float(dist.k), size=size, random_state=seed)
@do_sample_scipy.register(ExponentialDistribution)
def _(dist: ExponentialDistribution, size, seed):
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.expon.html#scipy.stats.expon
return scipy.stats.expon.rvs(scale=1 / float(dist.rate), size=size, random_state=seed)
@do_sample_scipy.register(GammaDistribution)
def _(dist: GammaDistribution, size, seed):
# https://stackoverflow.com/questions/42150965/how-to-plot-gamma-distribution-with-alpha-and-beta-parameters-in-python
return scipy.stats.gamma.rvs(a=float(dist.k), scale=float(dist.theta), size=size, random_state=seed)
@do_sample_scipy.register(LogNormalDistribution)
def _(dist: LogNormalDistribution, size, seed):
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.lognorm.html
return scipy.stats.lognorm.rvs(scale=float(exp(dist.mean)), s=float(dist.std), size=size, random_state=seed)
@do_sample_scipy.register(NormalDistribution)
def _(dist: NormalDistribution, size, seed):
return scipy.stats.norm.rvs(loc=float(dist.mean), scale=float(dist.std), size=size, random_state=seed)
@do_sample_scipy.register(ParetoDistribution)
def _(dist: ParetoDistribution, size, seed):
# https://stackoverflow.com/questions/42260519/defining-pareto-distribution-in-python-scipy
return scipy.stats.pareto.rvs(b=float(dist.alpha), scale=float(dist.xm), size=size, random_state=seed)
@do_sample_scipy.register(StudentTDistribution)
def _(dist: StudentTDistribution, size, seed):
return scipy.stats.t.rvs(df=float(dist.nu), size=size, random_state=seed)
@do_sample_scipy.register(UniformDistribution)
def _(dist: UniformDistribution, size, seed):
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.uniform.html
return scipy.stats.uniform.rvs(loc=float(dist.left), scale=float(dist.right - dist.left), size=size, random_state=seed)
@do_sample_scipy.register(BetaDistribution)
def _(dist: BetaDistribution, size, seed):
# same parametrisation
return scipy.stats.beta.rvs(a=float(dist.alpha), b=float(dist.beta), size=size, random_state=seed)
@do_sample_scipy.register(CauchyDistribution)
def _(dist: CauchyDistribution, size, seed):
return scipy.stats.cauchy.rvs(loc=float(dist.x0), scale=float(dist.gamma), size=size, random_state=seed)
# DRV:
@do_sample_scipy.register(DiscreteDistributionHandmade)
def _(dist: DiscreteDistributionHandmade, size, seed):
from scipy.stats import rv_discrete
z = Dummy('z')
handmade_pmf = lambdify(z, dist.pdf(z), ['numpy', 'scipy'])
class scipy_pmf(rv_discrete):
def _pmf(dist, x):
return handmade_pmf(x)
scipy_rv = scipy_pmf(a=float(dist.set._inf), b=float(dist.set._sup),
name='scipy_pmf')
return scipy_rv.rvs(size=size, random_state=seed)
@do_sample_scipy.register(GeometricDistribution)
def _(dist: GeometricDistribution, size, seed):
return scipy.stats.geom.rvs(p=float(dist.p), size=size, random_state=seed)
@do_sample_scipy.register(LogarithmicDistribution)
def _(dist: LogarithmicDistribution, size, seed):
return scipy.stats.logser.rvs(p=float(dist.p), size=size, random_state=seed)
@do_sample_scipy.register(NegativeBinomialDistribution)
def _(dist: NegativeBinomialDistribution, size, seed):
return scipy.stats.nbinom.rvs(n=float(dist.r), p=float(dist.p), size=size, random_state=seed)
@do_sample_scipy.register(PoissonDistribution)
def _(dist: PoissonDistribution, size, seed):
return scipy.stats.poisson.rvs(mu=float(dist.lamda), size=size, random_state=seed)
@do_sample_scipy.register(SkellamDistribution)
def _(dist: SkellamDistribution, size, seed):
return scipy.stats.skellam.rvs(mu1=float(dist.mu1), mu2=float(dist.mu2), size=size, random_state=seed)
@do_sample_scipy.register(YuleSimonDistribution)
def _(dist: YuleSimonDistribution, size, seed):
return scipy.stats.yulesimon.rvs(alpha=float(dist.rho), size=size, random_state=seed)
@do_sample_scipy.register(ZetaDistribution)
def _(dist: ZetaDistribution, size, seed):
return scipy.stats.zipf.rvs(a=float(dist.s), size=size, random_state=seed)
# FRV:
@do_sample_scipy.register(SingleFiniteDistribution)
def _(dist: SingleFiniteDistribution, size, seed):
# scipy can handle with custom distributions
from scipy.stats import rv_discrete
density_ = dist.dict
x, y = [], []
for k, v in density_.items():
x.append(int(k))
y.append(float(v))
scipy_rv = rv_discrete(name='scipy_rv', values=(x, y))
return scipy_rv.rvs(size=size, random_state=seed)