{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mark and Recapture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Think Bayes, Second Edition\n",
    "\n",
    "Copyright 2020 Allen B. Downey\n",
    "\n",
    "License: [Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)](https://creativecommons.org/licenses/by-nc-sa/4.0/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.028418Z",
     "iopub.status.busy": "2021-04-16T19:37:47.027988Z",
     "iopub.status.idle": "2021-04-16T19:37:47.030142Z",
     "shell.execute_reply": "2021-04-16T19:37:47.029674Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# If we're running on Colab, install empiricaldist\n",
    "# https://pypi.org/project/empiricaldist/\n",
    "\n",
    "import sys\n",
    "IN_COLAB = 'google.colab' in sys.modules\n",
    "\n",
    "if IN_COLAB:\n",
    "    !pip install empiricaldist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.033444Z",
     "iopub.status.busy": "2021-04-16T19:37:47.033007Z",
     "iopub.status.idle": "2021-04-16T19:37:47.034750Z",
     "shell.execute_reply": "2021-04-16T19:37:47.035100Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get utils.py\n",
    "\n",
    "from os.path import basename, exists\n",
    "\n",
    "def download(url):\n",
    "    filename = basename(url)\n",
    "    if not exists(filename):\n",
    "        from urllib.request import urlretrieve\n",
    "        local, _ = urlretrieve(url, filename)\n",
    "        print('Downloaded ' + local)\n",
    "    \n",
    "download('https://github.com/AllenDowney/ThinkBayes2/raw/master/soln/utils.py')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.038782Z",
     "iopub.status.busy": "2021-04-16T19:37:47.038205Z",
     "iopub.status.idle": "2021-04-16T19:37:47.706520Z",
     "shell.execute_reply": "2021-04-16T19:37:47.706106Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import set_pyplot_params\n",
    "set_pyplot_params()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This chapter introduces \"mark and recapture\" experiments, in which we sample individuals from a population, mark them somehow, and then take a second sample from the same population.  Seeing how many individuals in the second sample are marked, we can estimate the size of the population.\n",
    "\n",
    "Experiments like this were originally used in ecology, but turn out to be useful in many other fields.  Examples in this chapter include software engineering and epidemiology.\n",
    "\n",
    "Also, in this chapter we'll work with models that have three parameters, so we'll extend the joint distributions we've been using to three dimensions.\n",
    "\n",
    "But first, grizzly bears."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Grizzly Bear Problem\n",
    "\n",
    "In 1996 and 1997 researchers deployed bear traps in locations in British Columbia and Alberta, Canada, in an effort to estimate the population of grizzly bears.  They describe the experiment in [this article](https://www.researchgate.net/publication/229195465_Estimating_Population_Size_of_Grizzly_Bears_Using_Hair_Capture_DNA_Profiling_and_Mark-Recapture_Analysis).\n",
    "\n",
    "The \"trap\" consists of  a lure and several strands of barbed wire intended to capture samples of hair from bears that visit the lure.  Using the hair samples, the researchers use DNA analysis to identify individual bears.\n",
    "\n",
    "During the first session, the researchers deployed traps at 76 sites.  Returning 10 days later, they obtained 1043 hair samples and identified 23 different bears.  During a second 10-day session they obtained 1191 samples from 19 different bears, where 4 of the 19 were from bears they had identified in the first batch.\n",
    "\n",
    "To estimate the population of bears from this data, we need a model for the probability that each bear will be observed during each session.  As a starting place, we'll make the simplest assumption, that every bear in the population has the same (unknown) probability of being sampled during each session."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With these assumptions we can compute the probability of the data for a range of possible populations.\n",
    "\n",
    "As an example, let's suppose that the actual population of bears is 100.\n",
    "\n",
    "After the first session, 23 of the 100 bears have been identified.\n",
    "During the second session, if we choose 19 bears at random, what is the probability that 4 of them were previously identified?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I'll define\n",
    "\n",
    "* $N$: actual population size, 100.\n",
    "\n",
    "* $K$: number of bears identified in the first session, 23.\n",
    "\n",
    "* $n$: number of bears observed in the second session, 19 in the example.\n",
    "\n",
    "* $k$: number of bears in the second session that were previously identified, 4.\n",
    "\n",
    "For given values of $N$, $K$, and $n$, the probability of finding $k$ previously-identified bears is given by the [hypergeometric distribution](https://en.wikipedia.org/wiki/Hypergeometric_distribution):\n",
    "\n",
    "$$\\binom{K}{k} \\binom{N-K}{n-k}/ \\binom{N}{n}$$\n",
    "\n",
    "where the [binomial coefficient](https://en.wikipedia.org/wiki/Binomial_coefficient), $\\binom{K}{k}$, is the number of subsets of size $k$ we can choose from a population of size $K$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To understand why, consider: \n",
    "\n",
    "* The denominator, $\\binom{N}{n}$, is the number of subsets of $n$ we could choose from a population of $N$ bears.\n",
    "\n",
    "* The numerator is the number of subsets that contain $k$ bears from the previously identified $K$ and $n-k$ from the previously unseen $N-K$.\n",
    "\n",
    "SciPy provides `hypergeom`, which we can use to compute this probability for a range of values of $k$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.711709Z",
     "iopub.status.busy": "2021-04-16T19:37:47.711259Z",
     "iopub.status.idle": "2021-04-16T19:37:47.712868Z",
     "shell.execute_reply": "2021-04-16T19:37:47.713221Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.stats import hypergeom\n",
    "\n",
    "N = 100\n",
    "K = 23\n",
    "n = 19\n",
    "\n",
    "ks = np.arange(12)\n",
    "ps = hypergeom(N, K, n).pmf(ks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is the distribution of $k$ with given parameters $N$, $K$, and $n$.\n",
    "Here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.727971Z",
     "iopub.status.busy": "2021-04-16T19:37:47.725947Z",
     "iopub.status.idle": "2021-04-16T19:37:47.888438Z",
     "shell.execute_reply": "2021-04-16T19:37:47.888796Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from utils import decorate\n",
    "\n",
    "plt.bar(ks, ps)\n",
    "\n",
    "decorate(xlabel='Number of bears observed twice',\n",
    "         ylabel='PMF',\n",
    "         title='Hypergeometric distribution of k (known population 100)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most likely value of $k$ is 4, which is the value actually observed in the experiment.  \n",
    "That suggests that $N=100$ is a reasonable estimate of the population, given this data.\n",
    "\n",
    "We've computed the distribution of $k$ given $N$, $K$, and $n$.\n",
    "Now let's go the other way: given $K$, $n$, and $k$, how can we estimate the total population, $N$?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Update\n",
    "\n",
    "As a starting place, let's suppose that, prior to this study, an expert estimates that the local bear population is between 50 and 500, and equally likely to be any value in that range.\n",
    "\n",
    "I'll use `make_uniform` to make a uniform distribution of integers in this range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.893465Z",
     "iopub.status.busy": "2021-04-16T19:37:47.892728Z",
     "iopub.status.idle": "2021-04-16T19:37:47.900265Z",
     "shell.execute_reply": "2021-04-16T19:37:47.899629Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from utils import make_uniform\n",
    "\n",
    "qs = np.arange(50, 501)\n",
    "prior_N = make_uniform(qs, name='N')\n",
    "prior_N.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So that's our prior.\n",
    "\n",
    "To compute the likelihood of the data, we can use `hypergeom` with constants `K` and `n`, and a range of values of `N`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.906764Z",
     "iopub.status.busy": "2021-04-16T19:37:47.906124Z",
     "iopub.status.idle": "2021-04-16T19:37:47.908285Z",
     "shell.execute_reply": "2021-04-16T19:37:47.908789Z"
    }
   },
   "outputs": [],
   "source": [
    "Ns = prior_N.qs\n",
    "K = 23\n",
    "n = 19\n",
    "k = 4\n",
    "\n",
    "likelihood = hypergeom(Ns, K, n).pmf(k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the posterior in the usual way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.914099Z",
     "iopub.status.busy": "2021-04-16T19:37:47.913348Z",
     "iopub.status.idle": "2021-04-16T19:37:47.916623Z",
     "shell.execute_reply": "2021-04-16T19:37:47.916072Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N = prior_N * likelihood\n",
    "posterior_N.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:47.920279Z",
     "iopub.status.busy": "2021-04-16T19:37:47.919818Z",
     "iopub.status.idle": "2021-04-16T19:37:48.079352Z",
     "shell.execute_reply": "2021-04-16T19:37:48.078946Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_N.plot(color='C4')\n",
    "\n",
    "decorate(xlabel='Population of bears (N)',\n",
    "         ylabel='PDF',\n",
    "         title='Posterior distribution of N')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most likely value is 109."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.083122Z",
     "iopub.status.busy": "2021-04-16T19:37:48.082536Z",
     "iopub.status.idle": "2021-04-16T19:37:48.085201Z",
     "shell.execute_reply": "2021-04-16T19:37:48.084795Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N.max_prob()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But the distribution is skewed to the right, so the posterior mean is substantially higher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.088879Z",
     "iopub.status.busy": "2021-04-16T19:37:48.088365Z",
     "iopub.status.idle": "2021-04-16T19:37:48.091018Z",
     "shell.execute_reply": "2021-04-16T19:37:48.090654Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the credible interval is quite wide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.095317Z",
     "iopub.status.busy": "2021-04-16T19:37:48.094716Z",
     "iopub.status.idle": "2021-04-16T19:37:48.097561Z",
     "shell.execute_reply": "2021-04-16T19:37:48.097064Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N.credible_interval(0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This solution is relatively simple, but it turns out we can do a little better if we model the unknown probability of observing a bear explicitly."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two-Parameter Model\n",
    "\n",
    "Next we'll try a model with two parameters: the number of bears, `N`, and the probability of observing a bear, `p`.\n",
    "\n",
    "We'll assume that the probability is the same in both rounds, which is probably reasonable in this case because it is the same kind of trap in the same place.\n",
    "\n",
    "We'll also assume that the probabilities are independent; that is, the probability a bear is observed in the second round does not depend on whether it was observed in the first round.  This assumption might be less reasonable, but for now it is a necessary simplification.\n",
    "\n",
    "Here are the counts again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.100459Z",
     "iopub.status.busy": "2021-04-16T19:37:48.100002Z",
     "iopub.status.idle": "2021-04-16T19:37:48.102455Z",
     "shell.execute_reply": "2021-04-16T19:37:48.102065Z"
    }
   },
   "outputs": [],
   "source": [
    "K = 23\n",
    "n = 19\n",
    "k = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this model, I'll express the data in a notation that will make it easier to generalize to more than two rounds: \n",
    "\n",
    "* `k10` is the number of bears observed in the first round but not the second,\n",
    "\n",
    "* `k01` is the number of bears observed in the second round but not the first, and\n",
    "\n",
    "* `k11` is the number of bears observed in both rounds.\n",
    "\n",
    "Here are their values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.105589Z",
     "iopub.status.busy": "2021-04-16T19:37:48.105048Z",
     "iopub.status.idle": "2021-04-16T19:37:48.106939Z",
     "shell.execute_reply": "2021-04-16T19:37:48.107364Z"
    }
   },
   "outputs": [],
   "source": [
    "k10 = 23 - 4\n",
    "k01 = 19 - 4\n",
    "k11 = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we know the actual values of `N` and `p`.  We can use them to compute the likelihood of this data.\n",
    "\n",
    "For example, suppose we know that `N=100` and `p=0.2`.\n",
    "We can use `N` to compute `k00`, which is the number of unobserved bears."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.111393Z",
     "iopub.status.busy": "2021-04-16T19:37:48.110772Z",
     "iopub.status.idle": "2021-04-16T19:37:48.113489Z",
     "shell.execute_reply": "2021-04-16T19:37:48.113917Z"
    }
   },
   "outputs": [],
   "source": [
    "N = 100\n",
    "\n",
    "observed = k01 + k10 + k11\n",
    "k00 = N - observed\n",
    "k00"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the update, it will be convenient to store the data as a list that represents the number of bears in each category."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.117755Z",
     "iopub.status.busy": "2021-04-16T19:37:48.117071Z",
     "iopub.status.idle": "2021-04-16T19:37:48.119882Z",
     "shell.execute_reply": "2021-04-16T19:37:48.120304Z"
    }
   },
   "outputs": [],
   "source": [
    "x = [k00, k01, k10, k11]\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, if we know `p=0.2`, we can compute the probability a bear falls in each category.  For example, the probability of being observed in both rounds is `p*p`, and the probability of being unobserved in both rounds is `q*q` (where `q=1-p`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.124805Z",
     "iopub.status.busy": "2021-04-16T19:37:48.124072Z",
     "iopub.status.idle": "2021-04-16T19:37:48.127454Z",
     "shell.execute_reply": "2021-04-16T19:37:48.126952Z"
    }
   },
   "outputs": [],
   "source": [
    "p = 0.2\n",
    "q = 1-p\n",
    "y = [q*q, q*p, p*q, p*p]\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the probability of the data is given by the [multinomial distribution](https://en.wikipedia.org/wiki/Multinomial_distribution):\n",
    "\n",
    "$$\\frac{N!}{\\prod x_i!} \\prod y_i^{x_i}$$\n",
    "\n",
    "where $N$ is actual population, $x$ is a sequence with the counts in each category, and $y$ is a sequence of probabilities for each category.\n",
    "\n",
    "SciPy provides `multinomial`, which provides `pmf`, which computes this probability.\n",
    "Here is the probability of the data for these values of `N` and `p`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.132720Z",
     "iopub.status.busy": "2021-04-16T19:37:48.131899Z",
     "iopub.status.idle": "2021-04-16T19:37:48.135127Z",
     "shell.execute_reply": "2021-04-16T19:37:48.134684Z"
    }
   },
   "outputs": [],
   "source": [
    "from scipy.stats import multinomial\n",
    "\n",
    "likelihood = multinomial.pmf(x, N, y)\n",
    "likelihood"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's the likelihood if we know `N` and `p`, but of course we don't.  So we'll choose prior distributions for `N` and `p`, and use the likelihoods to update it. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Prior\n",
    "\n",
    "We'll use `prior_N` again for the prior distribution of `N`, and a uniform prior for the probability of observing a bear, `p`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.139275Z",
     "iopub.status.busy": "2021-04-16T19:37:48.138848Z",
     "iopub.status.idle": "2021-04-16T19:37:48.140998Z",
     "shell.execute_reply": "2021-04-16T19:37:48.140599Z"
    }
   },
   "outputs": [],
   "source": [
    "qs = np.linspace(0, 0.99, num=100)\n",
    "prior_p = make_uniform(qs, name='p')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can make a joint distribution in the usual way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.144506Z",
     "iopub.status.busy": "2021-04-16T19:37:48.143690Z",
     "iopub.status.idle": "2021-04-16T19:37:48.148207Z",
     "shell.execute_reply": "2021-04-16T19:37:48.147721Z"
    }
   },
   "outputs": [],
   "source": [
    "from utils import make_joint\n",
    "\n",
    "joint_prior = make_joint(prior_p, prior_N)\n",
    "joint_prior.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a Pandas `DataFrame` with values of `N` down the rows and values of `p` across the columns.\n",
    "However, for this problem it will be convenient to represent the prior distribution as a 1-D `Series` rather than a 2-D `DataFrame`.\n",
    "We can convert from one format to the other using `stack`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.152464Z",
     "iopub.status.busy": "2021-04-16T19:37:48.151856Z",
     "iopub.status.idle": "2021-04-16T19:37:48.161492Z",
     "shell.execute_reply": "2021-04-16T19:37:48.161060Z"
    }
   },
   "outputs": [],
   "source": [
    "from empiricaldist import Pmf\n",
    "\n",
    "joint_pmf = Pmf(joint_prior.stack())\n",
    "joint_pmf.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.164795Z",
     "iopub.status.busy": "2021-04-16T19:37:48.164220Z",
     "iopub.status.idle": "2021-04-16T19:37:48.166498Z",
     "shell.execute_reply": "2021-04-16T19:37:48.166851Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "type(joint_pmf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.170172Z",
     "iopub.status.busy": "2021-04-16T19:37:48.169707Z",
     "iopub.status.idle": "2021-04-16T19:37:48.171922Z",
     "shell.execute_reply": "2021-04-16T19:37:48.172268Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "type(joint_pmf.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.175413Z",
     "iopub.status.busy": "2021-04-16T19:37:48.174911Z",
     "iopub.status.idle": "2021-04-16T19:37:48.177293Z",
     "shell.execute_reply": "2021-04-16T19:37:48.177641Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "joint_pmf.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a `Pmf` whose index is a `MultiIndex`.\n",
    "A `MultiIndex` can have more than one column; in this example, the first column contains values of `N` and the second column contains values of `p`.\n",
    "\n",
    "The `Pmf` has one row (and one prior probability) for each possible pair of parameters `N` and `p`.\n",
    "So the total number of rows is the product of the lengths of `prior_N` and `prior_p`.\n",
    "\n",
    "Now we have to compute the likelihood of the data for each pair of parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Update\n",
    "\n",
    "To allocate space for the likelihoods, it is convenient to make a copy of `joint_pmf`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.181299Z",
     "iopub.status.busy": "2021-04-16T19:37:48.180508Z",
     "iopub.status.idle": "2021-04-16T19:37:48.182589Z",
     "shell.execute_reply": "2021-04-16T19:37:48.183028Z"
    }
   },
   "outputs": [],
   "source": [
    "likelihood = joint_pmf.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we loop through the pairs of parameters, we compute the likelihood of the data as in the previous section, and then store the result as an element of `likelihood`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:48.187709Z",
     "iopub.status.busy": "2021-04-16T19:37:48.187158Z",
     "iopub.status.idle": "2021-04-16T19:37:53.018364Z",
     "shell.execute_reply": "2021-04-16T19:37:53.017902Z"
    }
   },
   "outputs": [],
   "source": [
    "observed = k01 + k10 + k11\n",
    "\n",
    "for N, p in joint_pmf.index:\n",
    "    k00 = N - observed\n",
    "    x = [k00, k01, k10, k11]\n",
    "    q = 1-p\n",
    "    y = [q*q, q*p, p*q, p*p]\n",
    "    likelihood[N, p] = multinomial.pmf(x, N, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compute the posterior in the usual way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.021855Z",
     "iopub.status.busy": "2021-04-16T19:37:53.021063Z",
     "iopub.status.idle": "2021-04-16T19:37:53.026370Z",
     "shell.execute_reply": "2021-04-16T19:37:53.026762Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_pmf = joint_pmf * likelihood\n",
    "posterior_pmf.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use `plot_contour` again to visualize the joint posterior distribution.\n",
    "But remember that the posterior distribution we just computed is represented as a `Pmf`, which is a `Series`, and `plot_contour` expects a `DataFrame`.\n",
    "\n",
    "Since we used `stack` to convert from a `DataFrame` to a `Series`, we can use `unstack` to go the other way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.029734Z",
     "iopub.status.busy": "2021-04-16T19:37:53.029220Z",
     "iopub.status.idle": "2021-04-16T19:37:53.040079Z",
     "shell.execute_reply": "2021-04-16T19:37:53.039661Z"
    }
   },
   "outputs": [],
   "source": [
    "joint_posterior = posterior_pmf.unstack()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what the result looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.058967Z",
     "iopub.status.busy": "2021-04-16T19:37:53.055472Z",
     "iopub.status.idle": "2021-04-16T19:37:53.229340Z",
     "shell.execute_reply": "2021-04-16T19:37:53.228917Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from utils import plot_contour\n",
    "\n",
    "plot_contour(joint_posterior)\n",
    "\n",
    "decorate(title='Joint posterior distribution of N and p')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most likely values of `N` are near 100, as in the previous model. The most likely values of `p` are near 0.2.\n",
    "\n",
    "The shape of this contour indicates that these parameters are correlated.  If `p` is near the low end of the range, the most likely values of `N` are higher; if `p` is near the high end of the range, `N` is lower. \n",
    "\n",
    "Now that we have a posterior `DataFrame`, we can extract the marginal distributions in the usual way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.232985Z",
     "iopub.status.busy": "2021-04-16T19:37:53.232543Z",
     "iopub.status.idle": "2021-04-16T19:37:53.236313Z",
     "shell.execute_reply": "2021-04-16T19:37:53.235949Z"
    }
   },
   "outputs": [],
   "source": [
    "from utils import marginal\n",
    "\n",
    "posterior2_p = marginal(joint_posterior, 0)\n",
    "posterior2_N = marginal(joint_posterior, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Here's the posterior distribution for `p`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.268190Z",
     "iopub.status.busy": "2021-04-16T19:37:53.257803Z",
     "iopub.status.idle": "2021-04-16T19:37:53.391214Z",
     "shell.execute_reply": "2021-04-16T19:37:53.390708Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior2_p.plot(color='C1')\n",
    "\n",
    "decorate(xlabel='Probability of observing a bear',\n",
    "         ylabel='PDF',\n",
    "         title='Posterior marginal distribution of p')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "The most likely values are near 0.2."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the posterior distribution for `N` based on the two-parameter model, along with the posterior we got using the one-parameter (hypergeometric) model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.473706Z",
     "iopub.status.busy": "2021-04-16T19:37:53.456178Z",
     "iopub.status.idle": "2021-04-16T19:37:53.613688Z",
     "shell.execute_reply": "2021-04-16T19:37:53.613318Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N.plot(label='one-parameter model', color='C4')\n",
    "posterior2_N.plot(label='two-parameter model', color='C1')\n",
    "\n",
    "decorate(xlabel='Population of bears (N)',\n",
    "         ylabel='PDF',\n",
    "         title='Posterior marginal distribution of N')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the two-parameter model, the mean is a little lower and the 90% credible interval is a little narrower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.618046Z",
     "iopub.status.busy": "2021-04-16T19:37:53.617436Z",
     "iopub.status.idle": "2021-04-16T19:37:53.620131Z",
     "shell.execute_reply": "2021-04-16T19:37:53.619764Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(posterior_N.mean(), \n",
    "      posterior_N.credible_interval(0.9))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.624345Z",
     "iopub.status.busy": "2021-04-16T19:37:53.623861Z",
     "iopub.status.idle": "2021-04-16T19:37:53.626345Z",
     "shell.execute_reply": "2021-04-16T19:37:53.625922Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(posterior2_N.mean(), \n",
    "      posterior2_N.credible_interval(0.9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "The two-parameter model yields a narrower posterior distribution for `N`, compared to the one-parameter model, because it takes advantage of an additional source of information: the consistency of the two observations.\n",
    "\n",
    "To see how this helps, consider a scenario where `N` is relatively low, like 138 (the posterior mean of the two-parameter model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.629158Z",
     "iopub.status.busy": "2021-04-16T19:37:53.628703Z",
     "iopub.status.idle": "2021-04-16T19:37:53.630931Z",
     "shell.execute_reply": "2021-04-16T19:37:53.630463Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "N1 = 138"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Given that we saw 23 bears during the first trial and 19 during the second, we can estimate the corresponding value of `p`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.634899Z",
     "iopub.status.busy": "2021-04-16T19:37:53.634158Z",
     "iopub.status.idle": "2021-04-16T19:37:53.637423Z",
     "shell.execute_reply": "2021-04-16T19:37:53.636868Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "mean = (23 + 19) / 2\n",
    "p = mean/N1\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "With these parameters, how much variability do you expect in the number of bears from one trial to the next?  We can quantify that by computing the standard deviation of the binomial distribution with these parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.642568Z",
     "iopub.status.busy": "2021-04-16T19:37:53.641819Z",
     "iopub.status.idle": "2021-04-16T19:37:53.645073Z",
     "shell.execute_reply": "2021-04-16T19:37:53.644629Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from scipy.stats import binom\n",
    "\n",
    "binom(N1, p).std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now let's consider a second scenario where `N` is 173, the posterior mean of the one-parameter model.  The corresponding value of `p` is lower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.649081Z",
     "iopub.status.busy": "2021-04-16T19:37:53.648344Z",
     "iopub.status.idle": "2021-04-16T19:37:53.651511Z",
     "shell.execute_reply": "2021-04-16T19:37:53.650985Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "N2 = 173\n",
    "p = mean/N2\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "In this scenario, the variation we expect to see from one trial to the next is higher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.656928Z",
     "iopub.status.busy": "2021-04-16T19:37:53.656278Z",
     "iopub.status.idle": "2021-04-16T19:37:53.658960Z",
     "shell.execute_reply": "2021-04-16T19:37:53.659323Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "binom(N2, p).std()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "So if the number of bears we observe is the same in both trials, that would be evidence for lower values of `N`, where we expect more consistency.\n",
    "If the number of bears is substantially different between the two trials, that would be evidence for higher values of `N`.\n",
    "\n",
    "In the actual data, the difference between the two trials is low, which is why the posterior mean of the two-parameter model is lower.\n",
    "The two-parameter model takes advantage of additional information, which is why the credible interval is narrower."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Joint and Marginal Distributions\n",
    "\n",
    "Marginal distributions are called \"marginal\" because in a common visualization they appear in the margins of the plot.\n",
    "\n",
    "Seaborn provides a class called `JointGrid` that creates this visualization.\n",
    "The following function uses it to show the joint and marginal distributions in a single plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.664695Z",
     "iopub.status.busy": "2021-04-16T19:37:53.664254Z",
     "iopub.status.idle": "2021-04-16T19:37:53.667740Z",
     "shell.execute_reply": "2021-04-16T19:37:53.667301Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from seaborn import JointGrid\n",
    "\n",
    "def joint_plot(joint, **options):\n",
    "    \"\"\"Show joint and marginal distributions.\n",
    "    \n",
    "    joint: DataFrame that represents a joint distribution\n",
    "    options: passed to JointGrid\n",
    "    \"\"\"\n",
    "    # get the names of the parameters\n",
    "    x = joint.columns.name\n",
    "    x = 'x' if x is None else x\n",
    "\n",
    "    y = joint.index.name\n",
    "    y = 'y' if y is None else y\n",
    "\n",
    "    # make a JointGrid with minimal data\n",
    "    data = pd.DataFrame({x:[0], y:[0]})\n",
    "    g = JointGrid(x=x, y=y, data=data, **options)\n",
    "\n",
    "    # replace the contour plot\n",
    "    g.ax_joint.contour(joint.columns, \n",
    "                       joint.index, \n",
    "                       joint, \n",
    "                       cmap='viridis')\n",
    "    \n",
    "    # replace the marginals\n",
    "    marginal_x = marginal(joint, 0)\n",
    "    g.ax_marg_x.plot(marginal_x.qs, marginal_x.ps)\n",
    "    \n",
    "    marginal_y = marginal(joint, 1)\n",
    "    g.ax_marg_y.plot(marginal_y.ps, marginal_y.qs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.710725Z",
     "iopub.status.busy": "2021-04-16T19:37:53.699575Z",
     "iopub.status.idle": "2021-04-16T19:37:53.974546Z",
     "shell.execute_reply": "2021-04-16T19:37:53.974124Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "joint_plot(joint_posterior)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "A `JointGrid` is a concise way to represent the joint and marginal distributions visually."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Lincoln Index Problem\n",
    "\n",
    "In [an excellent blog post](http://www.johndcook.com/blog/2010/07/13/lincoln-index/), John D. Cook wrote about the Lincoln index, which is a way to estimate the\n",
    "number of errors in a document (or program) by comparing results from\n",
    "two independent testers.\n",
    "Here's his presentation of the problem:\n",
    "\n",
    "> \"Suppose you have a tester who finds 20 bugs in your program. You\n",
    "> want to estimate how many bugs are really in the program. You know\n",
    "> there are at least 20 bugs, and if you have supreme confidence in your\n",
    "> tester, you may suppose there are around 20 bugs. But maybe your\n",
    "> tester isn't very good. Maybe there are hundreds of bugs. How can you\n",
    "> have any idea how many bugs there are? There's no way to know with one\n",
    "> tester. But if you have two testers, you can get a good idea, even if\n",
    "> you don't know how skilled the testers are.\"\n",
    "\n",
    "Suppose the first tester finds 20 bugs, the second finds 15, and they\n",
    "find 3 in common; how can we estimate the number of bugs?\n",
    "\n",
    "This problem is similar to the Grizzly Bear problem, so I'll represent the data in the same way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.977529Z",
     "iopub.status.busy": "2021-04-16T19:37:53.977079Z",
     "iopub.status.idle": "2021-04-16T19:37:53.979111Z",
     "shell.execute_reply": "2021-04-16T19:37:53.978757Z"
    }
   },
   "outputs": [],
   "source": [
    "k10 = 20 - 3\n",
    "k01 = 15 - 3\n",
    "k11 = 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But in this case it is probably not reasonable to assume that the testers have the same probability of finding a bug.\n",
    "So I'll define two parameters, `p0` for the probability that the first tester finds a bug, and `p1` for the probability that the second tester finds a bug.\n",
    "\n",
    "I will continue to assume that the probabilities are independent, which is like assuming that all bugs are equally easy to find.  That might not be a good assumption, but let's stick with it for now.\n",
    "\n",
    "As an example, suppose we know that the probabilities are 0.2 and 0.15."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.982038Z",
     "iopub.status.busy": "2021-04-16T19:37:53.981602Z",
     "iopub.status.idle": "2021-04-16T19:37:53.984284Z",
     "shell.execute_reply": "2021-04-16T19:37:53.983762Z"
    }
   },
   "outputs": [],
   "source": [
    "p0, p1 = 0.2, 0.15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the array of probabilities, `y`, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.988890Z",
     "iopub.status.busy": "2021-04-16T19:37:53.988165Z",
     "iopub.status.idle": "2021-04-16T19:37:53.990550Z",
     "shell.execute_reply": "2021-04-16T19:37:53.990104Z"
    }
   },
   "outputs": [],
   "source": [
    "def compute_probs(p0, p1):\n",
    "    \"\"\"Computes the probability for each of 4 categories.\"\"\"\n",
    "    q0 = 1-p0\n",
    "    q1 = 1-p1\n",
    "    return [q0*q1, q0*p1, p0*q1, p0*p1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:53.994239Z",
     "iopub.status.busy": "2021-04-16T19:37:53.993623Z",
     "iopub.status.idle": "2021-04-16T19:37:53.997022Z",
     "shell.execute_reply": "2021-04-16T19:37:53.996576Z"
    }
   },
   "outputs": [],
   "source": [
    "y = compute_probs(p0, p1)\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With these probabilities, there is a \n",
    "68% chance that neither tester finds the bug and a\n",
    "3% chance that both do. \n",
    "\n",
    "Pretending that these probabilities are known, we can compute the posterior distribution for `N`.\n",
    "Here's a prior distribution that's uniform from 32 to 350 bugs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.005383Z",
     "iopub.status.busy": "2021-04-16T19:37:54.004613Z",
     "iopub.status.idle": "2021-04-16T19:37:54.008366Z",
     "shell.execute_reply": "2021-04-16T19:37:54.007866Z"
    }
   },
   "outputs": [],
   "source": [
    "qs = np.arange(32, 350, step=5) \n",
    "prior_N = make_uniform(qs, name='N')\n",
    "prior_N.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I'll put the data in an array, with 0 as a place-keeper for the unknown value `k00`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.012337Z",
     "iopub.status.busy": "2021-04-16T19:37:54.011753Z",
     "iopub.status.idle": "2021-04-16T19:37:54.014357Z",
     "shell.execute_reply": "2021-04-16T19:37:54.013863Z"
    }
   },
   "outputs": [],
   "source": [
    "data = np.array([0, k01, k10, k11])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here are the likelihoods for each value of `N`, with `ps` as a constant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.024013Z",
     "iopub.status.busy": "2021-04-16T19:37:54.023587Z",
     "iopub.status.idle": "2021-04-16T19:37:54.025396Z",
     "shell.execute_reply": "2021-04-16T19:37:54.025737Z"
    }
   },
   "outputs": [],
   "source": [
    "likelihood = prior_N.copy()\n",
    "observed = data.sum()\n",
    "x = data.copy()\n",
    "\n",
    "for N in prior_N.qs:\n",
    "    x[0] = N - observed\n",
    "    likelihood[N] = multinomial.pmf(x, N, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the posterior in the usual way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.029826Z",
     "iopub.status.busy": "2021-04-16T19:37:54.029301Z",
     "iopub.status.idle": "2021-04-16T19:37:54.031611Z",
     "shell.execute_reply": "2021-04-16T19:37:54.032006Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N = prior_N * likelihood\n",
    "posterior_N.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.066285Z",
     "iopub.status.busy": "2021-04-16T19:37:54.054991Z",
     "iopub.status.idle": "2021-04-16T19:37:54.200634Z",
     "shell.execute_reply": "2021-04-16T19:37:54.200132Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_N.plot(color='C4')\n",
    "\n",
    "decorate(xlabel='Number of bugs (N)',\n",
    "         ylabel='PMF',\n",
    "         title='Posterior marginal distribution of n with known p1, p2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.204715Z",
     "iopub.status.busy": "2021-04-16T19:37:54.204282Z",
     "iopub.status.idle": "2021-04-16T19:37:54.206388Z",
     "shell.execute_reply": "2021-04-16T19:37:54.206734Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(posterior_N.mean(), \n",
    "      posterior_N.credible_interval(0.9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the assumption that `p0` and `p1` are known to be `0.2` and `0.15`, the posterior mean is 102 with 90% credible interval (77, 127).\n",
    "But this result is based on the assumption that we know the probabilities, and we don't."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Three-Parameter Model\n",
    "\n",
    "What we need is a model with three parameters: `N`, `p0`, and `p1`.\n",
    "We'll use `prior_N` again for the prior distribution of `N`, and here are the priors for `p0` and `p1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.211539Z",
     "iopub.status.busy": "2021-04-16T19:37:54.211111Z",
     "iopub.status.idle": "2021-04-16T19:37:54.213388Z",
     "shell.execute_reply": "2021-04-16T19:37:54.212835Z"
    }
   },
   "outputs": [],
   "source": [
    "qs = np.linspace(0, 1, num=51)\n",
    "prior_p0 = make_uniform(qs, name='p0')\n",
    "prior_p1 = make_uniform(qs, name='p1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have to assemble them into a joint prior with three dimensions.\n",
    "I'll start by putting the first two into a `DataFrame`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.217788Z",
     "iopub.status.busy": "2021-04-16T19:37:54.217207Z",
     "iopub.status.idle": "2021-04-16T19:37:54.219947Z",
     "shell.execute_reply": "2021-04-16T19:37:54.220378Z"
    }
   },
   "outputs": [],
   "source": [
    "joint2 = make_joint(prior_p0, prior_N)\n",
    "joint2.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now I'll stack them, as in the previous example, and put the result in a `Pmf`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.225352Z",
     "iopub.status.busy": "2021-04-16T19:37:54.224407Z",
     "iopub.status.idle": "2021-04-16T19:37:54.234609Z",
     "shell.execute_reply": "2021-04-16T19:37:54.233999Z"
    }
   },
   "outputs": [],
   "source": [
    "joint2_pmf = Pmf(joint2.stack())\n",
    "joint2_pmf.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use `make_joint` again to add in the third parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.238175Z",
     "iopub.status.busy": "2021-04-16T19:37:54.237718Z",
     "iopub.status.idle": "2021-04-16T19:37:54.242307Z",
     "shell.execute_reply": "2021-04-16T19:37:54.242651Z"
    }
   },
   "outputs": [],
   "source": [
    "joint3 = make_joint(prior_p1, joint2_pmf)\n",
    "joint3.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a `DataFrame` with values of `N` and `p0` in a `MultiIndex` that goes down the rows and values of `p1` in an index that goes across the columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.247118Z",
     "iopub.status.busy": "2021-04-16T19:37:54.246694Z",
     "iopub.status.idle": "2021-04-16T19:37:54.266251Z",
     "shell.execute_reply": "2021-04-16T19:37:54.265904Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "joint3.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now I'll apply `stack` again:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.271012Z",
     "iopub.status.busy": "2021-04-16T19:37:54.270430Z",
     "iopub.status.idle": "2021-04-16T19:37:54.281206Z",
     "shell.execute_reply": "2021-04-16T19:37:54.280739Z"
    }
   },
   "outputs": [],
   "source": [
    "joint3_pmf = Pmf(joint3.stack())\n",
    "joint3_pmf.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a `Pmf` with a three-column `MultiIndex` containing all possible triplets of parameters.\n",
    "\n",
    "The number of rows is the product of the number of values in all three priors, which is almost 170,000."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.284532Z",
     "iopub.status.busy": "2021-04-16T19:37:54.283882Z",
     "iopub.status.idle": "2021-04-16T19:37:54.286757Z",
     "shell.execute_reply": "2021-04-16T19:37:54.286325Z"
    }
   },
   "outputs": [],
   "source": [
    "joint3_pmf.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's still small enough to be practical, but it will take longer to compute the likelihoods than in the previous examples.\n",
    "\n",
    "Here's the loop that computes the likelihoods; it's similar to the one in the previous section:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:37:54.291184Z",
     "iopub.status.busy": "2021-04-16T19:37:54.290454Z",
     "iopub.status.idle": "2021-04-16T19:38:12.997223Z",
     "shell.execute_reply": "2021-04-16T19:38:12.996695Z"
    }
   },
   "outputs": [],
   "source": [
    "likelihood = joint3_pmf.copy()\n",
    "observed = data.sum()\n",
    "x = data.copy()\n",
    "\n",
    "for N, p0, p1 in joint3_pmf.index:\n",
    "    x[0] = N - observed\n",
    "    y = compute_probs(p0, p1)\n",
    "    likelihood[N, p0, p1] = multinomial.pmf(x, N, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compute the posterior in the usual way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.000759Z",
     "iopub.status.busy": "2021-04-16T19:38:13.000131Z",
     "iopub.status.idle": "2021-04-16T19:38:13.006898Z",
     "shell.execute_reply": "2021-04-16T19:38:13.006422Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_pmf = joint3_pmf * likelihood\n",
    "posterior_pmf.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, to extract the marginal distributions, we could unstack the joint posterior as we did in the previous section.\n",
    "But `Pmf` provides a version of `marginal` that works with a `Pmf` rather than a `DataFrame`.\n",
    "Here's how we use it to get the posterior distribution for `N`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.010138Z",
     "iopub.status.busy": "2021-04-16T19:38:13.009641Z",
     "iopub.status.idle": "2021-04-16T19:38:13.015759Z",
     "shell.execute_reply": "2021-04-16T19:38:13.015325Z"
    }
   },
   "outputs": [],
   "source": [
    "posterior_N = posterior_pmf.marginal(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.040703Z",
     "iopub.status.busy": "2021-04-16T19:38:13.030126Z",
     "iopub.status.idle": "2021-04-16T19:38:13.177605Z",
     "shell.execute_reply": "2021-04-16T19:38:13.177145Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_N.plot(color='C4')\n",
    "\n",
    "decorate(xlabel='Number of bugs (N)',\n",
    "         ylabel='PDF',\n",
    "         title='Posterior marginal distributions of N')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.181343Z",
     "iopub.status.busy": "2021-04-16T19:38:13.180742Z",
     "iopub.status.idle": "2021-04-16T19:38:13.183133Z",
     "shell.execute_reply": "2021-04-16T19:38:13.183477Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_N.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The posterior mean is 105 bugs, which suggests that there are still many bugs the testers have not found.\n",
    "\n",
    "Here are the posteriors for `p0` and `p1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.187388Z",
     "iopub.status.busy": "2021-04-16T19:38:13.186841Z",
     "iopub.status.idle": "2021-04-16T19:38:13.340072Z",
     "shell.execute_reply": "2021-04-16T19:38:13.339679Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_p1 = posterior_pmf.marginal(1)\n",
    "posterior_p2 = posterior_pmf.marginal(2)\n",
    "\n",
    "posterior_p1.plot(label='p1')\n",
    "posterior_p2.plot(label='p2')\n",
    "\n",
    "decorate(xlabel='Probability of finding a bug',\n",
    "         ylabel='PDF',\n",
    "         title='Posterior marginal distributions of p1 and p2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.344502Z",
     "iopub.status.busy": "2021-04-16T19:38:13.344075Z",
     "iopub.status.idle": "2021-04-16T19:38:13.348520Z",
     "shell.execute_reply": "2021-04-16T19:38:13.348155Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_p1.mean(), posterior_p1.credible_interval(0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.352681Z",
     "iopub.status.busy": "2021-04-16T19:38:13.352235Z",
     "iopub.status.idle": "2021-04-16T19:38:13.356101Z",
     "shell.execute_reply": "2021-04-16T19:38:13.355713Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "posterior_p2.mean(), posterior_p2.credible_interval(0.9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Comparing the posterior distributions, the tester who found more bugs probably has a higher probability of finding bugs.  The posterior means are about 23% and 18%.  But the distributions overlap, so we should not be too sure."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the first example we've seen with three parameters.\n",
    "As the number of parameters increases, the number of combinations increases quickly.\n",
    "The method we've been using so far, enumerating all possible combinations, becomes impractical if the number of parameters is more than 3 or 4.\n",
    "\n",
    "However there are other methods that can handle models with many more parameters, as we'll see in <<_MCMC>>."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "The problems in this chapter are examples of [mark and recapture](https://en.wikipedia.org/wiki/Mark_and_recapture) experiments, which are used in ecology to estimate animal populations.  They also have applications in engineering, as in the Lincoln index problem.  And in the exercises you'll see that they are used in epidemiology, too.\n",
    "\n",
    "This chapter introduces two new probability distributions:\n",
    "\n",
    "* The hypergeometric distribution is a variation of the binomial distribution in which samples are drawn from the population without replacement. \n",
    "\n",
    "* The multinomial distribution is a generalization of the binomial distribution where there are more than two possible outcomes.\n",
    "\n",
    "Also in this chapter, we saw the first example of a model with three parameters.  We'll see more in subsequent chapters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** [In an excellent paper](http://chao.stat.nthu.edu.tw/wordpress/paper/110.pdf), Anne Chao explains how mark and recapture experiments are used in epidemiology to estimate the prevalence of a disease in a human population based on multiple incomplete lists of cases.\n",
    "\n",
    "One of the examples in that paper is a study \"to estimate the number of people who were infected by hepatitis in an outbreak that occurred in and around a college in northern Taiwan from April to July 1995.\"\n",
    "\n",
    "Three lists of cases were available:\n",
    "\n",
    "1. 135 cases identified using a serum test. \n",
    "\n",
    "2. 122 cases reported by local hospitals. \n",
    "\n",
    "3. 126 cases reported on questionnaires collected by epidemiologists.\n",
    "\n",
    "In this exercise, we'll use only the first two lists; in the next exercise we'll bring in the third list.\n",
    "\n",
    "Make a joint prior and update it using this data, then compute the posterior mean of `N` and a 90% credible interval."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "The following array contains 0 as a place-holder for the unknown value of `k00`, followed by known values of `k01`, `k10`, and `k11`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.359215Z",
     "iopub.status.busy": "2021-04-16T19:38:13.358791Z",
     "iopub.status.idle": "2021-04-16T19:38:13.360380Z",
     "shell.execute_reply": "2021-04-16T19:38:13.360718Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data2 = np.array([0, 73, 86, 49])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "These data indicate that there are 73 cases on the second list that are not on the first, 86 cases on the first list that are not on the second, and 49 cases on both lists.\n",
    "\n",
    "To keep things simple, we'll assume that each case has the same probability of appearing on each list.  So we'll use a two-parameter model where `N` is the total number of cases and `p` is the probability that any case appears on any list.\n",
    "\n",
    "Here are priors you can start with (but feel free to modify them)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.367733Z",
     "iopub.status.busy": "2021-04-16T19:38:13.367161Z",
     "iopub.status.idle": "2021-04-16T19:38:13.369591Z",
     "shell.execute_reply": "2021-04-16T19:38:13.369935Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "qs = np.arange(200, 500, step=5)\n",
    "prior_N = make_uniform(qs, name='N')\n",
    "prior_N.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.377139Z",
     "iopub.status.busy": "2021-04-16T19:38:13.376572Z",
     "iopub.status.idle": "2021-04-16T19:38:13.378957Z",
     "shell.execute_reply": "2021-04-16T19:38:13.379302Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "qs = np.linspace(0, 0.98, num=50)\n",
    "prior_p = make_uniform(qs, name='p')\n",
    "prior_p.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.397988Z",
     "iopub.status.busy": "2021-04-16T19:38:13.397393Z",
     "iopub.status.idle": "2021-04-16T19:38:13.400383Z",
     "shell.execute_reply": "2021-04-16T19:38:13.400013Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.407934Z",
     "iopub.status.busy": "2021-04-16T19:38:13.403583Z",
     "iopub.status.idle": "2021-04-16T19:38:13.410792Z",
     "shell.execute_reply": "2021-04-16T19:38:13.410389Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.415288Z",
     "iopub.status.busy": "2021-04-16T19:38:13.414841Z",
     "iopub.status.idle": "2021-04-16T19:38:13.994497Z",
     "shell.execute_reply": "2021-04-16T19:38:13.994870Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:13.998580Z",
     "iopub.status.busy": "2021-04-16T19:38:13.997855Z",
     "iopub.status.idle": "2021-04-16T19:38:14.001960Z",
     "shell.execute_reply": "2021-04-16T19:38:14.001563Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.004993Z",
     "iopub.status.busy": "2021-04-16T19:38:14.004419Z",
     "iopub.status.idle": "2021-04-16T19:38:14.008280Z",
     "shell.execute_reply": "2021-04-16T19:38:14.008628Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.027577Z",
     "iopub.status.busy": "2021-04-16T19:38:14.024250Z",
     "iopub.status.idle": "2021-04-16T19:38:14.153324Z",
     "shell.execute_reply": "2021-04-16T19:38:14.152812Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.173753Z",
     "iopub.status.busy": "2021-04-16T19:38:14.165487Z",
     "iopub.status.idle": "2021-04-16T19:38:14.299631Z",
     "shell.execute_reply": "2021-04-16T19:38:14.299087Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.304726Z",
     "iopub.status.busy": "2021-04-16T19:38:14.304159Z",
     "iopub.status.idle": "2021-04-16T19:38:14.309512Z",
     "shell.execute_reply": "2021-04-16T19:38:14.308987Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Now let's do the version of the problem with all three lists.  Here's the data from Chou's paper:\n",
    "\n",
    "```\n",
    "Hepatitis A virus list\n",
    "P    Q    E    Data\n",
    "1    1    1    k111 =28\n",
    "1    1    0    k110 =21\n",
    "1    0    1    k101 =17\n",
    "1    0    0    k100 =69\n",
    "0    1    1    k011 =18\n",
    "0    1    0    k010 =55\n",
    "0    0    1    k001 =63\n",
    "0    0    0    k000 =??\n",
    "```\n",
    "\n",
    "Write a loop that computes the likelihood of the data for each pair of parameters, then update the prior and compute the posterior mean of `N`.  How does it compare to the results using only the first two lists?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Here's the data in a NumPy array (in reverse order)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.313377Z",
     "iopub.status.busy": "2021-04-16T19:38:14.312777Z",
     "iopub.status.idle": "2021-04-16T19:38:14.316098Z",
     "shell.execute_reply": "2021-04-16T19:38:14.315626Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data3 = np.array([0, 63, 55, 18, 69, 17, 21, 28])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Again, the first value is a place-keeper for the unknown `k000`.  The second value is `k001`, which means there are 63 cases that appear on the third list but not the first two.  And the last value is `k111`, which means there are 28 cases that appear on all three lists.\n",
    "\n",
    "In the two-list version of the problem we computed `ps` by enumerating the combinations of `p` and `q`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.319095Z",
     "iopub.status.busy": "2021-04-16T19:38:14.318115Z",
     "iopub.status.idle": "2021-04-16T19:38:14.322989Z",
     "shell.execute_reply": "2021-04-16T19:38:14.322543Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "q = 1-p\n",
    "ps = [q*q, q*p, p*q, p*p]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "We could do the same thing for the three-list version, computing the probability for each of the eight categories.  But we can generalize it by recognizing that we are computing the cartesian product of `p` and `q`, repeated once for each list.\n",
    "\n",
    "And we can use the following function (based on [this StackOverflow answer](https://stackoverflow.com/questions/58242078/cartesian-product-of-arbitrary-lists-in-pandas/58242079#58242079)) to compute Cartesian products:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.326320Z",
     "iopub.status.busy": "2021-04-16T19:38:14.325903Z",
     "iopub.status.idle": "2021-04-16T19:38:14.327969Z",
     "shell.execute_reply": "2021-04-16T19:38:14.327558Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def cartesian_product(*args, **options):\n",
    "    \"\"\"Cartesian product of sequences.\n",
    "    \n",
    "    args: any number of sequences\n",
    "    options: passes to `MultiIndex.from_product`\n",
    "    \n",
    "    returns: DataFrame with one column per sequence\n",
    "    \"\"\"\n",
    "    index = pd.MultiIndex.from_product(args, **options)\n",
    "    return pd.DataFrame(index=index).reset_index()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Here's an example with `p=0.2`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.338979Z",
     "iopub.status.busy": "2021-04-16T19:38:14.332076Z",
     "iopub.status.idle": "2021-04-16T19:38:14.342183Z",
     "shell.execute_reply": "2021-04-16T19:38:14.341674Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "p = 0.2\n",
    "t = (1-p, p)\n",
    "df = cartesian_product(t, t, t)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "To compute the probability for each category, we take the product across the columns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.347732Z",
     "iopub.status.busy": "2021-04-16T19:38:14.347185Z",
     "iopub.status.idle": "2021-04-16T19:38:14.349713Z",
     "shell.execute_reply": "2021-04-16T19:38:14.350099Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "y = df.prod(axis=1)\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now you finish it off from there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:14.358021Z",
     "iopub.status.busy": "2021-04-16T19:38:14.355220Z",
     "iopub.status.idle": "2021-04-16T19:38:23.932525Z",
     "shell.execute_reply": "2021-04-16T19:38:23.932159Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:23.936587Z",
     "iopub.status.busy": "2021-04-16T19:38:23.935690Z",
     "iopub.status.idle": "2021-04-16T19:38:23.938938Z",
     "shell.execute_reply": "2021-04-16T19:38:23.939283Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:23.942623Z",
     "iopub.status.busy": "2021-04-16T19:38:23.941916Z",
     "iopub.status.idle": "2021-04-16T19:38:23.945820Z",
     "shell.execute_reply": "2021-04-16T19:38:23.945383Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:23.963282Z",
     "iopub.status.busy": "2021-04-16T19:38:23.962826Z",
     "iopub.status.idle": "2021-04-16T19:38:24.071099Z",
     "shell.execute_reply": "2021-04-16T19:38:24.070753Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:24.075027Z",
     "iopub.status.busy": "2021-04-16T19:38:24.074561Z",
     "iopub.status.idle": "2021-04-16T19:38:24.077867Z",
     "shell.execute_reply": "2021-04-16T19:38:24.077390Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:24.102710Z",
     "iopub.status.busy": "2021-04-16T19:38:24.093396Z",
     "iopub.status.idle": "2021-04-16T19:38:24.229139Z",
     "shell.execute_reply": "2021-04-16T19:38:24.229631Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:24.233828Z",
     "iopub.status.busy": "2021-04-16T19:38:24.233412Z",
     "iopub.status.idle": "2021-04-16T19:38:24.237539Z",
     "shell.execute_reply": "2021-04-16T19:38:24.238030Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-04-16T19:38:24.242140Z",
     "iopub.status.busy": "2021-04-16T19:38:24.241727Z",
     "iopub.status.idle": "2021-04-16T19:38:24.245735Z",
     "shell.execute_reply": "2021-04-16T19:38:24.245348Z"
    }
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}