{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lab 4: Belief Propagation\n",
    "\n",
    "This lab is built around the process of identifying the fault with a coffee machine.\n",
    "\n",
    "Your task is to:\n",
    " 1. Given the structure of a graphical model for the state of a coffee machine learn the distributions from data.\n",
    " 2. Implement belief propagation, so you can evaluate the probability of each failure given the available evidence.\n",
    " 3. Identify the most probable failure for a set of broken coffee machines."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Marking and Submission\n",
    "\n",
    "These lab exercises are marked, and contribute to your final grade. For this lab exercise there are 3 places where you are expected to enter your own code, for 15 marks overall. Every place you have to add code is indicated by\n",
    "\n",
    "`# **************************************************************** n marks`\n",
    "\n",
    "with instructions above the code block.\n",
    "\n",
    "Please submit your completed workbook using Moodle before 5pm on the 28th November 2018. The workbook you submit must be an `.ipynb` file, which is saved into the directory you're running Jupyter; alternatively you can download it from the menu above using `File -> Download As -> Notebook (.ipynb)`. Remember to save your work regularly (Save and checkpoint in the File menu, the icon of a floppy disk, or Ctrl-S); the version you submit should have all code blocks showing the results (if any) of execution below them. You will normally receive feedback within a week."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import zipfile\n",
    "import io\n",
    "import csv\n",
    "\n",
    "import numpy\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coffee Machine Data Set\n",
    "\n",
    "You can download from moodle a zip file that contains the data as a csv file. The below code will load the data directly from the zip file, so you don't have to unzip it.\n",
    "\n",
    "Each row of the file contains a fully observed coffee machine, with the state of every random variable. The random variables are all binary, with False represented by 0 and True represented by 1. The variables are:\n",
    "\n",
    "Failures (you're trying to detect these):  \n",
    "* 0. `he` - No electricity\n",
    "* 1. `fp` - Fried power supply unit\n",
    "* 2. `fc` - Fried circuit board  \n",
    "* 3. `wr` - Water reservoir empty\n",
    "* 4. `gs` - Group head gasket forms seal  \n",
    "* 5. `dp` - Dead pump  \n",
    "* 6. `fh` - Fried heating element  \n",
    "\n",
    "\n",
    "Mechanism (these are unobservable):  \n",
    "* 7. `pw` - Power supply unit works  \n",
    "* 8. `cb` - Circuit board works  \n",
    "* 9. `gw` - Get water out of group head   \n",
    "\n",
    "Diagnostic (these are the tests the mechanic can run - observable):\n",
    "* 10. `ls` - Room lights switch on\n",
    "* 11. `vp` - A voltage is measured across power supply unit\n",
    "* 12. `lo` - Power light switches on\n",
    "* 13. `wv` - Water visible in reservoir\n",
    "* 14. `hp` - Can hear pump\n",
    "* 15. `me` - Makes espresso\n",
    "* 16. `ta` - Makes a hot, tasty espresso\n",
    "\n",
    "For the above the number is the column number of the provided data (`dm`) and the two letter code a suggested variable name.\n",
    "\n",
    "For if you are unfamiliar with an espresso coffee machine here is a brief description of how one works (you can ignore this if you want):\n",
    "> The user puts ground coffee into a portafilter (round container with a handle and two spouts at the bottom), tamps it (compacts the coffee down), and clamps the portafilter into the group head at the front of the machine. A gasket (rubber ring) forms a seal between the portafilter and group head. A button is pressed. Water is drawn from a reservoir by a pump into a boiler. In the boiler a heating element raises the waters temperature, before the pump pushes it through the group head and into the portafilter at high pressure. The water passes through the coffee grinds and makes a tasty espresso.\n",
    "\n",
    "The graphical model showing how the variables are related is included on moodle as `coffee machine.pdf`; here it is given as conditional probabilities:\n",
    "\n",
    "Failures:\n",
    " * `P_he:` $P(\\texttt{no electricity})$\n",
    " * `P_fp:` $P(\\texttt{fried psu})$\n",
    " * `P_fc:` $P(\\texttt{fried circuit board})$\n",
    " * `P_wr:` $P(\\texttt{water reservoir empty})$\n",
    " * `P_gs:` $P(\\texttt{group head gasket seal broken})$\n",
    " * `P_dp:` $P(\\texttt{dead pump})$\n",
    " * `P_fh:` $P(\\texttt{fried heating element})$\n",
    "\n",
    "Mechanism:\n",
    " * `P_pw_he_fp:` $P(\\texttt{psu works}\\enspace|\\enspace\\texttt{no electricity},\\enspace\\texttt{fried psu})$\n",
    " * `P_cb_pw_fc:` $P(\\texttt{circuit board works}\\enspace|\\enspace\\texttt{psu works},\\enspace\\texttt{fried circuit board})$\n",
    " * `P_gw_cb_wr_dp:` $P(\\texttt{get water}\\enspace|\\enspace\\texttt{circuit board works},\\enspace\\texttt{water reservoir empty},\\enspace\\texttt{dead pump})$\n",
    "\n",
    "Diagnostic:\n",
    " * `P_ls_he:` $P(\\texttt{lights switch on}\\enspace|\\enspace\\texttt{no electricity})$\n",
    " * `P_vp_pw:` $P(\\texttt{voltage across psu}\\enspace|\\enspace\\texttt{psu works})$\n",
    " * `P_lo_cb:` $P(\\texttt{power light on}\\enspace|\\enspace\\texttt{circuit board works})$\n",
    " * `P_wv_wr:` $P(\\texttt{water visible}\\enspace|\\enspace\\texttt{water reservoir empty})$\n",
    " * `P_hp_dp:` $P(\\texttt{can hear pump}\\enspace|\\enspace\\texttt{dead pump})$\n",
    " * `P_me_gw_gs:` $P(\\texttt{makes espresso}\\enspace|\\enspace\\texttt{get water},\\enspace\\texttt{group head gasket seal broken})$\n",
    " * `P_ta_me_fh:` $P(\\texttt{tasty}\\enspace|\\enspace\\texttt{makes espresso},\\enspace\\texttt{fried heating element})$\n",
    "\n",
    "Note that while the model is close to what you may guess the probabilities are not absolute, to account for mistakes and unknown failures. For instance, the mechanic may make a mistake while brewing an espresso and erroneously conclude that the machine is broken when it is in fact awesome. The probabilities associated with each failure are not uniform.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# It may prove helpful to have a mapping between the suggested variable names and\n",
    "# column indices in the provided file...\n",
    "nti = dict() # 'name to index'\n",
    "\n",
    "nti['he'] = 0\n",
    "nti['fp'] = 1\n",
    "nti['fc'] = 2\n",
    "nti['wr'] = 3\n",
    "nti['gs'] = 4\n",
    "nti['dp'] = 5\n",
    "nti['fh'] = 6\n",
    "\n",
    "nti['pw'] = 7\n",
    "nti['cb'] = 8\n",
    "nti['gw'] = 9\n",
    "\n",
    "nti['ls'] = 10\n",
    "nti['vp'] = 11\n",
    "nti['lo'] = 12\n",
    "nti['wv'] = 13\n",
    "nti['hp'] = 14\n",
    "nti['me'] = 15\n",
    "nti['ta'] = 16\n",
    "\n",
    "\n",
    "\n",
    "# Opposite to the above - index to name...\n",
    "itn = ['he', 'fp', 'fc', 'wr', 'gs', 'dp', 'fh',\n",
    "       'pw', 'cb', 'gw',\n",
    "       'ls', 'vp', 'lo', 'wv', 'hp', 'me', 'ta']\n",
    "\n",
    "\n",
    "\n",
    "# For conveniance this code loads the data from the zip file,\n",
    "# so you don't have to decompress it (takes a few seconds to run)...\n",
    "with zipfile.ZipFile('coffee_machines.zip') as zf:\n",
    "    with zf.open('coffee_machines.csv') as f:\n",
    "        sf = io.TextIOWrapper(f)\n",
    "        reader = csv.reader(sf)\n",
    "        next(reader)\n",
    "        dm = []\n",
    "        for row in reader:\n",
    "            dm.append([int(v) for v in row])\n",
    "        dm = numpy.array(dm, dtype=numpy.int8)\n",
    "\n",
    "print('Data: {} exemplars, {} features'.format(dm.shape[0], dm.shape[1]))\n",
    "print('     Broken machines  =', dm.shape[0] - dm[:,nti['ta']].sum())\n",
    "print('     Working machines =', dm[:,nti['ta']].sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Learn Model\n",
    "\n",
    "Below a set of variables to represent conditional probability distributions have been defined. They are a Bernoulli trial for each combination of conditional variables, given as $P(\\texttt{False}|...)$ in `[0,...]` and $P(\\texttt{True}|...)$ in `[1,...]` (It may be easier to think of them as boring categorical distributions).\n",
    "\n",
    "To fit a maximum likelihood model you should first use them as counters - loop over the data set and count how many of each combination exist. You must then normalise them so that the sum over axis 0 is always 1. There is an extra mark for doing the right thing and including a prior. You may want to know that the conjugate prior to the Bernoulli trial represented as $\\left[P(\\texttt{False}), P(\\texttt{True})\\right]$ is a Beta distribution; a uniform prior would be a reasonable choice (it can be argued that the expected failure probability is low, and you should adjust the first 7 variables accordingly, but given the quantity of data available it's not going to matter).\n",
    "\n",
    "Hint:\n",
    " * The use of `0=False` and `1=True` both in the dm table and in the conditional probability distributions is very deliberate.\n",
    " * Consider putting all of the variables into a list with extra information about them (such as indices from `nti`) to make your code neater.\n",
    " * If you make a mistake you could easily end up with NaN or infinity - which would break the next step. Print them out so you can check they are valid!\n",
    " * Do not write unique code for each variable - that will be hundreds of lines of code and very tedious. It's possible to get all of the marks in 7 lines of code, if you're particularly sneaky.\n",
    "\n",
    "__(4 marks)__\n",
    " * 2 marks for counting all conditions\n",
    " * 1 mark for normalising distributions\n",
    " * 1 mark for including a sensible prior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A set of variables that will ultimately represent conditional probability distributions.\n",
    "# The naming convention is that P_he means P(he), or that P_ta_me_hw means P(ta|me,hw)...\n",
    "\n",
    "P_he          = numpy.zeros(2)\n",
    "P_fp          = numpy.zeros(2)\n",
    "P_fc          = numpy.zeros(2)\n",
    "P_wr          = numpy.zeros(2)\n",
    "P_gs          = numpy.zeros(2)\n",
    "P_dp          = numpy.zeros(2)\n",
    "P_fh          = numpy.zeros(2)\n",
    "\n",
    "\n",
    "P_pw_he_fp    = numpy.zeros((2,2,2))\n",
    "P_cb_pw_fc    = numpy.zeros((2,2,2))\n",
    "P_gw_cb_wr_dp = numpy.zeros((2,2,2,2))\n",
    "\n",
    "P_ls_he       = numpy.zeros((2,2))\n",
    "P_vp_pw       = numpy.zeros((2,2))\n",
    "P_lo_cb       = numpy.zeros((2,2))\n",
    "P_wv_wr       = numpy.zeros((2,2))\n",
    "P_hp_dp       = numpy.zeros((2,2))\n",
    "P_me_gw_gs    = numpy.zeros((2,2,2))\n",
    "P_ta_me_fh    = numpy.zeros((2,2,2))\n",
    "\n",
    "\n",
    "\n",
    "# It may prove conveniant to have a structure that describes the above in a computer\n",
    "# readable form, including labeling what each dimension is...\n",
    "ops = [(P_he, nti['he']),\n",
    "       (P_fp, nti['fp']),\n",
    "       (P_fc, nti['fc']),\n",
    "       (P_wr, nti['wr']),\n",
    "       (P_gs, nti['gs']),\n",
    "       (P_dp, nti['dp']),\n",
    "       (P_fh, nti['fh']),\n",
    "       (P_pw_he_fp, nti['pw'], nti['he'], nti['fp']),\n",
    "       (P_cb_pw_fc, nti['cb'], nti['pw'], nti['fc']),\n",
    "       (P_gw_cb_wr_dp, nti['gw'], nti['cb'], nti['wr'], nti['dp']),\n",
    "       (P_ls_he, nti['ls'], nti['he']),\n",
    "       (P_vp_pw, nti['vp'], nti['pw']),\n",
    "       (P_lo_cb, nti['lo'], nti['cb']),\n",
    "       (P_wv_wr, nti['wv'], nti['wr']),\n",
    "       (P_hp_dp, nti['hp'], nti['dp']),\n",
    "       (P_me_gw_gs, nti['me'], nti['gw'], nti['gs']),\n",
    "       (P_ta_me_fh, nti['ta'], nti['me'], nti['fh'])]\n",
    "\n",
    "\n",
    "\n",
    "# **************************************************************** 4 marks\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Brute Force\n",
    "The below code does exactly what you want to implement for step 2, but it brute forces it. Slow and memory inefficient of course, but lets you test it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def brute_marginals(known):\n",
    "    \"\"\"known is a dictionary, where a random variable index existing as a key in the dictionary\n",
    "    indicates it has been observed. The value obtained using the key is the value the\n",
    "    random variable has been observed as. Returns a 17x2 matrix, such that [rv,0] is the\n",
    "    probability of random variable rv being False, [rv, 1] the probability of being True.\"\"\"\n",
    "    \n",
    "    # Calculate the full joint (please don't ask)...\n",
    "    params = []\n",
    "    for op in ops:\n",
    "        params.append(op[0])\n",
    "        params.append(op[1:])\n",
    "    params.append(range(17))\n",
    "    joint = numpy.einsum(*params)\n",
    "    \n",
    "    # Multiply in the known states (zero out the dimensions we know it's not)...\n",
    "    for key, value in known.items():\n",
    "        other = abs(value-1)\n",
    "        index = [slice(None)] * len(joint.shape)\n",
    "        index[key] = other\n",
    "        joint[tuple(index)] = 0.0\n",
    "    \n",
    "    # Calculate and return all marginals...\n",
    "    ret = numpy.empty((len(joint.shape), 2))\n",
    "    for row in range(ret.shape[0]):\n",
    "        ret[row,:] = numpy.einsum(joint, range(len(joint.shape)), [row])\n",
    "    \n",
    "    ret /= ret.sum(axis=1)[:,None]\n",
    "    return ret\n",
    "\n",
    "print(brute_marginals({}))\n",
    "print(brute_marginals({nti['ta'] : 0}))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Belief Propagation\n",
    "\n",
    "Your task is to write a function that take known states and calculates the marginal probability for every state of the graphical model. This will require using the _sum-product algorithm_ and passing all messages on the graph.\n",
    "\n",
    "Here is the wikipedia page for reference: https://en.wikipedia.org/wiki/Belief_propagation (not the greatest, but not horrible)\n",
    "\n",
    "Hints:\n",
    " * The order of the variables above is such that each variable is dependent only on variables that occur before it. This should save you some hassle in terms of calculating a message passing order.\n",
    " * You may want to add code before the function to prepare. This might involve creating a list of edges in the graphical model, in the order they need to be processed for instance. It can also be done with dictionaries or classes - choose whatever you are comfortable with, but spend time thinking it through first. If you get it wrong your code will get messy!\n",
    " * A good approach to recording messages is a dictionary indexed `[from, to]`\n",
    " * The easiest way to include a known value in the model is often to add/replace the unary term for that value.\n",
    " * This problem is small enough that you can brute force it - code to do so has been provided above. Do test that your implementation is correct by comparing them.\n",
    "\n",
    "__(10 marks)__\n",
    " * 3 marks for sensible data structures\n",
    " * 5 marks for sending the messages correctly\n",
    " * 2 marks for correctly calculating the marginals at the end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def marginals(known):\n",
    "    \"\"\"known is a dictionary, where a random variable index existing as a key in the dictionary\n",
    "    indicates it has been observed. The value obtained using the key is the value the\n",
    "    random variable has been observed as. Returns a 17x2 matrix, such that [rv,0] is the\n",
    "    probability of random variable rv being False, [rv, 1] the probability of being True.\"\"\"\n",
    "    \n",
    "    # **************************************************************** 10 marks\n",
    "    \n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. What's broken?\n",
    "\n",
    "Simply print out what the most probable broken part of each of the below machines is. Please print it as\n",
    "`dictionary key: name of broken part`, e.g. `A: wr`. This will allow automarking to be used.\n",
    "\n",
    "If you can't get your belief propagation implementation to work feel free to use `brute_marginals` instead.\n",
    "\n",
    "__(1 mark)__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "repair = {}\n",
    "repair['A'] = {nti['me'] : True}\n",
    "repair['B'] = {nti['wv'] : True}\n",
    "repair['C'] = {nti['wv'] : False, nti['lo'] : True}\n",
    "repair['D'] = {nti['hp'] : False, nti['ta'] : False}\n",
    "repair['E'] = {nti['hp'] : True, nti['wv'] : True, nti['vp']: True}\n",
    "\n",
    "\n",
    "\n",
    "# **************************************************************** 1 mark\n",
    "\n",
    "\n"
   ]
  }
 ],
 "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": 2
}
