ztnb_em.py 文件源码

def ztnb_pmf(y, mu, alpha):
    r = 1.0 / alpha
    if y <= 0:
        raise Exception('y must be larger than 0.')
    p = mu/(mu+r+0.0)
    ztnbin_mpmath = lambda y, p, r: mpmath.gamma(y + r)/(mpmath.gamma(y+1)*mpmath.gamma(r))*np.power(1-p, r)*np.power(p, y)/(1-np.power(1-p, r))
    ztnbin = np.frompyfunc(ztnbin_mpmath, 3, 1)
    return float(ztnbin(y, p, r))