lab_2/12402559_12141666_lab2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "3a8eef87c3a3f57f3060a92e8a0535c4",
     "grade": false,
     "grade_id": "cell-66ada6339bda77a6",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### Lab 2: Inference in Graphical Models\n",
    "\n",
    "### Machine Learning 2 (2019)\n",
    "\n",
    "* The lab exercises can be done in groups of two people, or individually.\n",
    "* The deadline is Tuesday October 1st 2019 at 16:59.\n",
    "* Assignment should be submitted through Canvas! Make sure to include your and your teammates' names with the submission.\n",
    "* Attach the .IPYNB (IPython Notebook) file containing your code and answers. Naming of the file should be \"studentid1\\_studentid2\\_lab#\", for example, the attached file should be \"12345\\_12346\\_lab1.ipynb\". Only use underscores (\"\\_\") to connect ids, otherwise the files cannot be parsed.\n",
    "\n",
    "Notes on implementation:\n",
    "\n",
    "* You should write your code and answers in an IPython Notebook: http://ipython.org/notebook.html. If you have problems, please ask.\n",
    "* Use __one cell__ for code and markdown answers only!\n",
    "    * Put all code in the cell with the ```# YOUR CODE HERE``` comment and overwrite the ```raise NotImplementedError()``` line.\n",
    "    * For theoretical questions, put your solution using LaTeX style formatting in the YOUR ANSWER HERE cell.\n",
    "* Among the first lines of your notebook should be \"%pylab inline\". This imports all required modules, and your plots will appear inline.\n",
    "* Large parts of you notebook will be graded automatically. Therefore it is important that your notebook can be run completely without errors and within a reasonable time limit. To test your notebook before submission, select Kernel -> Restart \\& Run All."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "2a0482dac028fae34551f9e75cc82068",
     "grade": false,
     "grade_id": "cell-7c4914def40aeb03",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### Introduction\n",
    "In this assignment, we will implement the sum-product and max-sum algorithms for factor graphs over discrete variables. The relevant theory is covered in chapter 8 of Bishop's PRML book, in particular section 8.4. Read this chapter carefuly before continuing!\n",
    "\n",
    "We will implement sum-product and max-sum and apply it to a poly-tree structured medical diagnosis example.\n",
    "\n",
    "For this assignment you should use numpy ndarrays (constructed with np.array, np.zeros, np.ones, etc.). We need n-dimensional arrays in order to store conditional distributions with more than one conditioning variable. If you want to perform matrix multiplication on arrays, use the np.dot function; all infix operators including *, +, -, work element-wise on arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "b85b4dd3e9c765f6ed0d81b98161f0cd",
     "grade": false,
     "grade_id": "cell-7f9d339abedb92d3",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "## Part 0: Doing the math (5 points)\n",
    "We start with a set of three Bernoulli distributed variables X, Y, Z. Calculate the marginals $p(Y=1)$ and $p(Y=0)$ using only the sum and product rule where,\n",
    "\n",
    "$$\n",
    "p(X=1) = 0.05 \\\\\\\\\n",
    "p(Z=1) = 0.2 \\\\\\\\\n",
    "$$\n",
    "$$\n",
    "p(Y = 1 | X = 1, Z = 1) = 0.99 \\\\\\\\\n",
    "p(Y = 1 | X = 1, Z = 0) = 0.9 \\\\\\\\\n",
    "p(Y = 1 | X = 0, Z = 1) = 0.7 \\\\\\\\\n",
    "p(Y = 1 | X = 0, Z = 0) = 0.0001 \\\\\\\\\n",
    "$$\n",
    "\n",
    "While implementing the message passing algorithms you should be able to use the results of this question as a guidance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "cafc6734ba3ef70e630191e714f3c0b3",
     "grade": true,
     "grade_id": "cell-653065eb00c64d46",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "source": [
    "$\\normalsize  \\color{blue}{\\textbf{# YOUR ANSWER HERE}}$\n",
    "***\n",
    "\n",
    "Given that $p(X=0) = 1-p(X=1) = 0.995$ and $p(Z=0) = 0.8$ we have that\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "p(Y=1) \n",
    "&= \\sum_X \\sum_Z p(Y=1, X, Z) \n",
    "\\\\&= \\sum_X \\sum_Z p(Y=1| X, Z) p(X) p(Z)\n",
    "\\\\&= 0.2 * 0.99 * 0.05 + 0.05 * 0.8 * 0.9 + 0.95 * 0.8 * 0.0001 + 0.95 * 0.2 * 0.7\n",
    "\\\\&= 0.178976\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Thus $p(Y = 0) = 1 - p(Y = 1) = 0.821024$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "28024e8ece48c8b90f5fa0fdd0c13a6e",
     "grade": false,
     "grade_id": "cell-9afa97d7b5de1808",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "## Part 1: The sum-product algorithm\n",
    "\n",
    "We will implement a data structure to store a factor graph and to facilitate computations on this graph. Recall that a factor graph consists of two types of nodes, factors and variables. Below you will find some classes for these node types to get you started. Carefully inspect this code and make sure you understand what it does. Step by step will update its functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "1affe45cf22d1acc9bdb7a7994d07c6a",
     "grade": false,
     "grade_id": "cell-70843f7b3041a417",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Populating the interactive namespace from numpy and matplotlib\n"
     ]
    }
   ],
   "source": [
    "%pylab inline\n",
    "\n",
    "class Node(object):\n",
    "    \"\"\"\n",
    "    Base-class for Nodes in a factor graph. Only instantiate sub-classes of Node.\n",
    "    \"\"\"\n",
    "    def __init__(self, name):\n",
    "        # A name for this Node, for printing purposes\n",
    "        self.name = name\n",
    "        \n",
    "        # Neighbours in the graph, identified with their index in this list.\n",
    "        # i.e. self.neighbours contains neighbour 0 through len(self.neighbours) - 1.\n",
    "        self.neighbours = []\n",
    "        \n",
    "        # Reset the node-state (not the graph topology)\n",
    "        self.reset()\n",
    "        \n",
    "    def reset(self):\n",
    "        # Incomming messages; a dictionary mapping neighbours to messages.\n",
    "        # That is, it maps  Node -> np.ndarray.\n",
    "        self.in_msgs = {}\n",
    "        \n",
    "        # A set of neighbours for which this node has pending messages.\n",
    "        # We use a python set object so we don't have to worry about duplicates.\n",
    "        self.pending = set([])\n",
    "\n",
    "    def add_neighbour(self, nb):\n",
    "        self.neighbours.append(nb)\n",
    "\n",
    "    def send_sp_msg(self, other):\n",
    "        # To be implemented in subclass.\n",
    "        raise Exception('Method send_sp_msg not implemented in base-class Node')\n",
    "   \n",
    "    def send_ms_msg(self, other):\n",
    "        # To be implemented in subclass.\n",
    "        raise Exception('Method send_ms_msg not implemented in base-class Node')\n",
    "    \n",
    "    def receive_msg(self, other, msg):\n",
    "        # Store the incomming message, replacing previous messages from the same node\n",
    "        self.in_msgs[other] = msg\n",
    "\n",
    "        # TODO: add pending messages\n",
    "        # self.pending.update(...)\n",
    "    \n",
    "    def __str__(self):\n",
    "        # This is printed when using 'print node_instance'\n",
    "        return self.name\n",
    "\n",
    "\n",
    "class Variable(Node):\n",
    "    def __init__(self, name, num_states):\n",
    "        \"\"\"\n",
    "        Variable node constructor.\n",
    "        Args:\n",
    "            name: a name string for this node. Used for printing. \n",
    "            num_states: the number of states this variable can take.\n",
    "            Allowable states run from 0 through (num_states - 1).\n",
    "            For example, for a binary variable num_states=2,\n",
    "            and the allowable states are 0, 1.\n",
    "        \"\"\"\n",
    "        self.num_states = num_states\n",
    "        \n",
    "        # Call the base-class constructor\n",
    "        super(Variable, self).__init__(name)\n",
    "    \n",
    "    def set_observed(self, observed_state):\n",
    "        \"\"\"\n",
    "        Set this variable to an observed state.\n",
    "        Args:\n",
    "            observed_state: an integer value in [0, self.num_states - 1].\n",
    "        \"\"\"\n",
    "        # Observed state is represented as a 1-of-N variable\n",
    "        # Set all-but-one states to a very low probability;\n",
    "        # (a bit of a hack to avoid -inf values when taking logs)\n",
    "        self.observed_state[:] = 1e-10\n",
    "        self.observed_state[observed_state] = 1.0\n",
    "        \n",
    "    def set_latent(self):\n",
    "        \"\"\"\n",
    "        Erase an observed state for this variable and consider it latent again.\n",
    "        \"\"\"\n",
    "        # No state is preferred, so set all entries of observed_state to 1.0\n",
    "        # Using this representation we need not differentiate observed an latent\n",
    "        # variables when sending messages.\n",
    "        self.observed_state[:] = 1.0\n",
    "        \n",
    "    def reset(self):\n",
    "        super(Variable, self).reset()\n",
    "        self.observed_state = np.ones(self.num_states)\n",
    "        \n",
    "    def marginal(self):\n",
    "        \"\"\"\n",
    "        Compute the marginal distribution of this Variable.\n",
    "        It is assumed that message passing has completed when this function is called.\n",
    "        \"\"\"\n",
    "        # TODO: compute marginal\n",
    "        return None\n",
    "    \n",
    "    def unnormalized_log_marginal(self):\n",
    "        \"\"\"\n",
    "        Compute the unnormalized log marginal distribution of this Variable.\n",
    "        It is assumed that message passing has completed when this function is called.\n",
    "        \"\"\"\n",
    "        # TODO: compute unnormalized log marginal\n",
    "        return None\n",
    "    \n",
    "    def send_sp_msg(self, other):\n",
    "        # TODO: implement Variable -> Factor message for sum-product\n",
    "        pass\n",
    "   \n",
    "    def send_ms_msg(self, other):\n",
    "        # TODO: implement Variable -> Factor message for max-sum\n",
    "        pass\n",
    "\n",
    "class Factor(Node):\n",
    "    def __init__(self, name, f, neighbours):\n",
    "        \"\"\"\n",
    "        Factor node constructor.\n",
    "        Args:\n",
    "            name: a name string for this node. Used for printing\n",
    "            f: a numpy.ndarray with N axes, where N is the number of neighbours.\n",
    "               That is, the axes of f correspond to variables, and the index along that axes corresponds to a value of that variable.\n",
    "               Each axis of the array should have as many entries as the corresponding neighbour variable has states.\n",
    "            neighbours: a list of neighbouring Variables. Bi-directional connections are created.\n",
    "        \"\"\"\n",
    "        # Call the base-class constructor\n",
    "        super(Factor, self).__init__(name)\n",
    "        \n",
    "        f = np.array(f)  # For convenience, accept lists (of lists (of numbers or potentially lists of ...))\n",
    "        \n",
    "        assert len(neighbours) == f.ndim, 'Factor function f should accept as many arguments as this Factor node has neighbours'\n",
    "        \n",
    "        for nb_ind in range(len(neighbours)):\n",
    "            nb = neighbours[nb_ind]\n",
    "            assert f.shape[nb_ind] == nb.num_states, 'The range of the factor function f is invalid for input %i %s' % (nb_ind, nb.name)\n",
    "            self.add_neighbour(nb)\n",
    "            nb.add_neighbour(self)\n",
    "\n",
    "        self.f = f\n",
    "        \n",
    "    def send_sp_msg(self, other):\n",
    "        # TODO: implement Factor -> Variable message for sum-product\n",
    "        pass\n",
    "   \n",
    "    def send_ms_msg(self, other):\n",
    "        # TODO: implement Factor -> Variable message for max-sum\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "24aa9afccd84106fdc0727236c1cb186",
     "grade": false,
     "grade_id": "cell-e0ada624a5583fce",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 1.1. Initialize the graph (5 points)\n",
    "The equations in Part 0 can be represented by a factor graph. Instantiate this graph by creating Variable and Factor instances and linking them according to the graph structure. \n",
    "To instantiate the factor graph, first create the Variable nodes and then create Factor nodes, passing a list of neighbour Variables to each Factor. To get you started, we initialize the variable node for $X$ and the factor node corresponding to the prior $p(X)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "e8e1e5033150190f7ba288f5a2b004b4",
     "grade": false,
     "grade_id": "cell-85011c26139bed2e",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "X = Variable(name='X', num_states=2)\n",
    "X_prior = Factor(name='p(X)',\n",
    "                 f=np.array([0.95, 0.05]),\n",
    "                 neighbours=[X])\n",
    "\n",
    "# Please stick to the naming convention used below, otherwise the test functionality throughout the lab won't work                       \n",
    "# YOUR CODE HERE\n",
    "Z = Variable(name='Z', num_states=2)\n",
    "Z_prior = Factor(name='p(Z)',\n",
    "                 f=np.array([0.8, 0.2]),\n",
    "                 neighbours=[Z])\n",
    "\n",
    "                 \n",
    "Y = Variable(name='Y', num_states=2)\n",
    "Y_cond = Factor(name='p(Y|X,Z)',\n",
    "                f=np.array([[[0.9999, 0.3], [ 0.1, 0.01]], [[0.0001, 0.7], [0.9, 0.99]]]),\n",
    "                neighbours=[Y, X, Z])\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "1d247840a0cad095db98f369d0531275",
     "grade": true,
     "grade_id": "cell-1bf4d184d405f181",
     "locked": true,
     "points": 5,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "assert Z_prior.f.shape == (2,)\n",
    "assert Y_cond.f.shape == (2, 2, 2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "09f3f672778ffb77e22ba8702e9a2911",
     "grade": false,
     "grade_id": "cell-af817385e2927300",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "We will be doing a lot of marginalizations, i.e. obtain the distribution of a *single* variable from a distribution over multiple variables. Let's first write a function `marginalize` for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d81924d48029a907c55cea627ae5b56f",
     "grade": false,
     "grade_id": "cell-74ecd6f85f99a2d5",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.48685075 0.51314925]\n",
      "[0.36854731 0.36324221 0.26821048]\n",
      "[0.19012979 0.26282865 0.3301945  0.21684706]\n"
     ]
    }
   ],
   "source": [
    "# Write a function marginalize that given\n",
    "def marginalize(P, dim):\n",
    "    # YOUR CODE HERE\n",
    "    \"\"\"\n",
    "    Every dimention represents a random variable\n",
    "    and its length the number of possible values.\n",
    "    \n",
    "    Should return a list with the marginal\n",
    "    \n",
    "    Sum along all the other dimentions\n",
    "    \"\"\"\n",
    "    sum_axes = set(np.arange(len(P.shape))).difference({dim})\n",
    "    return P.sum(axis=tuple(sum_axes))\n",
    "\n",
    "\n",
    "# Lets try it\n",
    "test_P = np.random.rand(2, 3, 4)\n",
    "test_P = test_P / test_P.sum()  # Normalize for proper distribution\n",
    "\n",
    "# Do the marginal distributions look like you expect?\n",
    "print (marginalize(test_P, 0))\n",
    "print (marginalize(test_P, 1))\n",
    "print (marginalize(test_P, 2))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "83cf7c2b9bf87fc4a0268c36436195f6",
     "grade": false,
     "grade_id": "cell-56015dc94e6caaf4",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 1.2 Factor to variable messages (20 points)\n",
    "Write a method `send_sp_msg(self, other)` for the Factor class, that checks if all the information required to pass a message to `Variable` `other` is present, computes the message and sends it to `other`. \"Sending\" here simply means calling the `receive_msg` function of the receiving node (we will implement this later). The message itself should be represented as a numpy array (np.array) whose length is equal to the number of states of the variable.\n",
    "\n",
    "In the very end of 1.2 below we overwrite `send_sp_msg(self, other)` for the Factor class. In general, this is considered bad practise but in this lab it saves us from scrolling up and down all the time.\n",
    "\n",
    "You will implement a function `send_sp_msg` that sends a message from a factor to a variable for the max-sum algorith. This function implements Equation 8.66 from Bishop. The message should be a numpy array (np.array) whose length is equal to the number of states of the variable.\n",
    "\n",
    "It is a good idea to write a number of helper functions to implement these equations.\n",
    "\n",
    "Since it is always a good idea to include checks, you should first write a method `can_send_message` that checks whether a node `node` has all the information required to pass a message to `other`. This should work for both variable and factor nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "8b75b5269753ba1861f6c7dbd41fa1a5",
     "grade": false,
     "grade_id": "cell-4164822e95a6b68a",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "False\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "def can_send_message(sender, receiver):\n",
    "    # YOUR CODE HERE\n",
    "    return (set(sender.in_msgs).difference({receiver}) == set(sender.neighbours).difference({receiver}))\n",
    "\n",
    "\n",
    "# Do the results make sense?\n",
    "print (can_send_message(X, X_prior))\n",
    "print (can_send_message(X_prior, X))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "15c446225bb0473c69caa2219f4c80bd",
     "grade": false,
     "grade_id": "cell-816ded386be1b608",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "\n",
    "In Eq. 8.66, Bishop writes $f(x, x_1, ..., x_M)$, where $x$ corresponds to the variable that will receive the message. For now, assume this variable is the `index`-th neighbour of the factor. In order to ease implementation, it may be a good idea to write a function that rearanges the dimensions of `f` to comply with this notation, i.e. moves the dimension `index` to the front. Make sure to return a copy and keep all other dimensions in order! Use `np.moveaxis`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "296bbf4e2783a797fe753c30c7ad7877",
     "grade": false,
     "grade_id": "cell-702210e632c423bd",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def move_dimension_first(f, index):\n",
    "    # YOUR CODE HERE\n",
    "    return np.moveaxis(f, index, 0).copy()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should calculate the product of the incoming messages of all neighbours of the sending node except the receiving node. Therefore it may be useful to write a function `get_neighbour_messages` that gathers these messages in a list. If you want to not make things complicated, make sure the order of the messages in the list corresponds to the order of the variables in `neighbours`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "aec4aebe3b5cc9ee35ae118c1763ff26",
     "grade": false,
     "grade_id": "cell-9e7f7986263e1543",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def get_neighbour_messages(sender, receiver):\n",
    "    # YOUR CODE HERE\n",
    "    return [sender.in_msgs[neighbor] for neighbor in sender.neighbours if neighbor is not receiver]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before marginalizing, we need to calculate $\\prod_{m\\in\\text{ne}(f_s)\\setminus x} \\mu_{x_m\\rightarrow f_s}(x_m)$ (Eq. 8.66) for all possible combinations of $x_1, ..., x_M$ (values of the neighbour nodes except the receiving node). An elegant and efficient way to calculate these is using the n-way outer product of vectors. This product takes n vectors $\\mathbf{x}^{(1)}, \\ldots, \\mathbf{x}^{(n)}$ and computes a $n$-dimensional tensor (ndarray) whose element $i_0,i_1,...,i_n$ is given by $\\prod_j \\mathbf{x}^{(j)}_{i_j}$. In python, this is realized as `np.multiply.reduce(np.ix_(*vectors))` for a python list `vectors` of 1D numpy arrays. Try to figure out how this statement works -- it contains some useful functional programming techniques. What should `vectors` be? Try to see the link between the result and Eq. 8.66."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "e5e5d3c0d6252b785451adf1c41bd5fc",
     "grade": false,
     "grade_id": "cell-67269e06a554035c",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def calc_other_neighbour_msg_prod(sender, receiver):\n",
    "    # YOUR CODE HERE\n",
    "    msgs_vector = get_neighbour_messages(sender, receiver)\n",
    "    return np.multiply.reduce(np.ix_(*msgs_vector))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "00442c496ed4ebda7eff1672a64c9f59",
     "grade": false,
     "grade_id": "cell-d5366f0c524b30d0",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "Following Eq. 8.66, before marginalizing, you should calculate the product of $f(x, x_1, ..., x_M)$ with $\\prod_{m\\in\\text{ne}(f_s)\\setminus x} \\mu_{x_m\\rightarrow f_s}(x_m)$ for all configurations of $x, x_1, ..., x_M$. Since the second part does not depend on $x$, its tensor representations are of different dimensions. You can overcome this problem by using a loop, but preferably use numpy broadcasting by first aligning the dimensions of the tensors. You can use `np.expand_dims` or `X[None, ...]` to insert one dimension at the front. Write a function `calculate_factor` that, given `f` (which is reordered such that $x$ corresponds to the first dimension) and the (outer) product of the other neighbour messages, computes $f(x, x_1, ..., x_M) \\prod_{m\\in\\text{ne}(f_s)\\setminus x} \\mu_{x_m\\rightarrow f_s}(x_m)$ for all configurations of $x, x_1, ..., x_M$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "3ab59a95b589ff71451036c73b289138",
     "grade": false,
     "grade_id": "cell-d34359928f3a997d",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def calculate_factor(f_neighb_first, neighbour_msg_prod):\n",
    "    # YOUR CODE HERE\n",
    "    return f_neighb_first * neighbour_msg_prod\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "ef9e90a515d362d281f313955fbbc75f",
     "grade": false,
     "grade_id": "cell-d9b12fad1b574924",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "Put all the pieces together to define a function `calc_sum_product_factor_to_variable_msg` that calculates Eq. 8.66."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "f3a2dc41ed909c8d0eb88b213ce66a99",
     "grade": false,
     "grade_id": "cell-ec6b421860d40ae0",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def calc_sum_product_factor_to_variable_msg(factor, variable):\n",
    "    # YOUR CODE HERE\n",
    "    idx = factor.neighbours.index(variable)\n",
    "\n",
    "    f_neighb_first = move_dimension_first(factor.f, idx)\n",
    "    neighbour_msg_prod = calc_other_neighbour_msg_prod(factor, variable)\n",
    "    f =  calculate_factor(f_neighb_first, neighbour_msg_prod)\n",
    "\n",
    "    return  marginalize(f , 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "c65b2bb803c0f8b53c5eb427e8bb2906",
     "grade": false,
     "grade_id": "cell-7d334ab3a40918b0",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Finally, we will define the send message function for you\n",
    "def factor_send_sp_msg(self, variable):\n",
    "    \n",
    "    assert isinstance(variable, Variable), \"Factor can only send messages to variable!\"\n",
    "    assert can_send_message(self, variable), \"Cannot send message!\"\n",
    "    \n",
    "    out_msg = calc_sum_product_factor_to_variable_msg(self, variable)\n",
    "    \n",
    "    # Send the message\n",
    "    variable.receive_msg(self, out_msg)\n",
    "    \n",
    "    # Remove the pending sign if present\n",
    "    self.pending.discard(variable)\n",
    "    \n",
    "Factor.send_sp_msg = factor_send_sp_msg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "4ddb724608cd7205b51cbadda1a3664f",
     "grade": true,
     "grade_id": "cell-429fba08a94a3bd2",
     "locked": true,
     "points": 20,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "# message from X_prior to X\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "\n",
    "X_prior.send_sp_msg(X)\n",
    "assert np.allclose(list(X.in_msgs.values()), [0.95, 0.05])\n",
    "\n",
    "# message from Z_prior to Z\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "\n",
    "Z_prior.send_sp_msg(Z)\n",
    "assert np.allclose(list(Z.in_msgs.values()), [0.8, 0.2])\n",
    "\n",
    "# message from Y_cond to Y\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "Y_cond.receive_msg(X, X_prior.f) # simulating that Y_cond received all necessary messages from X\n",
    "Y_cond.receive_msg(Z, Z_prior.f) # simulating that Y_cond received all necessary messages from Z\n",
    "Y_cond.send_sp_msg(Y)\n",
    "assert np.allclose(list(Y.in_msgs.values()), [0.821024, 0.178976])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "2fd488d52a30d4081cdbc910c280a279",
     "grade": false,
     "grade_id": "cell-2ac03f39f1b1281c",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 1.3 Variable to factor messages (10 points)\n",
    "\n",
    "Write a method `calc_sum_product_variable_to_factor_msg(variable, factor)` that computes the message to be sent to a neighbour variable by a factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "5c3c49af77b02653acaac2455d864132",
     "grade": false,
     "grade_id": "cell-93d0a0d839eb99f3",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def calc_sum_product_variable_to_factor_msg(variable, factor):\n",
    "    # YOUR CODE HERE\n",
    "    neighbour_msg = get_neighbour_messages(variable, factor)\n",
    "\n",
    "    if len(neighbour_msg) > 0:\n",
    "        return np.prod(np.array(neighbour_msg), axis=0)\n",
    "    else:\n",
    "        return np.ones(variable.num_states)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "bb0a564a418df48edc0a627ff3df1bc4",
     "grade": false,
     "grade_id": "cell-8ce3afc9581048ad",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Finally, we will define the send message function for you\n",
    "def variable_send_sp_msg(self, factor):\n",
    "    \n",
    "    assert isinstance(factor, Factor), \"Variable can only send messages to factor!\"\n",
    "    assert can_send_message(self, factor), \"Cannot send message!\"\n",
    "    \n",
    "    out_msg = calc_sum_product_variable_to_factor_msg(self, factor)\n",
    "    \n",
    "    # Send the message\n",
    "    factor.receive_msg(self, out_msg)\n",
    "    \n",
    "    # Remove the pending sign if present\n",
    "    self.pending.discard(factor)\n",
    "    \n",
    "Variable.send_sp_msg = variable_send_sp_msg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "80a738a13d17bc6014f56719bc935ca4",
     "grade": true,
     "grade_id": "cell-0cbe0bfa635450f3",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "# First message from X to Y_cond\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "X_prior.send_sp_msg(X) # simulating that X received all necessary messages\n",
    "X.send_sp_msg(Y_cond) \n",
    "assert np.allclose(list(Y_cond.in_msgs.values()), [0.95, 0.05])\n",
    "\n",
    "# Second message from Z to Y_cond\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "Z_prior.send_sp_msg(Z) # simulating that Z received all necessary messages\n",
    "Z.send_sp_msg(Y_cond)\n",
    "assert np.allclose(list(Y_cond.in_msgs.values()), [[0.95, 0.05], [0.8, 0.2]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "72a4f47ec3cd02e4e32b7d68e9d0692b",
     "grade": false,
     "grade_id": "cell-6346fc68da031348",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### Testing a single forward pass and a single backward pass\n",
    "Before we go on we will make sure that messages can be passed in both directions the graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d50a8fd57d73f85f1c8f81be1ba13f69",
     "grade": false,
     "grade_id": "cell-f44819044a283ec9",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "X_prior.reset()\n",
    "X.reset()\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "# Forward pass\n",
    "X_prior.send_sp_msg(X)\n",
    "Z_prior.send_sp_msg(Z)\n",
    "X.send_sp_msg(Y_cond)\n",
    "Z.send_sp_msg(Y_cond)\n",
    "Y_cond.send_sp_msg(Y)\n",
    "assert np.allclose(list(Y.in_msgs.values()), [0.821024, 0.178976])\n",
    "\n",
    "# Backward pass\n",
    "Y.send_sp_msg(Y_cond)\n",
    "Y_cond.send_sp_msg(X)\n",
    "Y_cond.send_sp_msg(Z)\n",
    "X.send_sp_msg(X_prior)\n",
    "Z.send_sp_msg(Z_prior)\n",
    "assert np.allclose(list(X.in_msgs.values()), [[0.95, 0.05],[1., 1.]])\n",
    "assert np.allclose(list(Z.in_msgs.values()), [[0.8, 0.2],[1., 1.]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "27ce20145a58d03708168b08fbbfeadd",
     "grade": false,
     "grade_id": "cell-f1b2c01d81e8d007",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 1.4 Compute marginal (10 points)\n",
    "Later in this assignment, we will implement message passing schemes to do inference. Once the message passing has completed, we will want to compute local marginals for each variable.\n",
    "Write the method `marginal` for the `Variable` class, that computes a marginal distribution over that node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "8002994abb8ad6ada1ad19a0b433c273",
     "grade": false,
     "grade_id": "cell-0f6434f5c5d4dc9c",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def marginal(self):\n",
    "    # YOUR CODE HERE    \n",
    "    marginal = np.prod(np.array(list(self.in_msgs.values())), axis=0)\n",
    "    marginal /= marginal.sum()\n",
    "    return marginal\n",
    "\n",
    "Variable.marginal = marginal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "8dc0f9a18f2801cf7f3182732f45b52e",
     "grade": true,
     "grade_id": "cell-7f666a68bbbe81e5",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "# Simulate a single forward pass\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "X_prior.send_sp_msg(X)\n",
    "Z_prior.send_sp_msg(Z)\n",
    "X.send_sp_msg(Y_cond)\n",
    "Z.send_sp_msg(Y_cond)\n",
    "Y_cond.send_sp_msg(Y)\n",
    "\n",
    "assert np.allclose(X.marginal(), [0.95, 0.05])\n",
    "assert np.allclose(Z.marginal(), [0.8, 0.2])\n",
    "assert np.allclose(Y.marginal(), [0.821024, 0.178976])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "3d2d9dd5139b1ed6891107868305e57b",
     "grade": false,
     "grade_id": "cell-e38ca86d50f6c9c1",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 1.5 Receiving messages (10 points)\n",
    "In order to implement a message passing algorithms, we need some way to determine which nodes are ready to send messages to which neighbours. We make use of the concept of \"pending messages\", which is explained in Bishop (8.4.7): \n",
    "\"we will say that a (variable or factor)\n",
    "node a has a message pending on its link to a node b if node a has received any\n",
    "message on any of its other links since the last time it sent a message to b. Thus,\n",
    "when a node receives a message on one of its links, this creates pending messages\n",
    "on all of its other links.\"\n",
    "\n",
    "Before we say node a has a pending message for node b, we **must check that node a has received all messages needed to compute the message that is to be sent to b**.\n",
    "\n",
    "Modify the function `receive_msg`, so that it updates the self.pending variable as described above. The member self.pending is a set that is to be filled with Nodes to which self has pending messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d4256320e0a6fff7550f0a0c744ddb61",
     "grade": false,
     "grade_id": "cell-39b7b3a84cdf1e56",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "# ANSWER 1.5\n",
    "def receive_msg(self, other, msg):\n",
    "    self.in_msgs[other] = msg\n",
    "    # YOUR CODE HERE\n",
    "    self.pending.update(set([neighbour for neighbour in self.neighbours if can_send_message(self, neighbour)]))\n",
    "\n",
    "Node.receive_msg = receive_msg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "05683e9bfc7be13ffac9b88806845c3b",
     "grade": true,
     "grade_id": "cell-42c321c98bf9b6b1",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "assert X_prior.pending == set()\n",
    "\n",
    "X_prior.pending.add(X)\n",
    "assert str(list(X_prior.pending)[0]) == X.name\n",
    "\n",
    "X_prior.send_sp_msg(X)\n",
    "assert X_prior.pending == set()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.6 Inference Engine (10 points)\n",
    "Write a function `sum_product(node_list)` that runs the sum-product message passing algorithm on a tree-structured factor graph with given nodes. The input parameter `node_list` is a list of all Node instances in the graph, which is assumed to be ordered correctly. That is, the list starts with a leaf node, which can always send a message. Subsequent nodes in `node_list` should be capable of sending a message when the pending messages of preceding nodes in the list have been sent. The sum-product algorithm then proceeds by passing over the list from beginning to end, sending all pending messages at the nodes it encounters. Then, in reverse order, the algorithm traverses the list again and again sends all pending messages at each node as it is encountered. For this to work, we initialized pending messages for all the leaf nodes, e.g. `X_prior.pending.add(X)`, where `X_prior` is a Factor node corresponding the the prior, `X` is a `Variable` node and the only connection of `X_prior` goes to `X`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "0295fa5b1b0d1c2dd532935bd543056f",
     "grade": false,
     "grade_id": "cell-6a6fb2568d840c68",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def sum_product(node_list):\n",
    "    # YOUR CODE HERE\n",
    "    # forward\n",
    "    for node_a in node_list:\n",
    "        for node_b in list(node_a.pending):\n",
    "            node_a.send_sp_msg(node_b)\n",
    "    # reversed\n",
    "    for node_a in reversed(node_list):\n",
    "        for node_b in list(node_a.pending):\n",
    "            node_a.send_sp_msg(node_b)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "261e257a05c618b4e1c5d9568f7403d5",
     "grade": true,
     "grade_id": "cell-fb8315046f8fde42",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "nodes = [X_prior, X, Z_prior, Z, Y_cond, Y]\n",
    "for n in nodes:\n",
    "    n.reset()\n",
    "    \n",
    "X_prior.pending.add(X)\n",
    "Z_prior.pending.add(Z)\n",
    "Y.pending.add(Y_cond)\n",
    "\n",
    "sum_product(nodes)\n",
    "assert np.allclose(Y.marginal(), [0.821024, 0.178976])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.7 Observed variables and probabilistic queries (15 points)\n",
    "We will now use the inference engine to answer probabilistic queries. That is, we will set certain variables to observed values, and obtain the marginals over latent variables. We have already provided functions `set_observed` and `set_latent` that manage a member of Variable called `observed_state`. Modify the `calc_sum_product_variable_to_factor_msg` and `Variable.marginal` routines that you wrote before, to use `observed_state` so as to get the required marginals when some nodes are observed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "da29573eaddaf62947074fc00f714df1",
     "grade": false,
     "grade_id": "cell-55c50fb9c1836789",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def calc_sum_product_variable_to_factor_msg(variable, factor):\n",
    "    # YOUR CODE HERE\n",
    "    neighbour_msg = get_neighbour_messages(variable, factor)\n",
    "\n",
    "    if len(neighbour_msg) > 0:\n",
    "        msg = np.prod(np.array(neighbour_msg), axis=0)\n",
    "    else:\n",
    "        msg = np.ones(variable.num_states)\n",
    "    \n",
    "    return msg * variable.observed_state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d27bd6da3b20252943b2828ceb5bb472",
     "grade": true,
     "grade_id": "cell-13801e018b6f9d27",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test, test, test\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "Y_cond.reset()\n",
    "\n",
    "X_prior.send_sp_msg(X)\n",
    "X.set_observed(0)\n",
    "X.send_sp_msg(Y_cond)\n",
    "assert np.allclose(list(Y_cond.in_msgs.values()), [9.5e-01, 5.0e-12])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "55a9272099aabff3579cac4f41f08973",
     "grade": false,
     "grade_id": "cell-d0b91da808874e73",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def marginal(self):\n",
    "    # YOUR CODE HERE\n",
    "#     marginal = np.multiply.reduce(np.ix_(*[*self.in_msgs.values()]))\n",
    "    marginal = np.prod(np.array(list(self.in_msgs.values())), axis=0)\n",
    "    marginal /= np.sum(marginal)\n",
    "    return marginal*self.observed_state\n",
    "    \n",
    "Variable.marginal = marginal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "fa342104c5f9d55b65ad4c202475a2c5",
     "grade": true,
     "grade_id": "cell-2f8e5d9052ed73ef",
     "locked": true,
     "points": 5,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test, test, test\n",
    "# Simulate a single forward pass\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "X.set_observed(0)\n",
    "Z.set_observed(0)\n",
    "\n",
    "X_prior.send_sp_msg(X)\n",
    "Z_prior.send_sp_msg(Z)\n",
    "X.send_sp_msg(Y_cond)\n",
    "Z.send_sp_msg(Y_cond)\n",
    "Y_cond.send_sp_msg(Y)\n",
    "\n",
    "assert np.allclose(Y.marginal(), [9.99900000e-01, 1.00000022e-04])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "97cce5eb6885d02431cf74e2affba069",
     "grade": false,
     "grade_id": "cell-2277868e4ec25b1d",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "## Part 2: The max-sum algorithm\n",
    "Next, we implement the max-sum algorithm as described in section 8.4.5 of Bishop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "34cda2c1859ee99af59b1e6f6609b5b6",
     "grade": false,
     "grade_id": "cell-f1dbc98349fb824f",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 2.1 Factor to variable messages (10 points)\n",
    "Implement the function `Factor.send_ms_msg` that sends Factor -> Variable messages for the max-sum algorithm. It is analogous to the `Factor.send_sp_msg` function you implemented before. Make sure it works for observed and unobserved nodes. Consider using a number of helper functions as seen in Part 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "c16471807291ab66dbe760bddc2f1480",
     "grade": false,
     "grade_id": "cell-389621edefd825a3",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def marginalize_max(P, dim):\n",
    "    sum_axes = set(np.arange(len(P.shape))).difference({dim})\n",
    "    return P.max(axis=tuple(sum_axes))\n",
    "\n",
    "def calc_other_neighbour_msg_sum(sender, receiver):\n",
    "    msgs_vector = get_neighbour_messages(sender, receiver)\n",
    "    return np.add.reduce(np.ix_(*msgs_vector))\n",
    "\n",
    "def calculate_factor_max(f_neighb_first, neighbour_msg_prod):\n",
    "    return np.log(f_neighb_first) + neighbour_msg_prod\n",
    "\n",
    "def calc_max_sum_factor_to_variable_msg(factor, variable):\n",
    "    idx = factor.neighbours.index(variable)\n",
    "\n",
    "    f_neighb_first = move_dimension_first(factor.f, idx)\n",
    "    neighbour_msg_prod = calc_other_neighbour_msg_sum(factor, variable)\n",
    "    f =  calculate_factor_max(f_neighb_first, neighbour_msg_prod)\n",
    "\n",
    "    return  marginalize_max(f , 0)\n",
    "\n",
    "def factor_send_ms_msg(self, variable):\n",
    "    assert isinstance(variable, Variable), \"Factor can only send messages to variable!\"\n",
    "    assert can_send_message(self, variable), \"Cannot send message!\"\n",
    "    \n",
    "    out_msg = calc_max_sum_factor_to_variable_msg(self, variable)\n",
    "    \n",
    "    # Send the message\n",
    "    variable.receive_msg(self, out_msg)\n",
    "    \n",
    "    # Remove the pending sign if present\n",
    "    self.pending.discard(variable)\n",
    "    \n",
    "Factor.send_ms_msg = factor_send_ms_msg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "36080ee2c75dc501b57f32f527c65656",
     "grade": true,
     "grade_id": "cell-c6156e1505633bc4",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "# message from X_prior to X\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "\n",
    "X_prior.send_ms_msg(X)\n",
    "assert np.allclose(list(X.in_msgs.values()), [-0.05129329, -2.99573227])\n",
    "\n",
    "# message from Z_prior to Z\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "\n",
    "Z_prior.send_ms_msg(Z)\n",
    "assert np.allclose(list(Z.in_msgs.values()), [-0.22314355, -1.60943791])\n",
    "\n",
    "# message from Y_cond to Y\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "Y_cond.receive_msg(X, X_prior.f) # simulating that Y_cond received all necessary messages from X\n",
    "Y_cond.receive_msg(Z, Z_prior.f) # simulating that Y_cond received all necessary messages from Z\n",
    "Y_cond.send_ms_msg(Y)\n",
    "assert np.allclose(list(Y.in_msgs.values()), [1.74989999, 0.79332506])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "cfbc002bdf7c6ff9b40c06ef10403001",
     "grade": false,
     "grade_id": "cell-6eff0dd8439243ca",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 2.2 Variable to factor messages (10 points)\n",
    "Implement the `Variable.send_ms_msg` function that sends Variable -> Factor messages for the max-sum algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "a8e8c05dffcf10717b1e43299b96b627",
     "grade": false,
     "grade_id": "cell-183948f06650a7e0",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "# ANSWER 2.2\n",
    "\n",
    "def calc_max_sum_product_variable_to_factor_msg(variable, factor):\n",
    "    # YOUR CODE HERE\n",
    "    neighbour_msg = get_neighbour_messages(variable, factor)\n",
    "    \n",
    "    if len(neighbour_msg) > 0:\n",
    "        msg = np.sum(np.array(neighbour_msg), axis=0)\n",
    "    else:\n",
    "        msg = np.zeros(variable.num_states)\n",
    "\n",
    "    return msg + np.log(variable.observed_state)\n",
    "\n",
    "def variable_send_ms_msg(self, factor):\n",
    "    \n",
    "    # YOUR CODE HERE\n",
    "    assert isinstance(factor, Factor), \"Variable can only send messages to factor!\"\n",
    "    assert can_send_message(self, factor), \"Cannot send message!\"\n",
    "    \n",
    "    out_msg = calc_max_sum_product_variable_to_factor_msg(self, factor)\n",
    "    \n",
    "    # Send the message\n",
    "    factor.receive_msg(self, out_msg)\n",
    "    \n",
    "    # Remove the pending sign if present\n",
    "    self.pending.discard(factor)\n",
    "\n",
    "Variable.send_ms_msg = variable_send_ms_msg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b597b2dc5896b79f1fb39862d19f2ea4",
     "grade": true,
     "grade_id": "cell-bbf08d58e9b309c9",
     "locked": true,
     "points": 10,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "# First message from X to Y_cond\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "X_prior.send_ms_msg(X) # simulating that X received all necessary messages\n",
    "X.send_ms_msg(Y_cond)\n",
    "assert np.allclose(list(Y_cond.in_msgs.values()), [-0.05129329, -2.99573227])\n",
    "\n",
    "# Second message from Z to Y_cond\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "Z_prior.send_ms_msg(Z) # simulating that Z received all necessary messages\n",
    "Z.send_ms_msg(Y_cond)\n",
    "assert np.allclose(list(Y_cond.in_msgs.values()), [[-0.05129329, -2.99573227], [-0.22314355, -1.60943791]])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "d94749234e759cb16863184e99919a57",
     "grade": false,
     "grade_id": "cell-6e1b3e4ee8013fcb",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 2.3 Implement unnormalized log marginal (5 points)\n",
    "Write the method `unnormalized_log_marginal` for the `Variable` class, that computes a unnormalized log marginal distribution over that node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "2a82b636f2eb1326ea6c7b4cce15f428",
     "grade": false,
     "grade_id": "cell-96de1de3d9b1da94",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def unnormalized_log_marginal(self):\n",
    "    # YOUR CODE HERE\n",
    "    np_values = np.array(list(self.in_msgs.values()))\n",
    "    unnormalized_log_marginal = np.sum(np_values, axis=0)\n",
    "    unnormalized_log_marginal = unnormalized_log_marginal + np.log(self.observed_state)\n",
    "    return unnormalized_log_marginal\n",
    "\n",
    "Variable.unnormalized_log_marginal = unnormalized_log_marginal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "977c4965c9797bd48b917039e495038a",
     "grade": true,
     "grade_id": "cell-d0c08fd4b0857dde",
     "locked": true,
     "points": 5,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test\n",
    "# Simulate a single forward pass\n",
    "X_prior.reset()\n",
    "X.reset()\n",
    "Z_prior.reset()\n",
    "Z.reset()\n",
    "Y_cond.reset()\n",
    "Y.reset()\n",
    "\n",
    "X_prior.send_ms_msg(X)\n",
    "Z_prior.send_ms_msg(Z)\n",
    "X.send_ms_msg(Y_cond)\n",
    "Z.send_ms_msg(Y_cond)\n",
    "Y_cond.send_ms_msg(Y)\n",
    "\n",
    "assert np.allclose(X.unnormalized_log_marginal(), [-0.05129329, -2.99573227])\n",
    "assert np.allclose(Z.unnormalized_log_marginal(), [-0.22314355, -1.60943791])\n",
    "assert np.allclose(Y.unnormalized_log_marginal(), [-0.27453685, -2.01740615])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "811de1e611d918a437d3436251ae8f56",
     "grade": false,
     "grade_id": "cell-eb540987e2fb70d6",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 2.4 Find a MAP state (10 points)\n",
    "\n",
    "Using the same message passing schedule we used for sum-product, implement the max-sum algorithm. For simplicity, we will ignore issues relating to non-unique maxima. So there is no need to implement backtracking; the MAP state is obtained by a per-node maximization (eq. 8.98 in Bishop). Make sure your algorithm works with both latent and observed variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "17fb320cc9db70730f6e7a6fd0702903",
     "grade": false,
     "grade_id": "cell-f94559550902ac16",
     "locked": false,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "def max_sum(node_list):\n",
    "    # forward\n",
    "    for node_a in node_list:\n",
    "        for node_b in list(node_a.pending):\n",
    "            node_a.send_ms_msg(node_b)\n",
    "    \n",
    "    for node_a in reversed(node_list):\n",
    "        for node_b in list(node_a.pending):\n",
    "            node_a.send_ms_msg(node_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "49674d2f737491907099ead2f661ab25",
     "grade": true,
     "grade_id": "cell-ac2d26dcb27a8d75",
     "locked": true,
     "points": 5,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test: unobserved\n",
    "nodes = [X_prior, X, Z_prior, Z, Y_cond, Y]\n",
    "for n in nodes:\n",
    "    n.reset()\n",
    "    \n",
    "X_prior.pending.add(X)\n",
    "Z_prior.pending.add(Z)\n",
    "Y.pending.add(Y_cond)\n",
    "\n",
    "max_sum(nodes)\n",
    "\n",
    "assert np.allclose(Y.unnormalized_log_marginal(), [-0.27453685, -2.01740615] )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d00995d3a4fb74ce8ed9481349e9044d",
     "grade": true,
     "grade_id": "cell-f06b8b4bec85b4dd",
     "locked": true,
     "points": 5,
     "schema_version": 3,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "### Test test test: partiallY observed\n",
    "nodes = [X_prior, X, Z_prior, Z, Y_cond, Y]\n",
    "for n in nodes:\n",
    "    n.reset()\n",
    "    \n",
    "X_prior.pending.add(X)\n",
    "Z_prior.pending.add(Z)\n",
    "Y.pending.add(Y_cond)\n",
    "\n",
    "Z.set_observed(1)\n",
    "\n",
    "max_sum(nodes)\n",
    "\n",
    "assert np.allclose(Y.unnormalized_log_marginal(), [-2.86470401, -2.01740615])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "6e5b7cb4dc60f898345d08a07889f32c",
     "grade": false,
     "grade_id": "cell-f79ce013522649ed",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "Given the max-marginals what do you have to do in order to find the global optimum? Why can we neglect the normalization constant here?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "1bd8a1937fd5e81b184e36bc4c4b1b68",
     "grade": true,
     "grade_id": "cell-605e729ee910dbc8",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "source": [
    "Given the max-marginals we use them by keeping track from exactly which values they came from. This way we know which optima are local and which global.\n",
    "\n",
    "Furthermore, we can neglect the normalization constant because we are aiming to find the global optimum. This is independent from the normalization constant exactly because it is constant and it won't change, therefore the task is the same as trying to find the global optimum to the other part of the equation as well."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "80bbe06f90fa8e5f278119c02cf5081d",
     "grade": false,
     "grade_id": "cell-d0592bafd25122c0",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### Part 3: Medical graph\n",
    "Now that we implemented the sum-product and max-sum algorithm. We will apply them to a poly-tree structured medical diagnosis example.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "9c99fffc963ea4db625082689e3b7c55",
     "grade": false,
     "grade_id": "cell-f0022d301754179a",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 3.1 Initialize the graph (5 points)\n",
    "\n",
    "Convert the directed graphical model (\"Bayesian Network\") shown below to a factor graph. Instantiate this graph by creating Variable and Factor instances and linking them according to the graph structure. \n",
    "To instantiate the factor graph, first create the Variable nodes and then create Factor nodes, passing a list of neighbour Variables to each Factor.\n",
    "Use the following prior and conditional probabilities.\n",
    "\n",
    "$$\n",
    "p(\\verb+Influenza+) = 0.05 \\\\\\\\\n",
    "p(\\verb+Smokes+) = 0.2 \\\\\\\\\n",
    "$$\n",
    "$$\n",
    "p(\\verb+SoreThroat+ = 1 | \\verb+Influenza+ = 1) = 0.3 \\\\\\\\\n",
    "p(\\verb+SoreThroat+ = 1 | \\verb+Influenza+ = 0) = 0.001 \\\\\\\\\n",
    "p(\\verb+Fever+ = 1| \\verb+Influenza+ = 1) = 0.9 \\\\\\\\\n",
    "p(\\verb+Fever+ = 1| \\verb+Influenza+ = 0) 0.05 \\\\\\\\\n",
    "p(\\verb+Bronchitis+ = 1 | \\verb+Influenza+ = 1, \\verb+Smokes+ = 1) = 0.99 \\\\\\\\\n",
    "p(\\verb+Bronchitis+ = 1 | \\verb+Influenza+ = 1, \\verb+Smokes+ = 0) = 0.9 \\\\\\\\\n",
    "p(\\verb+Bronchitis+ = 1 | \\verb+Influenza+ = 0, \\verb+Smokes+ = 1) = 0.7 \\\\\\\\\n",
    "p(\\verb+Bronchitis+ = 1 | \\verb+Influenza+ = 0, \\verb+Smokes+ = 0) = 0.0001 \\\\\\\\\n",
    "p(\\verb+Coughing+ = 1| \\verb+Bronchitis+ = 1) = 0.8 \\\\\\\\\n",
    "p(\\verb+Coughing+ = 1| \\verb+Bronchitis+ = 0) = 0.07 \\\\\\\\\n",
    "p(\\verb+Wheezing+ = 1| \\verb+Bronchitis+ = 1) = 0.6 \\\\\\\\\n",
    "p(\\verb+Wheezing+ = 1| \\verb+Bronchitis+ = 0) = 0.001 \\\\\\\\\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "4335e3cbfd46c7579a4b7d368c5353cc",
     "grade": true,
     "grade_id": "cell-216d36a73771fdb7",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "I = Variable(name='Influenza', num_states=2)\n",
    "S = Variable(name='Smokes', num_states=2)\n",
    "ST = Variable(name='SoreThroat', num_states=2)\n",
    "F = Variable(name='Fever', num_states=2)\n",
    "B = Variable(name='Bronchitis', num_states=2)\n",
    "C = Variable(name='Coughing', num_states=2)\n",
    "W = Variable(name='Wheezing', num_states=2)\n",
    "\n",
    "I_prior = Factor(name='p(Influenza)',\n",
    "                 f=np.array([0.95, 0.05]),\n",
    "                 neighbours=[I])\n",
    "S_prior = Factor(name='p(Smokes)',\n",
    "                 f=np.array([0.8, 0.2]),\n",
    "                 neighbours=[S])\n",
    "\n",
    "ST_I_cond = Factor(name='p(ST|I)',\n",
    "                f=np.array([[0.999, 0.7], [0.001, 0.3]]),\n",
    "                neighbours=[ST, I])\n",
    "\n",
    "F_I_cond = Factor(name='p(F|I)',\n",
    "                f=np.array([[0.95, 0.1], [0.05, 0.9]]),\n",
    "                neighbours=[F, I])\n",
    "\n",
    "B_I_S_cond = Factor(name='p(B|I,S)',\n",
    "                f=np.array([[[0.9999, 0.3], [0.1, 0.01]],\n",
    "                            [[0.0001, 0.7], [0.9, 0.99]]]),\n",
    "                neighbours=[B, I, S])\n",
    "\n",
    "C_B_cond = Factor(name='p(C|B)',\n",
    "                f=np.array([[0.93, 0.2], [0.07, 0.8]]),\n",
    "                neighbours=[C, B])\n",
    "\n",
    "W_B_cond = Factor(name='p(W|B)',\n",
    "                f=np.array([[0.999, 0.4], [0.001, 0.6]]),\n",
    "                neighbours=[W, B])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "8e443821bd73340b3e8bc3eb58db013c",
     "grade": false,
     "grade_id": "cell-65ee0e0d85c57f32",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 3.2. Sum-product algorithm (10 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "48d2bd6a0a7dba5d29de2c939ca5c67f",
     "grade": false,
     "grade_id": "cell-fed678d3f9f5ce50",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "3.2.1 Run the sum-product algorithm on an unobserved graph. Print the marginal for every variable node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "a7fadef738bc786fd10febd2b25f15f5",
     "grade": true,
     "grade_id": "cell-07681153d48af573",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SoreThroat: [0.98405 0.01595]\n",
      "Fever: [0.9075 0.0925]\n",
      "Coughing: [0.79934752 0.20065248]\n",
      "Wheezing: [0.89179338 0.10820662]\n",
      "Influenza: [0.95 0.05]\n",
      "Bronchitis: [0.821024 0.178976]\n",
      "Smokes: [0.8 0.2]\n"
     ]
    }
   ],
   "source": [
    "def run_algorithm(sum_function, middleware=None):\n",
    "    medicine_nodes = [ST, F, C, W, I_prior, ST_I_cond, F_I_cond, C_B_cond, \n",
    "                     W_B_cond, I, B, B_I_S_cond, S, S_prior]\n",
    "\n",
    "\n",
    "    for node in medicine_nodes:\n",
    "        node.reset()\n",
    "\n",
    "    C.pending.add(C_B_cond)\n",
    "    W.pending.add(W_B_cond)\n",
    "    I_prior.pending.add(I)\n",
    "    S_prior.pending.add(S)\n",
    "    ST.pending.add(ST_I_cond)\n",
    "    F.pending.add(F_I_cond)\n",
    "\n",
    "    if middleware:\n",
    "        middleware()\n",
    "    \n",
    "    sum_function(medicine_nodes)\n",
    "\n",
    "    if sum_function.__name__ == 'sum_product':\n",
    "        for node in medicine_nodes:\n",
    "            if isinstance(node, Variable):\n",
    "                print(f'{node.name}: {node.marginal()}')\n",
    "    elif sum_function.__name__ == 'max_sum':\n",
    "        for node in medicine_nodes:\n",
    "            if isinstance(node, Variable):\n",
    "                print(f'{node.name}: {node.unnormalized_log_marginal()}')\n",
    "        \n",
    "run_algorithm(sum_product)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "7ab62562db323a454298e365f0989326",
     "grade": false,
     "grade_id": "cell-9cf9ad87ef338231",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "3.2.2 Rerun the sum-product algorithm on an partially observed graph, where the variable 'Influenza' is set to 0. Print the marginal for every variable node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "0bf552eaf0812c495b3cb3c79f2fc431",
     "grade": true,
     "grade_id": "cell-83214d3cc28cb7db",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SoreThroat: [0.999 0.001]\n",
      "Fever: [0.95 0.05]\n",
      "Coughing: [0.8277416 0.1722584]\n",
      "Wheezing: [0.91509208 0.08490792]\n",
      "Influenza: [9.5e-01 5.0e-12]\n",
      "Bronchitis: [0.85992 0.14008]\n",
      "Smokes: [0.8 0.2]\n"
     ]
    }
   ],
   "source": [
    "run_algorithm(sum_product, lambda: I.set_observed(0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "d5c513d1f3fcdd0e153834c8cd7013e7",
     "grade": false,
     "grade_id": "cell-c2dbbc1def846b5e",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "### 3.3. max-sum algorithm (10 points)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "e3397f5416e8ef7a287fd0e1844d5ec6",
     "grade": false,
     "grade_id": "cell-2aeb9d6f8bb0d6d9",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "3.3.1 Run the max_sum algorithm on an unobserved graph. Print the marginal for every variable node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "350a548fee4372022d9d7b51805ae71a",
     "grade": true,
     "grade_id": "cell-e5088e48afe006b9",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SoreThroat: [-0.40040184 -5.36753884]\n",
      "Fever: [-0.40040184 -3.34484082]\n",
      "Coughing: [-0.40040184 -2.80366912]\n",
      "Wheezing: [-0.40040184 -2.80366912]\n",
      "Influenza: [-0.40040184 -4.52024098]\n",
      "Bronchitis: [-0.40040184 -2.80366912]\n",
      "Smokes: [-0.40040184 -2.80366912]\n"
     ]
    }
   ],
   "source": [
    "run_algorithm(max_sum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "083775bc40177f2573543fd168e86e09",
     "grade": false,
     "grade_id": "cell-7e7771be2511ebd2",
     "locked": true,
     "schema_version": 3,
     "solution": false
    }
   },
   "source": [
    "3.3.2 Rerun the max_sum algorithm on an partially observed graph, where the variable 'Influenza' is set to 0. Print the marginal for every variable node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b467a95a6dbe003435c78b970ed0e318",
     "grade": true,
     "grade_id": "cell-f74e50a37f4e31b9",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SoreThroat: [-0.40040184 -7.30715662]\n",
      "Fever: [-0.40040184 -3.34484082]\n",
      "Coughing: [-0.40040184 -2.80366912]\n",
      "Wheezing: [-0.40040184 -2.80366912]\n",
      "Influenza: [ -0.40040184 -27.54609191]\n",
      "Bronchitis: [-0.40040184 -2.80366912]\n",
      "Smokes: [-0.40040184 -2.80366912]\n"
     ]
    }
   ],
   "source": [
    "run_algorithm(max_sum, lambda: I.set_observed(0))"
   ]
  }
 ],
 "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}