{
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {
            "display_name": "Python",
            "language": "python",
            "name": "python"
        },
        "language_info": {
            "name": "python"
        }
    },
    "cells": [
        {
            "id": "1a48de01",
            "cell_type": "markdown",
            "source": "<div class='alert alert-warning'>\n\nSciPy's interactive examples with Jupyterlite are experimental and may not always work as expected. Execution of cells containing imports may result in large downloads (up to 60MB of content for the first import from SciPy). Load times when importing from SciPy may take roughly 10-20 seconds. If you notice any problems, feel free to open an [issue](https://github.com/scipy/scipy/issues/new/choose).\n\n</div>",
            "metadata": {}
        },
        {
            "id": "b65f1cca",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "from scipy.stats import dirichlet_multinomial",
            "outputs": []
        },
        {
            "id": "5444eabe",
            "cell_type": "markdown",
            "source": "Get the PMF\n",
            "metadata": {}
        },
        {
            "id": "360e3f65",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "n = 6  # number of trials\nalpha = [3, 4, 5]  # concentration parameters\nx = [1, 2, 3]  # counts\ndirichlet_multinomial.pmf(x, alpha, n)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "0.08484162895927604"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "ee9b731d",
            "cell_type": "markdown",
            "source": "If the sum of category counts does not equal the number of trials,\nthe probability mass is zero.\n",
            "metadata": {}
        },
        {
            "id": "3e07d988",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dirichlet_multinomial.pmf(x, alpha, n=7)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "0.0"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "5d08aa86",
            "cell_type": "markdown",
            "source": "Get the log of the PMF\n",
            "metadata": {}
        },
        {
            "id": "084e6033",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dirichlet_multinomial.logpmf(x, alpha, n)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "-2.4669689491013327"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "22e08595",
            "cell_type": "markdown",
            "source": "Get the mean\n",
            "metadata": {}
        },
        {
            "id": "6117f1a0",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dirichlet_multinomial.mean(alpha, n)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([1.5, 2. , 2.5])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "b14c52e5",
            "cell_type": "markdown",
            "source": "Get the variance\n",
            "metadata": {}
        },
        {
            "id": "2ab9f4ee",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dirichlet_multinomial.var(alpha, n)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([1.55769231, 1.84615385, 2.01923077])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "01005fff",
            "cell_type": "markdown",
            "source": "Get the covariance\n",
            "metadata": {}
        },
        {
            "id": "c6f37377",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dirichlet_multinomial.cov(alpha, n)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([[ 1.55769231, -0.69230769, -0.86538462],\n       [-0.69230769,  1.84615385, -1.15384615],\n       [-0.86538462, -1.15384615,  2.01923077]])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "93fe0ff6",
            "cell_type": "markdown",
            "source": "Alternatively, the object may be called (as a function) to fix the\n`alpha` and `n` parameters, returning a \"frozen\" Dirichlet multinomial\nrandom variable.\n",
            "metadata": {}
        },
        {
            "id": "fb36628b",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dm = dirichlet_multinomial(alpha, n)\ndm.pmf(x)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "0.08484162895927579"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "3d47a33f",
            "cell_type": "markdown",
            "source": "All methods are fully vectorized. Each element of `x` and `alpha` is\na vector (along the last axis), each element of `n` is an\ninteger (scalar), and the result is computed element-wise.\n",
            "metadata": {}
        },
        {
            "id": "d15b5147",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "x = [[1, 2, 3], [4, 5, 6]]\nalpha = [[1, 2, 3], [4, 5, 6]]\nn = [6, 15]\ndirichlet_multinomial.pmf(x, alpha, n)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([0.06493506, 0.02626937])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "1d48a99d",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "dirichlet_multinomial.cov(alpha, n).shape  # both covariance matrices",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "(2, 3, 3)"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "b4684b9c",
            "cell_type": "markdown",
            "source": "Broadcasting according to standard NumPy conventions is supported. Here,\nwe have four sets of concentration parameters (each a two element vector)\nfor each of three numbers of trials (each a scalar).\n",
            "metadata": {}
        },
        {
            "id": "ff429992",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "alpha = [[3, 4], [4, 5], [5, 6], [6, 7]]\nn = [[6], [7], [8]]\ndirichlet_multinomial.mean(alpha, n).shape",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "(3, 4, 2)"
                    },
                    "execution_count": null
                }
            ]
        }
    ]
}