alternative parametrization of the negative binomial in scipy

前端 未结 2 851
一个人的身影
一个人的身影 2021-02-14 20:07

In scipy the negative binomial distribution is defined as:

nbinom.pmf(k) = choose(k+n-1, n-1) * p**n * (1-p)**k

This is the common definition,

2条回答
  •  余生分开走
    2021-02-14 20:36

    from scipy.stats import nbinom
    
    
    def convert_params(mu, theta):
        """
        Convert mean/dispersion parameterization of a negative binomial to the ones scipy supports
    
        See https://en.wikipedia.org/wiki/Negative_binomial_distribution#Alternative_formulations
        """
        r = theta
        var = mu + 1 / r * mu ** 2
        p = (var - mu) / var
        return r, 1 - p
    
    
    def pmf(counts, mu, theta):
        """
        >>> import numpy as np
        >>> from scipy.stats import poisson
        >>> np.isclose(pmf(10, 10, 10000), poisson.pmf(10, 10), atol=1e-3)
        True
        """
        return nbinom.pmf(counts, *convert_params(mu, theta))
    
    
    def logpmf(counts, mu, theta):
        return nbinom.logpmf(counts, *convert_params(mu, theta))
    
    
    def cdf(counts, mu, theta):
        return nbinom.cdf(counts, *convert_params(mu, theta))
    
    
    def sf(counts, mu, theta):
        return nbinom.sf(counts, *convert_params(mu, theta))
    

提交回复
热议问题