Repository: scipy/scipy
File: scipy/stats/_discrete_distns.py (Python)

Code before:

    for :math:`k \in \{0, 1, \dots, n\}`, :math:`n \geq 0`, :math:`a > 0`,
    :math:`b > 0`, where :math:`B(a, b)` is the beta function.

    `betabinom` takes :math:`n`, :math:`a`, and :math:`b` as shape parameters.

    %(after_notes)s

    References
    ----------
    .. [1] https://en.wikipedia.org/wiki/Beta-binomial_distribution

    .. versionadded:: 1.4.0

    %(example)s
    """
    def _shape_info(self):
        return [_ShapeInfo("n", True, (0, np.inf), (True, False)),
                _ShapeInfo("a", False, (0, np.inf), (False, False)),
                _ShapeInfo("b", False, (0, np.inf), (False, False))]

    def _rvs(self, n, a, b, size=None, random_state=None):
        p = random_state.beta(a, b, size)
        if not np.all(n == np.floor(n)):
            raise ValueError("`n` must be integral.")
        return random_state.binomial(np.asarray(n, dtype=np.int64), p, size)

    def _get_support(self, n, a, b):
        return 0, n

    def _argcheck(self, n, a, b):
        return (n >= 0) & _isintegral(n) & (a > 0) & (b > 0)

    def _logpmf(self, x, n, a, b):
        k = floor(x)
        combiln = -log(n + 1) - betaln(n - k + 1, k + 1)
        return combiln + betaln(k + a, n - k + b) - betaln(a, b)

    def _pmf(self, x, n, a, b):
        return exp(self._logpmf(x, n, a, b))

    def _stats(self, n, a, b, moments='mv'):
        e_p = a / (a + b)
        e_q = 1 - e_p
        mu = n * e_p
        var = n * (a + b + n) * e_p * e_q / (a + b + 1)
        g1, g2 = None, None
        if 's' in moments:
            g1 = 1.0 / sqrt(var)
            g1 *= (a + b + 2 * n) * (b - a)
            g1 /= (a + b + 2) * (a + b)

Review comment (suggestion):
```suggestion
        return random_state.binomial(np.asarray(n, dtype=int), p, size)
```

Code after:

    for :math:`k \in \{0, 1, \dots, n\}`, :math:`n \geq 0`, :math:`a > 0`,
    :math:`b > 0`, where :math:`B(a, b)` is the beta function.

    `betabinom` takes :math:`n`, :math:`a`, and :math:`b` as shape parameters.

    %(after_notes)s

    References
    ----------
    .. [1] https://en.wikipedia.org/wiki/Beta-binomial_distribution

    .. versionadded:: 1.4.0

    %(example)s
    """
    def _shape_info(self):
        return [_ShapeInfo("n", True, (0, np.inf), (True, False)),
                _ShapeInfo("a", False, (0, np.inf), (False, False)),
                _ShapeInfo("b", False, (0, np.inf), (False, False))]

    def _rvs(self, n, a, b, size=None, random_state=None):
        p = random_state.beta(a, b, size)
        if not np.all(n == np.floor(n)):
            raise ValueError("`n` must be integral.")
        return random_state.binomial(np.asarray(n, dtype=int), p, size)

    def _get_support(self, n, a, b):
        return 0, n

    def _argcheck(self, n, a, b):
        return (n >= 0) & _isintegral(n) & (a > 0) & (b > 0)

    def _logpmf(self, x, n, a, b):
        k = floor(x)
        combiln = -log(n + 1) - betaln(n - k + 1, k + 1)
        return combiln + betaln(k + a, n - k + b) - betaln(a, b)

    def _pmf(self, x, n, a, b):
        return exp(self._logpmf(x, n, a, b))

    def _stats(self, n, a, b, moments='mv'):
        e_p = a / (a + b)
        e_q = 1 - e_p
        mu = n * e_p
        var = n * (a + b + n) * e_p * e_q / (a + b + 1)
        g1, g2 = None, None
        if 's' in moments:
            g1 = 1.0 / sqrt(var)
            g1 *= (a + b + 2 * n) * (b - a)
            g1 /= (a + b + 2) * (a + b)