PythonCounterPmf.ipynb

{
 "metadata": {
  "name": "PythonCounterPmf"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Using Counters\n",
      "--------------\n",
      "\n",
      "This notebook demonstrates the use of the Counter class to implement Multisets, PMFs, and suites of Bayesian hypotheses."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from collections import Counter"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "A counter is a map from values to their frequencies.  If you initialize a counter with a string, you get a map from each letter to the number of times it appears.  If two words are anagrams, they yield equal Counters, so you can use Counters to test anagrams in linear time."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def is_anagram(word1, word2):\n",
      "    \"\"\"Checks whether the words are anagrams.\n",
      "\n",
      "    word1: string\n",
      "    word2: string\n",
      "\n",
      "    returns: boolean\n",
      "    \"\"\"\n",
      "    return Counter(word1) == Counter(word2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "is_anagram('tachymetric', 'mccarthyite')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 3,
       "text": [
        "True"
       ]
      }
     ],
     "prompt_number": 3
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "is_anagram('banana', 'peach')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 4,
       "text": [
        "False"
       ]
      }
     ],
     "prompt_number": 4
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "**Multisets**\n",
      "\n",
      "A Counter is a natural representation of a multiset, which is a set where the elements can appear more than once.  You can extend Counter with set operations like is_subset:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class Multiset(Counter):\n",
      "    \"\"\"A multiset is a set where elements can appear more than once.\"\"\"\n",
      "\n",
      "    def is_subset(self, other):\n",
      "        \"\"\"Checks whether self is a subset of other.\n",
      "\n",
      "        other: Multiset\n",
      "\n",
      "        returns: boolean\n",
      "        \"\"\"\n",
      "        for char, count in self.items():\n",
      "            if other[char] < count:\n",
      "                return False\n",
      "        return True\n",
      "    \n",
      "    # map the <= operator to is_subset\n",
      "    __le__ = is_subset"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 33
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "You could use <tt>is_subset</tt> in a game like Scrabble to see if a given set of tiles can be used to spell a given word."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def can_spell(word, tiles):\n",
      "    \"\"\"Checks whether a set of tiles can spell a word.\n",
      "\n",
      "    word: string\n",
      "    tiles: string\n",
      "\n",
      "    returns: boolean\n",
      "    \"\"\"\n",
      "    return Multiset(word) <= Multiset(tiles)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 34
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "can_spell('SYZYGY', 'AGSYYYZ')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "pyout",
       "prompt_number": 35,
       "text": [
        "True"
       ]
      }
     ],
     "prompt_number": 35
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Probability Mass Functions\n",
      "--------------------------\n",
      "\n",
      "You can also extend Counter to represent a probability mass function (PMF).\n",
      "\n",
      "`normalize` computes the total of the frequencies and divides through, yielding probabilities that add to 1.\n",
      "\n",
      "`__add__` enumerates all pairs of value and returns a new Pmf that represents the distribution of the sum.\n",
      "\n",
      "`__hash__` and `__id__` make Pmfs hashable; this is not the best way to do it, because they are mutable.  So this implementation comes with a warning that if you use a Pmf as a key, you should not modify it.  A better alternative would be to define a frozen Pmf.\n",
      "\n",
      "`render` returns the values and probabilities in a form ready for plotting"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class Pmf(Counter):\n",
      "    \"\"\"A Counter with probabilities.\"\"\"\n",
      "\n",
      "    def normalize(self):\n",
      "        \"\"\"Normalizes the PMF so the probabilities add to 1.\"\"\"\n",
      "        total = float(sum(self.values()))\n",
      "        for key in self:\n",
      "            self[key] /= total\n",
      "\n",
      "    def __add__(self, other):\n",
      "        \"\"\"Adds two distributions.\n",
      "\n",
      "        The result is the distribution of sums of values from the\n",
      "        two distributions.\n",
      "\n",
      "        other: Pmf\n",
      "\n",
      "        returns: new Pmf\n",
      "        \"\"\"\n",
      "        pmf = Pmf()\n",
      "        for key1, prob1 in self.items():\n",
      "            for key2, prob2 in other.items():\n",
      "                pmf[key1 + key2] += prob1 * prob2\n",
      "        return pmf\n",
      "\n",
      "    def __hash__(self):\n",
      "        \"\"\"Returns an integer hash value.\"\"\"\n",
      "        return id(self)\n",
      "    \n",
      "    def __eq__(self, other):\n",
      "        return self is other\n",
      "\n",
      "    def render(self):\n",
      "        \"\"\"Returns values and their probabilities, suitable for plotting.\"\"\"\n",
      "        return zip(*sorted(self.items()))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 8
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "As an example, we can make a Pmf object that represents a 6-sided die."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "d6 = Pmf([1,2,3,4,5,6])\n",
      "d6.normalize()\n",
      "d6.name = 'one die'\n",
      "print(d6)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Pmf({1: 0.16666666666666666, 2: 0.16666666666666666, 3: 0.16666666666666666, 4: 0.16666666666666666, 5: 0.16666666666666666, 6: 0.16666666666666666})\n"
       ]
      }
     ],
     "prompt_number": 9
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Using the add operator, we can compute the distribution for the sum of two dice."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "d6_twice = d6 + d6\n",
      "d6_twice.name = 'two dice'\n",
      "\n",
      "for key, prob in d6_twice.items():\n",
      "    print(key, prob)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "(2, 0.027777777777777776)\n",
        "(3, 0.05555555555555555)\n",
        "(4, 0.08333333333333333)\n",
        "(5, 0.1111111111111111)\n",
        "(6, 0.1388888888888889)\n",
        "(7, 0.16666666666666669)\n",
        "(8, 0.1388888888888889)\n",
        "(9, 0.1111111111111111)\n",
        "(10, 0.08333333333333333)\n",
        "(11, 0.05555555555555555)\n",
        "(12, 0.027777777777777776)\n"
       ]
      }
     ],
     "prompt_number": 10
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Using numpy.sum, we can compute the distribution for the sum of three dice."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# if we use the built-in sum we have to provide a Pmf additive identity value\n",
      "# pmf_ident = Pmf([0])\n",
      "# d6_thrice = sum([d6]*3, pmf_ident)\n",
      "\n",
      "# with --numpy inline we get numpy.sum, which does not require an identity\n",
      "d6_thrice = sum([d6]*3)\n",
      "d6_thrice.name = 'three dice'"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 17
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "And then plot the results (using Pmf.render)"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import matplotlib.pyplot as pyplot"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 18
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "for die in [d6, d6_twice, d6_thrice]:\n",
      "    xs, ys = die.render()\n",
      "    pyplot.plot(xs, ys, label=die.name, linewidth=3, alpha=0.5)\n",
      "    \n",
      "pyplot.xlabel('Total')\n",
      "pyplot.ylabel('Probability')\n",
      "pyplot.legend()\n",
      "pyplot.show()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Bayesian statistics\n",
      "-------------------\n",
      "\n",
      "A Suite is a Pmf that represents a set of hypotheses and their probabilities; it provides `bayesian_update`, which updates the probability of the hypotheses based on new data.\n",
      "\n",
      "Suite is an abstract parent class; child classes should provide a likelihood method that evaluates the likelihood of the data under a given hypothesis.  `update_bayesian` loops through the hypothesis, evaluates the likelihood of the data under each hypothesis, and updates the probabilities accordingly.  Then it re-normalizes the PMF."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class Suite(Pmf):\n",
      "    \"\"\"Map from hypothesis to probability.\"\"\"\n",
      "\n",
      "    def bayesian_update(self, data):\n",
      "        \"\"\"Performs a Bayesian update.\n",
      "        \n",
      "        Note: called bayesian_update to avoid overriding dict.update\n",
      "\n",
      "        data: result of a die roll\n",
      "        \"\"\"\n",
      "        for hypo in self:\n",
      "            like = self.likelihood(data, hypo)\n",
      "            self[hypo] *= like\n",
      "\n",
      "        self.normalize()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 24
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "As an example, I'll use Suite to solve the \"Dice Problem,\" from Chapter 3 of <i>Think Bayes</i>.\n",
      "\n",
      "\"Suppose I have a box of dice that contains a 4-sided die, a 6-sided die, an 8-sided die, a 12-sided die, and a 20-sided die. If you have ever played Dungeons & Dragons, you know what I am talking about.  Suppose I select a die from the box at random, roll it, and get a 6. What is the probability that I rolled each die?\"\n",
      "\n",
      "I'll start by making a list of Pmfs to represent the dice:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def make_die(num_sides):\n",
      "    die = Pmf(range(1, num_sides+1))\n",
      "    die.name = 'd%d' % num_sides\n",
      "    die.normalize()\n",
      "    return die\n",
      "\n",
      "\n",
      "dice = [make_die(x) for x in [4, 6, 8, 12, 20]]\n",
      "print(dice)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "[Pmf({1: 0.25, 2: 0.25, 3: 0.25, 4: 0.25}), Pmf({1: 0.16666666666666666, 2: 0.16666666666666666, 3: 0.16666666666666666, 4: 0.16666666666666666, 5: 0.16666666666666666, 6: 0.16666666666666666}), Pmf({1: 0.125, 2: 0.125, 3: 0.125, 4: 0.125, 5: 0.125, 6: 0.125, 7: 0.125, 8: 0.125}), Pmf({1: 0.08333333333333333, 2: 0.08333333333333333, 3: 0.08333333333333333, 4: 0.08333333333333333, 5: 0.08333333333333333, 6: 0.08333333333333333, 7: 0.08333333333333333, 8: 0.08333333333333333, 9: 0.08333333333333333, 10: 0.08333333333333333, 11: 0.08333333333333333, 12: 0.08333333333333333}), Pmf({1: 0.05, 2: 0.05, 3: 0.05, 4: 0.05, 5: 0.05, 6: 0.05, 7: 0.05, 8: 0.05, 9: 0.05, 10: 0.05, 11: 0.05, 12: 0.05, 13: 0.05, 14: 0.05, 15: 0.05, 16: 0.05, 17: 0.05, 18: 0.05, 19: 0.05, 20: 0.05})]\n"
       ]
      }
     ],
     "prompt_number": 36
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Next I'll define DiceSuite, which inherits `bayesian_update` from Suite and provides `likelihood`.\n",
      "\n",
      "`data` is the observed die roll, 6 in the example.\n",
      "\n",
      "`hypo` is the hypothetical die I might have rolled; to get the likelihood of the data, I select, from the given die, the probability of the given value. "
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "class DiceSuite(Suite):\n",
      "    \n",
      "    def likelihood(self, data, hypo):\n",
      "        \"\"\"Computes the likelihood of the data under the hypothesis.\n",
      "\n",
      "        data: result of a die roll\n",
      "        hypo: Die object\n",
      "        \"\"\"\n",
      "        return hypo[data]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 27
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Finally, I use the list of dice to instantiate a Suite that maps from each die to its prior probability.  By default, all dice have the same prior.\n",
      "\n",
      "Then I update the distribution with the given value and print the results:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dice_suite = DiceSuite(dice)\n",
      "\n",
      "dice_suite.bayesian_update(6)\n",
      "\n",
      "for die, prob in sorted(dice_suite.items()):\n",
      "    print die.name, prob"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "d4 0.0\n",
        "d6 0.392156862745\n",
        "d8 0.294117647059\n",
        "d12 0.196078431373\n",
        "d20 0.117647058824\n"
       ]
      }
     ],
     "prompt_number": 29
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "As expected, the 4-sided die has been eliminated; it now has 0 probability.  The 6-sided die is the most likely, but the 8-sided die is still quite possible.\n",
      "\n",
      "Now suppose I roll the die again and get an 8.  We can update the Suite again with the new data"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "dice_suite.bayesian_update(8)\n",
      "\n",
      "for die, prob in sorted(dice_suite.items()):\n",
      "    print die.name, prob"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "d4 0.0\n",
        "d6 0.0\n",
        "d8 0.623268698061\n",
        "d12 0.277008310249\n",
        "d20 0.0997229916898\n"
       ]
      }
     ],
     "prompt_number": 37
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Now the 6-sided die has been eliminated, the 8-sided die is most likely, and there is less than a 10% chance that I am rolling a 20-sided die.\n",
      "\n",
      "These examples demonstrate the versatility of the Counter class, one of Python's underused data structures."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}