elephants_soln.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Think Bayes\n",
    "\n",
    "Copyright 2018 Allen B. Downey\n",
    "\n",
    "MIT License: https://opensource.org/licenses/MIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# import classes from thinkbayes2\n",
    "from thinkbayes2 import Pmf, Cdf, Suite, Joint\n",
    "\n",
    "from thinkbayes2 import MakePoissonPmf, EvalBinomialPmf, MakeMixture\n",
    "\n",
    "import thinkplot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cats and rats and elephants\n",
    "\n",
    "Suppose there are six species that might be in a zoo: lions and tigers and bears, and cats and rats and elephants.  Every zoo has a subset of these species, and every subset is equally likely.\n",
    "\n",
    "One day we visit a zoo and see 3 lions, 2 tigers, and one bear.  Assuming that every animal in the zoo has an equal chance to be seen, what is the probability that the next animal we see is an elephant?\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solution\n",
    "\n",
    "I'll start by enumerating all possible zoos with `itertools`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import combinations\n",
    "\n",
    "def power_set(s):\n",
    "    n = len(s)\n",
    "    for r in range(1, n+1):\n",
    "        for combo in combinations(s, r):\n",
    "            yield ''.join(combo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can enumerate only the zoos that are possible, given a set of animals known to be present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enumerate_zoos(all_species, present):\n",
    "    \"\"\"Enumerate all zoos that contain `present`.\n",
    "    \n",
    "    all_species: sequence of all species\n",
    "    present: sequence of species present\n",
    "    \n",
    "    yields: possible zoos\n",
    "    \"\"\"\n",
    "    present = set(present)\n",
    "    for combo in power_set(species):\n",
    "        intersect = set(combo) & present\n",
    "        if len(intersect) == len(present):\n",
    "            yield len(combo), combo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the possible zoos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3 LTB\n",
      "4 LTBC\n",
      "4 LTBR\n",
      "4 LTBE\n",
      "5 LTBCR\n",
      "5 LTBCE\n",
      "5 LTBRE\n",
      "6 LTBCRE\n"
     ]
    }
   ],
   "source": [
    "species = 'LTBCRE'\n",
    "present = 'LTB'\n",
    "\n",
    "for n, zoo in enumerate_zoos(species, present):\n",
    "    print(n, zoo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To represent the prior and posterior distributions I'll use a hierarchical model with one Dirichlet object for each possible zoo.\n",
    "\n",
    "At the bottom of the hierarchy, it is easy to update each Dirichlet object just by adding the observed frequencies to the parameters.\n",
    "\n",
    "In order to update the top of the hierarchy, we need the total probability of the data for each hypothetical zoo.  When we do an update using grid algorithms, we get the probability of the data free, since it is the normalizing constant.\n",
    "\n",
    "But when we do an update using a conjugate distribution, we don't get the total probability of the data, and for a Dirichlet distribution it is not easy to compute.\n",
    "\n",
    "However, we can estimate it by drawing samples from the Dirichlet distribution, and then computing the probability of the data for each sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dirichlet(object):\n",
    "    \"\"\"Represents a Dirichlet distribution.\n",
    "\n",
    "    See http://en.wikipedia.org/wiki/Dirichlet_distribution\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n, conc=1, label=None):\n",
    "        \"\"\"Initializes a Dirichlet distribution.\n",
    "\n",
    "        n: number of dimensions\n",
    "        conc: concentration parameter\n",
    "        label: string label\n",
    "        \"\"\"\n",
    "        if n < 2:\n",
    "            raise ValueError('A Dirichlet distribution with '\n",
    "                             'n<2 makes no sense')\n",
    "\n",
    "        self.n = n\n",
    "        self.params = np.ones(n, dtype=np.float) * conc\n",
    "        self.label = label\n",
    "\n",
    "    def update(self, data):\n",
    "        \"\"\"Updates a Dirichlet distribution.\n",
    "\n",
    "        data: sequence of observations\n",
    "        \"\"\"\n",
    "        m = len(data)\n",
    "        self.params[:m] += data\n",
    "\n",
    "    def random(self):\n",
    "        \"\"\"Generates a random variate from this distribution.\n",
    "\n",
    "        Returns: normalized vector of fractions\n",
    "        \"\"\"\n",
    "        p = np.random.gamma(self.params)\n",
    "        return p / p.sum()\n",
    "\n",
    "    def mean(self):\n",
    "        \"\"\"Array of means.\"\"\"\n",
    "        return self.params / self.params.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example that represents a zoo with 4 animals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.Dirichlet at 0x7f209e113278>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d4 = Dirichlet(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's a sample from it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.04846145, 0.59748328, 0.30808513, 0.04597014])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p = d4.random()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compute the probability of the data, given these prevalences, using the multinomial distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0007510394590242562"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from scipy.stats import multinomial\n",
    "\n",
    "data = [3, 2, 1, 0]\n",
    "m = sum(data)\n",
    "multinomial(m, p).pmf(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since I only observed 3 species, and my hypothetical zoo has 4, I had to zero-pad the data.  Here's a function that makes that easier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zero_pad(a, n):\n",
    "    \"\"\"Why does np.pad have to be so complicated?\n",
    "    \"\"\"\n",
    "    res = np.zeros(n)\n",
    "    res[:len(a)] = a\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3., 2., 1., 0.])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = [3, 2, 1]\n",
    "zero_pad(data, 4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's pull all that together.  Here's a function that estimates the total probability of the data by sampling from the dirichlet distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_likelihood(dirichlet, data, iters=1000):\n",
    "    \"\"\"Estimate the total probability of the data.\n",
    "    \n",
    "    dirichlet: Dirichlet object\n",
    "    data: array of observed frequencies\n",
    "    iters: number of samples to draw\n",
    "    \"\"\"\n",
    "    data = zero_pad(data, dirichlet.n)\n",
    "    m = np.sum(data)\n",
    "    likes = [multinomial(m, dirichlet.random()).pmf(data) \n",
    "             for i in range(iters)]\n",
    "    return np.mean(likes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.011496028801106875"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_likelihood(d4, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to solve the problem.\n",
    "\n",
    "Here's a Suite that represents the set of possible zoos.  The likelihood of any zoo is just the total probability of the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Zoo(Suite):\n",
    "    \n",
    "    def Likelihood(self, data, hypo):\n",
    "        \"\"\"\n",
    "        data: sequence of counts\n",
    "        hypo: Dirichlet object\n",
    "        \"\"\"\n",
    "        return sample_likelihood(hypo, data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can construct the prior by enumerating the possible zoos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "suite = Zoo([Dirichlet(n, label=''.join(zoo))\n",
    "             for n, zoo in enumerate_zoos(species, present)]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.125 LTB\n",
      "0.125 LTBC\n",
      "0.125 LTBR\n",
      "0.125 LTBE\n",
      "0.125 LTBCR\n",
      "0.125 LTBCE\n",
      "0.125 LTBRE\n",
      "0.125 LTBCRE\n"
     ]
    }
   ],
   "source": [
    "def print_zoos(suite):\n",
    "    for d, p in suite.Items():\n",
    "        print(p, d.label)\n",
    "        \n",
    "print_zoos(suite)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can update the top level of the hierarchy by calling `Update`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.011250586599734533"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "suite.Update(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have to update the bottom level explicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "for hypo in suite:\n",
    "    hypo.update(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the posterior for the top level."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.43679617920687586 LTB\n",
      "0.12516851264765871 LTBC\n",
      "0.1304178179372656 LTBR\n",
      "0.12616846874234114 LTBE\n",
      "0.048267930148971114 LTBCR\n",
      "0.056066029260621264 LTBCE\n",
      "0.052742406633655016 LTBRE\n",
      "0.024372655422611408 LTBCRE\n"
     ]
    }
   ],
   "source": [
    "print_zoos(suite)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we can get the posterior distribution of `n`, the number of species."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "pmf_n = Pmf()\n",
    "for d, p in suite.Items():\n",
    "    pmf_n[d.n] += p"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.769025497681595\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAagAAAEYCAYAAAAJeGK1AAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAAFS9JREFUeJzt3X20XXV95/H3xwSqBYQZyVQhgWChzkTq06SAY4dxLHWBD8GZ0kparTo61I5RZtnVFlsWKoPtqFMtCn1AsaU+gJRqG2kUbS2doV1QgmXUgIyZCE0araHIk1gx+J0/zr6Zw/Hknpvk7nt/J3m/1ror++G39/6en7I/dz/c30lVIUlSax6z2AVIkjSOASVJapIBJUlqkgElSWqSASVJapIBJUlqkgGlA1qSB5M8uad9PzfJtqH5TUmeO0/7/pkknx6aryTHz8e+u/311i+zHPNxST6R5L4kf7iQx1abDCgtiCR3JvlWd+L7hyS/l+TQfdjfyu6kvHRf6qqqQ6tqy77sYw+O9dSqun62NnP9XFX14ap6/nzUleT6JK8Z2f+C9cuQs4AfAJ5QVT+5wMdWgwwoLaQXV9WhwLOAHwHOX6xC9jXY9nX7aT12z44F/k9V7VzsQtQGA0oLrqr+HvgkcCJAkqOSrE9yT5LNSf7zTNskJyXZmOT+7srrXd2q/9n9e293Vfbsrv1/SnJ7km8kuS7JsUP7qiSvS/Jl4MtDy47vpg9P8gdJdiS5K8n5SR7TrXtlkr9K8u4k9wBvGf1c3S2q3++OfRuDEB5ef2eS0/b0c407drfshpESXpBkS5K7k7xzqPa3JPnQUB27rtKSvA34t8Al3fEu2Yt+uSHJ/+g+91eSnLG7/+2T/Kvuiu3e7pbnmm75W4ELgJd2dbx6zLZvSXJ1V8sD3fard3csTb/99TcxNSzJCuAFwMe6RVcCm4CjgH8JfCbJlqr6c+Bi4OKq+mB3S/DEbptTga8AR8z8xp3kJcCvAC9mEEDndfv+N0OHfwlwMvCtMaW9FzgceDLwBODTwFeBy7v1JwNXAf8COGjM9m8GfrD7OYRBCO/Onnyup4w59kvH7PM/AKuBQ4E/A+4A3j9LDVTVryZ5DvChqtpd27n0yxXAkcA5wOVJjq6RcdSSHAR8AvgA8HzgR4E/SbK6qt6cpIDjq+pls5S8BviPwKuAi4BLgFNm+4yaXl5BaSH9cZJ7gRuAvwR+rQurHwV+uar+qapuZXBSfXm3zXeA45McWVUPVtWNs+z/54Bfr6rbu5P7rwHPGL6K6tbfU1WPCqgkSxic9N9UVQ9U1Z3AbwzVAbC9qt5bVTtHt+/8FPC2bv9bgffMUuuefK65HBvg7d2x/w74TWDthH1ONMd+uauq3ldVjzAIqicxeJY06hQG4fnfq+rhqvoscO0e1nlDVW3ojvVB4Ol7/KE0NQwoLaSXVNURVXVsVf2X7kR7FHBPVT0w1O4u4Ohu+tXADwFfSnJzkhfNsv9jgYu720f3AvcAGdoXwNbdbHskcHB37HF1zLbtjKNG2ty1u4bs2eeay7FH29zV1bOv5tIvX5uZqKqHuslxL8AcBWytqu/Osq9JvjY0/RDw2P34mdwBz4DSYtsO/PMkhw0tOwb4e4Cq+nJVrWVwa+vtwDVJDgHGDcO/Ffi5LgRnfh5XVX891GZ3w/ffzeCqZvhqa1cdE7ad8VVgxcj2Y+3h55rLsRlz7O3d9DeB7x9a98Q92Pdc+mWutgMrZp5f7eO+dAAwoLSoulthfw38epLHJnkag6uLDwMkeVmSZd1v3fd2mz0C7AC+y+C5yIzfAd6U5KndtocnmdPryt0to6uBtyU5rLst+EbgQ7Nv+ShXd8f/Z0mWA6/fXcM9/Fxz9YvdsVcA5wIf7ZbfCpya5JgkhwNvGtnuH3Z3vHnqlxk3MQjLX0pyUAZ/E/ZiBs/WpO9hQKkFa4GVDH7D/jjw5qr6TLfudGBTkgcZvFhwdves6iHgbcBfdbf0TqmqjzO4Grkqyf3AF4HdvlE2xusZnEC3MHhO9hEGD/Tn6q0Mbll9hcGLBB+cpe2cP9ceHP9PgFsYBNKf0r3E0PXlR4HPd+uvHdnuYuCs7i28cc/N9rVf6Op4mMFLDmcwuDL7LeBnq+pLe7ovHRjiFxZKklrkFZQkqUkGlCSpSQaUJKlJBpQkqUlT9wduRx55ZK1cuXKxy5Ak7aVbbrnl7qpaNqnd1AXUypUr2bhx42KXIUnaS0lmG2VlF2/xSZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmjR1I0nMh3UXXbnYJUylS85fu9glSDqAeAUlSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWpSrwGV5PQkdyTZnOS8WdqdlaSSrO6zHknS9OgtoJIsAS4FzgBWAWuTrBrT7jDgDcBNfdUiSZo+fV5BnQRsrqotVfUwcBVw5ph2/w14B/BPPdYiSZoyfQbU0cDWoflt3bJdkjwTWFFV1862oyTnJNmYZOOOHTvmv1JJUnP6DKiMWVa7ViaPAd4N/MKkHVXVZVW1uqpWL1u2bB5LlCS1qs+A2gasGJpfDmwfmj8MOBG4PsmdwCnAel+UkCRBvwF1M3BCkuOSHAycDayfWVlV91XVkVW1sqpWAjcCa6pqY481SZKmRG8BVVU7gXXAdcDtwNVVtSnJhUnW9HVcSdL+YWmfO6+qDcCGkWUX7Kbtc/usRZI0XRxJQpLUJANKktQkA0qS1KRen0HpwLLuoisXu4SpdMn5axe7BKlJXkFJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmtRrQCU5PckdSTYnOW/M+tcm+UKSW5PckGRVn/VIkqZHbwGVZAlwKXAGsApYOyaAPlJVP1xVzwDeAbyrr3okSdOlzyuok4DNVbWlqh4GrgLOHG5QVfcPzR4CVI/1SJKmyNIe9300sHVofhtw8mijJK8D3ggcDDxv3I6SnAOcA3DMMcfMe6GSpPb0eQWVMcu+5wqpqi6tqh8Efhk4f9yOquqyqlpdVauXLVs2z2VKklrUZ0BtA1YMzS8Hts/S/irgJT3WI0maIn0G1M3ACUmOS3IwcDawfrhBkhOGZl8IfLnHeiRJU6S3Z1BVtTPJOuA6YAnwgaralORCYGNVrQfWJTkN+A7wDeAVfdUjSZoufb4kQVVtADaMLLtgaPrcPo8vSZpejiQhSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlq0qwBleT3h6Zf0Xs1kiR1Jl1BPX1o+tw+C5EkadikgKoFqUKSpBFLJ6xfnuQ9QIamd6mqN/RWmSTpgDYpoH5xaHpjn4VIkjRs1oCqqisWqhBJkobNGlBJ1s+2vqrWzG85kiQNTLrF92xgK3AlcBODZ1GSJPVuUkA9EfhxYC3w08CfAldW1aa+C5MkHdhmfc28qh6pqk9V1SuAU4DNwPVJXr8g1UmSDliTrqBI8n3ACxlcRa0E3gN8rN+yJEkHukkvSVwBnAh8EnhrVX1xQaqSJB3wJl1BvRz4JvBDwLlJZkaWCFBV9fg+i5MkHbgm/R2Uo51LkhbFpFt8jwVeCxwPfB74QFXtXIjCJEkHtklXSFcAq4EvAC8AfqP3iiRJYvIzqFVV9cMASS4H/qb/kiRJmnwF9Z2ZCW/tSZIW0sQvLExyf/fzAPC0mekk90/aeZLTk9yRZHOS88asf2OS25J8PsmfJzl2bz+IJGn/MuktviV7u+MkS4BLGQyVtA24Ocn6qrptqNnfAqur6qEkPw+8A3jp3h5TkrT/6PM18pOAzVW1paoeBq4CzhxuUFV/UVUPdbM3Ast7rEeSNEX6DKijGYyEPmNbt2x3Xs1gxIrvkeScJBuTbNyxY8c8lihJalWfATXuqzlqzDKSvIzB6+zvHLe+qi6rqtVVtXrZsmXzWKIkqVUTB4vdB9uAFUPzy4Hto42SnAb8KvDvqurbPdYjSZoifV5B3QyckOS4JAcDZwOP+obeJM8EfhdYU1Vf77EWSdKU6S2gur+bWgdcB9wOXF1Vm5JcmGTmq+LfCRwK/GGSWyd9xbwk6cDR5y0+qmoDsGFk2QVD06f1eXxJ0vRytHJJUpMMKElSkwwoSVKTDChJUpMMKElSkwwoSVKTDChJUpMMKElSkwwoSVKTDChJUpMMKElSkwwoSVKTDChJUpMMKElSkwwoSVKTev0+KEn7bt1FVy52CVPpkvPXLnYJ2kdeQUmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKaZEBJkppkQEmSmmRASZKa1GtAJTk9yR1JNic5b8z6U5N8LsnOJGf1WYskabr0FlBJlgCXAmcAq4C1SVaNNPs74JXAR/qqQ5I0nZb2uO+TgM1VtQUgyVXAmcBtMw2q6s5u3Xd7rEOSNIX6vMV3NLB1aH5bt2yPJTknycYkG3fs2DEvxUmS2tZnQGXMstqbHVXVZVW1uqpWL1u2bB/LkiRNgz4DahuwYmh+ObC9x+NJkvYjfQbUzcAJSY5LcjBwNrC+x+NJkvYjvQVUVe0E1gHXAbcDV1fVpiQXJlkDkORHkmwDfhL43SSb+qpHkjRd+nyLj6raAGwYWXbB0PTNDG79SZL0KI4kIUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJatLSxS5AkhbLuouuXOwSps4l569dsGN5BSVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJUlqUq8BleT0JHck2ZzkvDHrvy/JR7v1NyVZ2Wc9kqTp0VtAJVkCXAqcAawC1iZZNdLs1cA3qup44N3A2/uqR5I0Xfq8gjoJ2FxVW6rqYeAq4MyRNmcCV3TT1wA/liQ91iRJmhKpqn52nJwFnF5Vr+nmXw6cXFXrhtp8sWuzrZv/v12bu0f2dQ5wTjf7FOCOPSznSODuia3aYK39sNZ+WGs/pqlW2PN6j62qZZMa9TlY7LgrodE0nEsbquoy4LK9LiTZWFWr93b7hWSt/bDWflhrP6apVuiv3j5v8W0DVgzNLwe2765NkqXA4cA9PdYkSZoSfQbUzcAJSY5LcjBwNrB+pM164BXd9FnAZ6uve46SpKnS2y2+qtqZZB1wHbAE+EBVbUpyIbCxqtYDlwMfTLKZwZXT2T2Vs9e3BxeBtfbDWvthrf2Yplqhp3p7e0lCkqR94UgSkqQmGVCSpCbtFwGVZEWSv0hye5JNSc4d0+a5Se5Lcmv3c8Fi1NrV8tgkf5Pkf3f1vnVMmyaGgZpjra9MsmOob1+zGLUO1bMkyd8muXbMuib6daie2Wptpl+T3JnkC10dG8esT5L3dP36+STPWow6u1om1drSueCIJNck+VJ3/nr2yPqW+nVSrfPer33+HdRC2gn8QlV9LslhwC1JPlNVt420+19V9aJFqG/Ut4HnVdWDSQ4Cbkjyyaq6cajNrmGgkpzNYBiolzZaK8BHh/8Ie5GdC9wOPH7Mulb6dcZstUJb/frvR/+IfsgZwAndz8nAb3f/LpbZaoV2zgUXA5+qqrO6t52/f2R9S/06qVaY537dL66gquqrVfW5bvoBBv/BH724Ve1eDTzYzR7U/Yy+rdLEMFBzrLUZSZYDLwTev5smTfQrzKnWaXIm8Afd/19uBI5I8qTFLqplSR4PnMrgbWaq6uGqunekWRP9Osda591+EVDDuls2zwRuGrP62d2tqk8meeqCFjaiu7VzK/B14DNVNVrv0cBWGLyyD9wHPGFhqxyYQ60AP9HdgrgmyYox6xfKbwK/BHx3N+ub6Vcm1wrt9GsBn05ySwZDj43a1a+dbSzeL4mTaoU2zgVPBnYAv9fd5n1/kkNG2rTSr3OpFea5X/ergEpyKPBHwH+tqvtHVn+OwfhPTwfeC/zxQtc3rKoeqapnMBhh46QkJ440mdMwUAthDrV+AlhZVU8D/oz/f4WyoJK8CPh6Vd0yW7Mxyxa8X+dYaxP92nlOVT2LwS2n1yU5dWR9E/3amVRrK+eCpcCzgN+uqmcC3wRGv5aolX6dS63z3q/7TUB1z0f+CPhwVX1sdH1V3T9zq6qqNgAHJTlygcv8Ht1l8vXA6SOrmhsGane1VtU/VtW3u9n3Af96gUub8RxgTZI7GYye/7wkHxpp00q/Tqy1oX6lqrZ3/34d+DiDbysYNpehzRbEpFobOhdsA7YN3ZG4hkEIjLZpoV8n1tpHv+4XAdU9Q7gcuL2q3rWbNk+cedaQ5CQGn/0fF67KR9WyLMkR3fTjgNOAL400a2IYqLnUOnJPfA2DZ4ALrqreVFXLq2olg1FJPltVLxtp1kS/zqXWVvo1ySHdy0d0t3WeD3xxpNl64Ge7t85OAe6rqq8ucKlzqrWVc0FVfQ3YmuQp3aIfA0Zf7GqiX+dSax/9ur+8xfcc4OXAF7pnJQC/AhwDUFW/w+Bk9PNJdgLfAs5exHH/ngRckcGXOj4GuLqqrs3iDAM1H7W+IckaBm9T3gO8cpFqHavRfh2r0X79AeDj3blnKfCRqvpUktfCrv++NgAvADYDDwGvarjWls4Frwc+3L0VtwV4VaP9CpNrnfd+dagjSVKT9otbfJKk/Y8BJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJapIBJTUiycoMvmfnfRl899anu9E7pAOSASW15QTg0qp6KnAv8BOLXI+0aAwoqS1fqaqZ4bpuAVYuYi3SojKgpLZ8e2j6Efaf8TKlPWZASZKaZEBJkprkaOaSpCZ5BSVJapIBJUlqkgElSWqSASVJapIBJUlqkgElSWqSASVJatL/A3FtKfu+rgcbAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "thinkplot.Hist(pmf_n)\n",
    "print(pmf_n.Mean())\n",
    "thinkplot.decorate(xlabel='n', \n",
    "                   ylabel='PMF', \n",
    "                   title='Posterior distribution of n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, to answer the question, we have to compute the posterior distribution of the prevalence of elephants.  Here's a function that computes it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enumerate_posterior(suite):\n",
    "    for d, p in suite.Items():\n",
    "        mean = d.mean()\n",
    "        index = d.label.find('E')\n",
    "        p_elephant = 0 if index == -1 else mean[index]\n",
    "        yield d, p, p_elephant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the possible zoos, the posterior probability of each, and the conditional prevalence of elephants for each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LTB 0.43679617920687586 0\n",
      "LTBC 0.12516851264765871 0\n",
      "LTBR 0.1304178179372656 0\n",
      "LTBE 0.12616846874234114 0.1\n",
      "LTBCR 0.048267930148971114 0\n",
      "LTBCE 0.056066029260621264 0.09090909090909091\n",
      "LTBRE 0.052742406633655016 0.09090909090909091\n",
      "LTBCRE 0.024372655422611408 0.08333333333333333\n"
     ]
    }
   ],
   "source": [
    "for d, p, p_elephant in enumerate_posterior(suite):\n",
    "    print(d.label, p, p_elephant)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can use the law of total probability to compute the probability of seeing an elephant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.024539577483173817"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total = np.sum(p * p_elephant \n",
    "               for d, p, p_elephant in enumerate_posterior(suite))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}