CODEWITHSUNDEEP
pythonnumpyscipy
mvd
mvd

Reputation: 1199

how to calculate probability mass function for multinomial in scipy?

is there a way to calculate the Multinomial PMF in python, using numpy or scipy? the PMF is described here: https://en.wikipedia.org/wiki/Multinomial_distribution

scipy.stats.binom is only for binomial random variables.

Upvotes: 0

Views: 2244

Answers (2)

ev-br
ev-br

Reputation: 26070

There is no multinomial distribution in scipy just yet, although it might be available in the future version (0.18 or up).

Meanwhile, you can DIY it fairly easily:

def logpmf(self, x, n, p):
    """Log of the multinomial probability mass function.
    Parameters
    ----------
    x : array_like
        Quantiles.
    n : int
        Number of trials
    p : array_like, shape (k,)
        Probabilities. These should sum to one. If they do not, then
        ``p[-1]`` is modified to account for the remaining probability so
        that ``sum(p) == 1``.
    Returns
    -------
    logpmf : float
        Log of the probability mass function evaluated at `x`. 
    """

    x = np.asarray(x)
    if p.shape[0] != x.shape[-1]:
        raise ValueError("x & p shapes do not match.")

    coef = gammaln(n + 1) - gammaln(x + 1.).sum(axis=-1)
    val = coef + np.sum(xlogy(x, p), axis=-1)

    # insist on that the support is a set of *integers*
    mask = np.logical_and.reduce(np.mod(x, 1) == 0, axis=-1)

    mask &= (x.sum(axis=-1) == n)
    out = np.where(mask, val, -np.inf)
    return out

Here gammaln is scipy.special.gammaln, and xlogy is scipy.special.xlogy. As you see the main chunk of work is making sure the pmf is zero for non-integer values.

Upvotes: 1

John Howard
John Howard

Reputation: 64205

There is no Multinomial PMF function provided in scipy. However, you can make your own making use of the numpy.random.multinomial class to draw samples.

Upvotes: 0

Related Questions