{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Programming Assignment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Naive Bayes and logistic regression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instructions\n",
    "\n",
    "In this notebook, you will write code to develop a Naive Bayes classifier model to the Iris dataset using Distribution objects from TensorFlow Probability. You will also explore the connection between the Naive Bayes classifier and logistic regression.\n",
    "\n",
    "Some code cells are provided you in the notebook. You should avoid editing provided code, and make sure to execute the cells in order to avoid unexpected errors. Some cells begin with the line: \n",
    "\n",
    "`#### GRADED CELL ####`\n",
    "\n",
    "Don't move or edit this first line - this is what the automatic grader looks for to recognise graded cells. These cells require you to write your own code to complete them, and are automatically graded when you submit the notebook. Don't edit the function name or signature provided in these cells, otherwise the automatic grader might not function properly.\n",
    "\n",
    "### How to submit\n",
    "\n",
    "Complete all the tasks you are asked for in the worksheet. When you have finished and are happy with your code, press the **Submit Assignment** button at the top of this notebook.\n",
    "\n",
    "### Let's get started!\n",
    "\n",
    "We'll start running some imports, and loading the dataset. Do not edit the existing imports in the following cell. If you would like to make further Tensorflow imports, you should add them here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### PACKAGE IMPORTS ####\n",
    "\n",
    "# Run this cell first to import all required packages. Do not make any imports elsewhere in the notebook\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "tfd = tfp.distributions\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn import datasets, model_selection\n",
    "%matplotlib inline\n",
    "\n",
    "# If you would like to make further imports from TensorFlow or TensorFlow Probability, add them here\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<tr>\n",
    "<td><img src=\"data/iris_setosa.jpg\" alt=\"Drawing\" style=\"height: 270px;\"/></td>\n",
    "<td><img src=\"data/iris_versicolor.jpg\" alt=\"Drawing\" style=\"height: 270px;\"/></td>\n",
    "<td><img src=\"data/iris_virginica.jpg\" alt=\"Drawing\" style=\"height: 270px;\"/></td>\n",
    "</tr>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The Iris dataset\n",
    "\n",
    "In this assignment, you will use the [Iris dataset](https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html). It consists of 50 samples from each of three species of Iris (Iris setosa, Iris virginica and Iris versicolor). Four features were measured from each sample: the length and the width of the sepals and petals, in centimeters. For a reference, see the following papers:\n",
    "\n",
    "- R. A. Fisher. \"The use of multiple measurements in taxonomic problems\". Annals of Eugenics. 7 (2): 179–188, 1936.\n",
    "\n",
    "Your goal is to construct a Naive Bayes classifier model that predicts the correct class from the sepal length and sepal width features. Under certain assumptions about this classifier model, you will explore the relation to logistic regression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load and prepare the data\n",
    "\n",
    "We will first read in the Iris dataset, and split the dataset into training and test sets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the dataset\n",
    "\n",
    "iris = datasets.load_iris()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use only the first two features: sepal length and width\n",
    "\n",
    "data = iris.data[:, :2]\n",
    "targets = iris.target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Randomly shuffle the data and make train and test splits\n",
    "\n",
    "x_train, x_test, y_train, y_test = model_selection.train_test_split(data, targets, test_size=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19428 characters]\n",
      "text/plain": [
       "<Figure size 576x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the training data\n",
    "\n",
    "labels = {0: 'Iris-Setosa', 1: 'Iris-Versicolour', 2: 'Iris-Virginica'}\n",
    "label_colours = ['blue', 'orange', 'green']\n",
    "\n",
    "def plot_data(x, y, labels, colours):\n",
    "    for c in np.unique(y):\n",
    "        inx = np.where(y == c)\n",
    "        plt.scatter(x[inx, 0], x[inx, 1], label=labels[c], c=colours[c])\n",
    "    plt.title(\"Training set\")\n",
    "    plt.xlabel(\"Sepal length (cm)\")\n",
    "    plt.ylabel(\"Sepal width (cm)\")\n",
    "    plt.legend()\n",
    "    \n",
    "plt.figure(figsize=(8, 5))\n",
    "plot_data(x_train, y_train, labels, label_colours)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Naive Bayes classifier\n",
    "\n",
    "We will briefly review the Naive Bayes classifier model. The fundamental equation for this classifier is Bayes' rule:\n",
    "\n",
    "$$\n",
    "P(Y=y_k | X_1,\\ldots,X_d) = \\frac{P(X_1,\\ldots,X_d | Y=y_k)P(Y=y_k)}{\\sum_{k=1}^K P(X_1,\\ldots,X_d | Y=y_k)P(Y=y_k)}\n",
    "$$\n",
    "\n",
    "In the above, $d$ is the number of features or dimensions in the inputs $X$ (in our case $d=2$), and $K$ is the number of classes (in our case $K=3$). The distribution $P(Y)$ is the class prior distribution, which is a discrete distribution over $K$ classes. The distribution $P(X | Y)$ is the class-conditional distribution over inputs.\n",
    "\n",
    "The Naive Bayes classifier makes the assumption that the data features $X_i$ are conditionally independent give the class $Y$ (the 'naive' assumption). In this case, the class-conditional distribution decomposes as\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "P(X | Y=y_k) &= P(X_1,\\ldots,X_d | Y=y_k)\\\\\n",
    "&= \\prod_{i=1}^d P(X_i | Y=y_k)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "This simplifying assumption means that we typically need to estimate far fewer parameters for each of the distributions $P(X_i | Y=y_k)$ instead of the full joint distribution $P(X | Y=y_k)$.\n",
    "\n",
    "Once the class prior distribution and class-conditional densities are estimated, the Naive Bayes classifier model can then make a class prediction $\\hat{Y}$ for a new data input $\\tilde{X} := (\\tilde{X}_1,\\ldots,\\tilde{X}_d)$ according to\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\hat{Y} &= \\text{argmax}_{y_k} P(Y=y_k | \\tilde{X}_1,\\ldots,\\tilde{X}_d) \\\\\n",
    "&= \\text{argmax}_{y_k}\\frac{P(\\tilde{X}_1,\\ldots,\\tilde{X}_d | Y=y_k)P(Y=y_k)}{\\sum_{k=1}^K P(\\tilde{X}_1,\\ldots,\\tilde{X}_d | Y=y_k)P(Y=y_k)}\\\\\n",
    "&= \\text{argmax}_{y_k} P(\\tilde{X}_1,\\ldots,\\tilde{X}_d | Y=y_k)P(Y=y_k)\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define the class prior distribution\n",
    " \n",
    "We will begin by defining the class prior distribution. To do this we will simply take the maximum likelihood estimate, given by\n",
    "\n",
    "$$\n",
    "P(Y=y_k) = \\frac{\\sum_{n=1}^N \\delta(Y^{(n)}=y_k)}{N},\n",
    "$$\n",
    "\n",
    "where the superscript $(n)$ indicates the $n$-th dataset example, $\\delta(Y^{(n)}=y_k) = 1$ if $Y^{(n)}=y_k$ and 0 otherwise, and $N$ is the total number of examples in the dataset. The above is simply the proportion of data examples belonging to class $k$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now write a function that builds the prior distribution from the training data, and returns it as a `Categorical` Distribution object.\n",
    "\n",
    "* The input to your function `y` will be a numpy array of shape `(num_samples,)`\n",
    "* The entries in `y` will be integer labels $k=0, 1,\\ldots, K-1$\n",
    "* Your function should build and return the prior distribution as a `Categorical` distribution object\n",
    "  * The probabilities for this distribution will be a length-$K$ vector, with entries corresponding to $P(Y = y_k)$ for $k=0,1,\\ldots,K-1$\n",
    "  * Your function should work for any value of $K\\ge 1$\n",
    "  * This Distribution will have an empty batch shape and empty event shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def get_prior(y):\n",
    "    \"\"\"\n",
    "    This function takes training labels as a numpy array y of shape (num_samples,) as an input.\n",
    "    Your function should \n",
    "    This function should build a Categorical Distribution object with empty batch shape \n",
    "    and event shape, with the probability of each class given as above. \n",
    "    Your function should return the Distribution object.\n",
    "    \"\"\"\n",
    "    _, counts = np.unique(y, return_counts=True)\n",
    "    priors = counts / len(y)\n",
    "    dist = tfd.Categorical(probs=priors)\n",
    "    return dist    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run your function to get the prior\n",
    "\n",
    "prior = get_prior(y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10700 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the prior distribution\n",
    "\n",
    "labels = ['Iris-Setosa', 'Iris-Versicolour', 'Iris-Virginica']\n",
    "plt.bar([0, 1, 2], prior.probs.numpy(), color=label_colours)\n",
    "plt.xlabel(\"Class\")\n",
    "plt.ylabel(\"Prior probability\")\n",
    "plt.title(\"Class prior distribution\")\n",
    "plt.xticks([0, 1, 2], labels)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define the class-conditional densities\n",
    "\n",
    "We now turn to the definition of the class-conditional distributions $P(X_i | Y=y_k)$ for $i=0, 1$ and $k=0, 1, 2$. In our model, we will assume these distributions to be univariate Gaussian:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "P(X_i | Y=y_k) &= N(X_i | \\mu_{ik}, \\sigma_{ik})\\\\\n",
    "&= \\frac{1}{\\sqrt{2\\pi\\sigma_{ik}^2}} \\exp\\left\\{-\\frac{1}{2} \\left(\\frac{x - \\mu_{ik}}{\\sigma_{ik}}\\right)^2\\right\\}\n",
    "\\end{align}\n",
    "$$\n",
    "with mean parameters $\\mu_{ik}$ and standard deviation parameters $\\sigma_{ik}$, twelve parameters in all. We will again estimate these parameters using maximum likelihood. In this case, the estimates are given by\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\hat{\\mu}_{ik} &= \\frac{\\sum_n X_i^{(n)} \\delta(Y^{(n)}=y_k)}{\\sum_n \\delta(Y^{(n)}=y_k)} \\\\\n",
    "\\hat{\\sigma}^2_{ik} &= \\frac{\\sum_n (X_i^{(n)} - \\hat{\\mu}_{ik})^2 \\delta(Y^{(n)}=y_k)}{\\sum_n \\delta(Y^{(n)}=y_k)}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "Note that the above are just the means and variances of the sample data points for each class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now write a function the computes the class-conditional Gaussian densities, using the maximum likelihood parameter estimates given above, and returns them in a single, batched `MultivariateNormalDiag` Distribution object. \n",
    "\n",
    "* The inputs to the function are \n",
    "  * a numpy array `x` of shape `(num_samples, num_features)` for the data inputs\n",
    "  * a numpy array `y` of shape `(num_samples,)` for the target labels\n",
    "* Your function should work for any number of classes $K\\ge 1$ and any number of features $d\\ge 1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def get_class_conditionals(x, y):\n",
    "    \"\"\"\n",
    "    This function takes training data samples x and labels y as inputs.\n",
    "    This function should build the class-conditional Gaussian distributions above. \n",
    "    It should construct a batch of distributions for each feature and each class, using the \n",
    "    parameter estimates above for the means and standard deviations.\n",
    "    The batch shape of this distribution should be rank 2, where the first dimension corresponds\n",
    "    to the number of classes and the second corresponds to the number of features.\n",
    "    Your function should then return the Distribution object.\n",
    "    \"\"\"\n",
    "    n_features = x.shape[1]\n",
    "    classes = np.unique(y)\n",
    "    mu = np.zeros((len(classes),n_features))\n",
    "    sig = np.zeros((len(classes),n_features))\n",
    "    for k in classes:\n",
    "        mu[k] = np.mean(x[y==k,:],axis=0)\n",
    "        sig[k] = np.std(x[y==k,:],axis=0)\n",
    "    dist = tfd.MultivariateNormalDiag(loc=mu, scale_diag=sig)\n",
    "    return dist\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag' batch_shape=[3] event_shape=[2] dtype=float64>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run your function to get the class-conditional distributions\n",
    "\n",
    "class_conditionals = get_class_conditionals(x_train, y_train)\n",
    "class_conditionals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualise the class-conditional densities with contour plots by running the cell below. Notice how the contours of each distribution correspond to a Gaussian distribution with diagonal covariance matrix, since the model assumes that each feature is independent given the class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 93104 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the training data with the class-conditional density contours\n",
    "\n",
    "def get_meshgrid(x0_range, x1_range, num_points=100):\n",
    "    x0 = np.linspace(x0_range[0], x0_range[1], num_points)\n",
    "    x1 = np.linspace(x1_range[0], x1_range[1], num_points)\n",
    "    return np.meshgrid(x0, x1)\n",
    "\n",
    "def contour_plot(x0_range, x1_range, prob_fn, batch_shape, colours, levels=None, num_points=100):\n",
    "    X0, X1 = get_meshgrid(x0_range, x1_range, num_points=num_points)\n",
    "    Z = prob_fn(np.expand_dims(np.array([X0.ravel(), X1.ravel()]).T, 1))\n",
    "    Z = np.array(Z).T.reshape(batch_shape, *X0.shape)\n",
    "    for batch in np.arange(batch_shape):\n",
    "        if levels:\n",
    "            plt.contourf(X0, X1, Z[batch], alpha=0.2, colors=colours, levels=levels)\n",
    "        else:\n",
    "            plt.contour(X0, X1, Z[batch], colors=colours[batch], alpha=0.3)\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plot_data(x_train, y_train, labels, label_colours)\n",
    "x0_min, x0_max = x_train[:, 0].min(), x_train[:, 0].max()\n",
    "x1_min, x1_max = x_train[:, 1].min(), x_train[:, 1].max()\n",
    "contour_plot((x0_min, x0_max), (x1_min, x1_max), class_conditionals.prob, 3, label_colours)\n",
    "plt.title(\"Training set with class-conditional density contours\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Make predictions from the model\n",
    "\n",
    "Now the prior and class-conditional distributions are defined, you can use them to compute the model's class probability predictions for an unknown test input $\\tilde{X} = (\\tilde{X}_1,\\ldots,\\tilde{X}_d)$, according to\n",
    "\n",
    "$$\n",
    "P(Y=y_k | \\tilde{X}_1,\\ldots,\\tilde{X}_d) = \\frac{P(\\tilde{X}_1,\\ldots,\\tilde{X}_d | Y=y_k)P(Y=y_k)}{\\sum_{k=1}^K P(\\tilde{X}_1,\\ldots,\\tilde{X}_d | Y=y_k)P(Y=y_k)}\n",
    "$$\n",
    "\n",
    "The class prediction can then be taken as the class with the maximum probability:\n",
    "\n",
    "$$\n",
    "\\hat{Y} = \\text{argmax}_{y_k} P(Y=y_k | \\tilde{X}_1,\\ldots,\\tilde{X}_d)\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now write a function to return the model's class probabilities for a given batch of test inputs of shape `(batch_shape, 2)`, where the `batch_shape` has rank at least one. \n",
    "\n",
    "* The inputs to the function are the `prior` and `class_conditionals` distributions, and the inputs `x`\n",
    "* Your function should use these distributions to compute the probabilities for each class $k$ as above\n",
    "  * As before, your function should work for any number of classes $K\\ge 1$\n",
    "* It should then compute the prediction by taking the class with the highest probability\n",
    "* The predictions should be returned in a numpy array of shape `(batch_shape)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def predict_class(prior, class_conditionals, x):\n",
    "    \"\"\"\n",
    "    This function takes the prior distribution, class-conditional distribution, and \n",
    "    a batch of inputs in a numpy array of shape (batch_shape, 2).\n",
    "    This function should compute the class probabilities for each input in the batch, using\n",
    "    the prior and class-conditional distributions, according to the above equation.\n",
    "    Note that the batch_shape of x could have rank higher than one!\n",
    "    Your function should then return the class predictions by taking the class with the \n",
    "    maximum probability in a numpy array of shape (batch_shape,).\n",
    "    \"\"\"\n",
    "#     cond_probs = class_conditionals.log_prob(x)[ np.newaxis, :]\n",
    "#     joint_likelihood = tf.add(np.log(prior.probs), cond_probs)\n",
    "    prior_prob = np.log(prior.probs)[ np.newaxis, :]\n",
    "    cond_prob = class_conditionals.log_prob(x[:,np.newaxis])\n",
    "    joint_likelihood = tf.add(cond_prob, prior_prob)\n",
    "    norm_factor = tf.math.reduce_logsumexp(joint_likelihood, axis=-1, keepdims=True)\n",
    "    log_prob = tf.exp(joint_likelihood - norm_factor)\n",
    "    yy = tf.argmax(log_prob,axis=0-1)\n",
    "    return yy.numpy()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Get the class predictions\n",
    "\n",
    "predictions = predict_class(prior, class_conditionals, x_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test accuracy: 0.7333\n"
     ]
    }
   ],
   "source": [
    "# Evaluate the model accuracy on the test set\n",
    "\n",
    "accuracy = accuracy_score(y_test, predictions)\n",
    "print(\"Test accuracy: {:.4f}\".format(accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26884 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the model's decision regions\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plot_data(x_train, y_train, labels, label_colours)\n",
    "x0_min, x0_max = x_train[:, 0].min(), x_train[:, 0].max()\n",
    "x1_min, x1_max = x_train[:, 1].min(), x_train[:, 1].max()\n",
    "contour_plot((x0_min, x0_max), (x1_min, x1_max),\n",
    "             lambda x: predict_class(prior, class_conditionals, x), \n",
    "             1, label_colours, levels=[-0.5, 0.5, 1.5, 2.5],\n",
    "             num_points=300)\n",
    "plt.title(\"Training set with decision regions\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Binary classifier\n",
    "\n",
    "We will now draw a connection between the Naive Bayes classifier and logistic regression.\n",
    "\n",
    "First, we will update our model to be a binary classifier. In particular, the model will output the probability that a given input data sample belongs to the 'Iris-Setosa' class: $P(Y=y_0 | \\tilde{X}_1,\\ldots,\\tilde{X}_d)$. The remaining two classes will be pooled together with the label $y_1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redefine the dataset to have binary labels\n",
    "\n",
    "y_train_binary = np.array(y_train)\n",
    "y_train_binary[np.where(y_train_binary == 2)] = 1\n",
    "\n",
    "y_test_binary = np.array(y_test)\n",
    "y_test_binary[np.where(y_test_binary == 2)] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17852 characters]\n",
      "text/plain": [
       "<Figure size 576x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the training data\n",
    "\n",
    "labels_binary = {0: 'Iris-Setosa', 1: 'Iris-Versicolour / Iris-Virginica'}\n",
    "label_colours_binary = ['blue', 'red']\n",
    "\n",
    "plt.figure(figsize=(8, 5))\n",
    "plot_data(x_train, y_train_binary, labels_binary, label_colours_binary)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will also make an extra modelling assumption that for each class $k$, the class-conditional distribution $P(X_i | Y=y_k)$ for each feature $i=0, 1$, has standard deviation $\\sigma_i$, which is the same for each class $k$. \n",
    "\n",
    "This means there are now six parameters in total: four for the means $\\mu_{ik}$ and two for the standard deviations $\\sigma_i$ ($i, k=0, 1$). \n",
    "\n",
    "We will again use maximum likelihood to estimate these parameters. The prior distribution will be as before, with the class prior probabilities given by\n",
    "\n",
    "$$\n",
    "P(Y=y_k) = \\frac{\\sum_{n=1}^N \\delta(Y^{(n)}=y_k)}{N},\n",
    "$$\n",
    "\n",
    "We will use your previous function `get_prior` to redefine the prior distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redefine the prior\n",
    "\n",
    "prior_binary = get_prior(y_train_binary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8208 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the prior distribution\n",
    "\n",
    "plt.bar([0, 1], prior_binary.probs.numpy(), color=label_colours_binary)\n",
    "plt.xlabel(\"Class\")\n",
    "plt.ylabel(\"Prior probability\")\n",
    "plt.title(\"Class prior distribution\")\n",
    "plt.xticks([0, 1], labels_binary)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the class-conditional densities, the maximum likelihood estimate for the means are again given by\n",
    "\n",
    "$$\n",
    "\\hat{\\mu}_{ik} = \\frac{\\sum_n X_i^{(n)} \\delta(Y^{(n)}=y_k)}{\\sum_n \\delta(Y^{(n)}=y_k)} \\\\\n",
    "$$\n",
    "\n",
    "However, the estimate for the standard deviations $\\sigma_i$ is updated. There is also a closed-form solution for the shared standard deviations, but we will instead learn these from the data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now write a function that takes the training inputs and target labels as input, as well as an optimizer object, number of epochs and a TensorFlow Variable. This function should be written according to the following spec:\n",
    "\n",
    "* The inputs to the function are:\n",
    "  * a numpy array `x` of shape `(num_samples, num_features)` for the data inputs\n",
    "  * a numpy array `y` of shape `(num_samples,)` for the target labels\n",
    "  * a `tf.Variable` object `scales` of length 2 for the standard deviations $\\sigma_i$\n",
    "  * `optimiser`: an optimiser object\n",
    "  * `epochs`: the number of epochs to run the training for\n",
    "* The function should first compute the means $\\mu_{ik}$ of the class-conditional Gaussians according to the above equation\n",
    "* Then create a batched multivariate Gaussian distribution object using `MultivariateNormalDiag` with the means set to $\\mu_{ik}$ and the scales set to `scales`\n",
    "* Run a custom training loop for `epochs` number of epochs, in which:\n",
    "  * the average per-example negative log likelihood for the whole dataset is computed as the loss\n",
    "  * the gradient of the loss with respect to the `scales` variables is computed\n",
    "  * the `scales` variables are updated by the `optimiser` object\n",
    "* At each iteration, save the values of the `scales` variable and the loss\n",
    "* The function should return a tuple of three objects:\n",
    "   * a numpy array of shape `(epochs,)` of loss values\n",
    "   * a numpy array of shape `(epochs, 2)` of values for the `scales` variable at each iteration\n",
    "   * the final learned batched `MultivariateNormalDiag` distribution object\n",
    "   \n",
    "_NB: ideally, we would like to constrain the `scales` variable to have positive values. We are not doing that here, but in later weeks of the course you will learn how this can be implemented._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 312,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "    \n",
    "def learn_stdevs(x, y, scales, optimiser, epochs):\n",
    "    \"\"\"\n",
    "    This function takes the data inputs, targets, scales variable, optimiser and number of\n",
    "    epochs as inputs.\n",
    "    This function should set up and run a custom training loop according to the above \n",
    "    specifications, by setting up the class conditional distributions as a MultivariateNormalDiag\n",
    "    object, and updating the trainable variables (the scales) in a custom training loop.\n",
    "    Your function should then return the a tuple of three elements: a numpy array of loss values\n",
    "    during training, a numpy array of scales variables during training, and the final learned\n",
    "    MultivariateNormalDiag distribution object.\n",
    "    \"\"\"\n",
    "    \n",
    "    # Keep results for plotting\n",
    "    train_loss_results = []\n",
    "    train_scales_results = []\n",
    "\n",
    "    \n",
    "    mu = np.zeros((x.shape[1], 2), dtype='float32')\n",
    "    for k in range(2):\n",
    "        mu[k] = np.mean(x[y==k,:],axis=0)\n",
    "\n",
    "    dist = tfd.MultivariateNormalDiag(loc=mu, scale_diag=scales)\n",
    "#     dist = get_class_conditionals(x, y)\n",
    "\n",
    "    for i in range(epochs):\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            tape.watch(dist.trainable_variables)\n",
    "            x_new = np.expand_dims(x_train, 1)\n",
    "            loss = -tf.reduce_mean(dist.log_prob(x_new))\n",
    "            grads = tape.gradient(loss, dist.trainable_variables)\n",
    "            optimiser.apply_gradients(zip(grads, dist.trainable_variables))\n",
    "\n",
    "            train_loss_results.append(loss)\n",
    "            train_scales_results.append(dist.parameters['scale_diag'].numpy())\n",
    "\n",
    "    return np.array(train_loss_results), np.array(train_scales_results), dist\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 314,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the inputs to your function\n",
    "\n",
    "scales = tf.Variable([1., 1.])\n",
    "opt = tf.keras.optimizers.Adam(learning_rate=0.01)\n",
    "epochs =500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 315,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2.5346925, 2.526281 , 2.5180852, 2.5100954, 2.502296 , 2.4946644,\n",
       "       2.487172 , 2.4797843, 2.4724627, 2.465171 , 2.457878 , 2.4505625,\n",
       "       2.443214 , 2.4358306, 2.4284205, 2.4209955, 2.413571 , 2.4061623,\n",
       "       2.3987863, 2.391457 , 2.3841875, 2.3769877, 2.3698654, 2.362826 ,\n",
       "       2.3558733, 2.3490114, 2.3422444, 2.335579 , 2.3290231, 2.3225896,\n",
       "       2.3162928, 2.3101487, 2.3041759, 2.298393 , 2.2928193, 2.287473 ,\n",
       "       2.282371 , 2.2775302, 2.2729657, 2.268692 , 2.264723 , 2.2610698,\n",
       "       2.2577446, 2.254756 , 2.2521107, 2.2498121, 2.2478597, 2.2462485,\n",
       "       2.2449672, 2.2439985, 2.2433202, 2.2429025, 2.2427125, 2.24271  ,\n",
       "       2.2428558, 2.243107 , 2.2434232, 2.2437654, 2.244101 , 2.2444022,\n",
       "       2.244647 , 2.2448223, 2.2449214, 2.244944 , 2.2448947, 2.2447836,\n",
       "       2.2446218, 2.2444232, 2.2442021, 2.243972 , 2.2437437, 2.2435272,\n",
       "       2.243331 , 2.2431595, 2.2430155, 2.2429001, 2.2428126, 2.2427511,\n",
       "       2.2427123, 2.2426925, 2.242688 , 2.2426956, 2.2427106, 2.2427306,\n",
       "       2.2427526, 2.242774 , 2.242793 , 2.2428086, 2.2428198, 2.2428262,\n",
       "       2.242828 , 2.2428253, 2.2428186, 2.2428088, 2.242797 , 2.2427828,\n",
       "       2.2427683, 2.2427542, 2.2427404, 2.2427278, 2.242717 , 2.2427075,\n",
       "       2.2426996, 2.2426946, 2.2426906, 2.2426882, 2.2426875, 2.2426875,\n",
       "       2.2426884, 2.24269  , 2.2426915, 2.2426932, 2.2426949, 2.2426958,\n",
       "       2.2426972, 2.2426977, 2.2426975, 2.2426972, 2.2426965, 2.2426956,\n",
       "       2.2426946, 2.2426937, 2.2426922, 2.242691 , 2.2426903, 2.2426896,\n",
       "       2.2426887, 2.242688 , 2.242688 , 2.2426877, 2.2426877, 2.2426872,\n",
       "       2.2426875, 2.2426875, 2.2426877, 2.242688 , 2.242688 , 2.242688 ,\n",
       "       2.2426884, 2.2426882, 2.2426882, 2.2426882, 2.2426882, 2.242688 ,\n",
       "       2.2426882, 2.2426877, 2.242688 , 2.2426877, 2.2426877, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426877, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426872, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426872, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426872, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426872, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426872, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426872, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426872, 2.2426872, 2.2426875, 2.2426872, 2.2426872, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426872, 2.2426872, 2.2426872, 2.2426875,\n",
       "       2.2426872, 2.2426872, 2.2426875, 2.2426872, 2.2426875, 2.2426872,\n",
       "       2.2426872, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875, 2.2426875,\n",
       "       2.2426875, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872, 2.2426872,\n",
       "       2.2426875, 2.2426875], dtype=float32)"
      ]
     },
     "execution_count": 315,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run your function to learn the class-conditional standard deviations\n",
    "\n",
    "nlls, scales_arr, class_conditionals_binary = learn_stdevs(x_train, y_train_binary, scales, opt, epochs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 316,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Class conditional means:\n",
      "[[4.987805  3.4292684]\n",
      " [6.249367  2.8670886]]\n",
      "\n",
      "Class conditional standard deviations:\n",
      "[[1.0596133 0.5204331]\n",
      " [1.0596133 0.5204331]]\n"
     ]
    }
   ],
   "source": [
    "# View the distribution parameters\n",
    "\n",
    "print(\"Class conditional means:\")\n",
    "print(class_conditionals_binary.loc.numpy())\n",
    "print(\"\\nClass conditional standard deviations:\")\n",
    "print(class_conditionals_binary.stddev().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 317,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32160 characters]\n",
      "text/plain": [
       "<Figure size 1008x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the loss and convergence of the standard deviation parameters\n",
    "\n",
    "fig, ax = plt.subplots(1, 2, figsize=(14, 5))\n",
    "ax[0].plot(nlls)\n",
    "ax[0].set_title(\"Loss vs epoch\")\n",
    "ax[0].set_xlabel(\"Epoch\")\n",
    "ax[0].set_ylabel(\"Average negative log-likelihood\")\n",
    "for k in [0, 1]:\n",
    "    ax[1].plot(scales_arr[:, k], color=label_colours_binary[k], label=labels_binary[k])\n",
    "ax[1].set_title(\"Standard deviation ML estimates vs epoch\")\n",
    "ax[1].set_xlabel(\"Epoch\")\n",
    "ax[1].set_ylabel(\"Standard deviation\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot the contours of the class-conditional Gaussian distributions as before, this time with just binary labelled data. Notice the contours are the same for each class, just with a different centre location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 318,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84912 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the training data with the class-conditional density contours\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plot_data(x_train, y_train_binary, labels_binary, label_colours_binary)\n",
    "x0_min, x0_max = x_train[:, 0].min(), x_train[:, 0].max()\n",
    "x1_min, x1_max = x_train[:, 1].min(), x_train[:, 1].max()\n",
    "contour_plot((x0_min, x0_max), (x1_min, x1_max), class_conditionals_binary.prob, 2, label_colours_binary)\n",
    "plt.title(\"Training set with class-conditional density contours\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also plot the decision regions for this binary classifier model, notice that the decision boundary is now linear."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 319,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22900 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the model's decision regions\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plot_data(x_train, y_train_binary, labels_binary, label_colours_binary)\n",
    "x0_min, x0_max = x_train[:, 0].min(), x_train[:, 0].max()\n",
    "x1_min, x1_max = x_train[:, 1].min(), x_train[:, 1].max()\n",
    "contour_plot((x0_min, x0_max), (x1_min, x1_max), \n",
    "             lambda x: predict_class(prior_binary, class_conditionals_binary, x), \n",
    "             1, label_colours_binary, levels=[-0.5, 0.5, 1.5],\n",
    "             num_points=500)\n",
    "plt.title(\"Training set with decision regions\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Link to logistic regression\n",
    "\n",
    "In fact, we can see that our predictive distribution $P(Y=y_0 | X)$ can be written as follows:\n",
    "\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "P(Y=y_0 | X) =& ~\\frac{P(X | Y=y_0)P(Y=y_0)}{P(X | Y=y_0)P(Y=y_0) + P(X | Y=y_1)P(Y=y_1)}\\\\\n",
    "=&  ~\\frac{1}{1 + \\frac{P(X | Y=y_1)P(Y=y_1)}{P(X | Y=y_0)P(Y=y_0)}}\\\\\n",
    "=& ~\\sigma(a)\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "where $\\sigma(a) = \\frac{1}{1 + e^{-a}}$ is the sigmoid function, and $a = \\log\\frac{P(X | Y=y_0)P(Y=y_0)}{P(X | Y=y_1)P(Y=y_1)}$ is the _log-odds_.\n",
    "\n",
    "With our additional modelling assumption of a shared covariance matrix $\\Sigma$, it can be shown (using the Gaussian pdf) that $a$ is in fact a linear function of $X$: \n",
    "\n",
    "$$\n",
    "a = w^T X + w_0\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "w =& ~\\Sigma^{-1} (\\mu_0 - \\mu_1)\\\\\n",
    "w_0 =& -\\frac{1}{2}\\mu_0^T \\Sigma^{-1}\\mu_0 + \\frac{1}{2}\\mu_1^T\\Sigma^{-1}\\mu_1 + \\log\\frac{P(Y=y_0)}{P(Y=y_1)}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "The model therefore takes the form $P(Y=y_0 | X) = \\sigma(w^T X + w_0)$, with weights $w\\in\\mathbb{R}^2$ and bias $w_0\\in\\mathbb{R}$. This is the form used by logistic regression, and explains why the decision boundary above is linear. \n",
    "\n",
    "In the above we have outlined the derivation of the generative logistic regression model. The parameters are typically estimated with maximum likelihood, as we have done. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we will use the above equations to directly parameterise the output Bernoulli distribution of the generative logistic regression model.\n",
    "\n",
    "You should now write the following function, according to the following specification:\n",
    "\n",
    "* The inputs to the function are:\n",
    "  * the prior distribution `prior` over the two classes\n",
    "  * the (batched) class-conditional distribution `class_conditionals`\n",
    "* The function should use the parameters of the above distributions to compute the weights and bias terms $w$ and $w_0$ as above\n",
    "* The function should then return a tuple of two numpy arrays for $w$ and $w_0$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 326,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### GRADED CELL ####\n",
    "\n",
    "# Complete the following function. \n",
    "# Make sure to not change the function name or arguments.\n",
    "\n",
    "def get_logistic_regression_params(prior, class_conditionals):\n",
    "    \"\"\"\n",
    "    This function takes the prior distribution and class-conditional distribution as inputs.\n",
    "    This function should compute the weights and bias terms of the generative logistic\n",
    "    regression model as above, and return them in a 2-tuple of numpy arrays of shapes\n",
    "    (2,) and () respectively.\n",
    "    \"\"\"\n",
    "    sigma = tf.pow(class_conditionals.parameters['scale_diag'],2)\n",
    "    inv_sigma = np.eye(2) * (1 / sigma)\n",
    "    mu0 = class_conditionals.parameters['loc'][0]\n",
    "    mu1 = class_conditionals.parameters['loc'][1]\n",
    "    \n",
    "    w = np.matmul(inv_sigma, (mu0 - mu1))\n",
    "    w0 = - 0.5 * (np.matmul(tf.transpose(mu0), np.matmul(inv_sigma, mu0)))\\\n",
    "         + 0.5 * (np.matmul(tf.transpose(mu1), np.matmul(inv_sigma, mu1)))\\\n",
    "         + np.log(prior.parameters['probs'][0] / prior.parameters['probs'][1])\n",
    "    return w, w0\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 327,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run your function to get the logistic regression parameters\n",
    "\n",
    "w, w0 = get_logistic_regression_params(prior_binary, class_conditionals_binary)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now use these parameters to make a contour plot to display the predictive distribution of our logistic regression model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 328,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61072 characters]\n",
      "text/plain": [
       "<Figure size 720x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Plot the training data with the logistic regression prediction contours\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 6))\n",
    "plot_data(x_train, y_train_binary, labels_binary, label_colours_binary)\n",
    "x0_min, x0_max = x_train[:, 0].min(), x_train[:, 0].max()\n",
    "x1_min, x1_max = x_train[:, 1].min(), x_train[:, 1].max()\n",
    "X0, X1 = get_meshgrid((x0_min, x0_max), (x1_min, x1_max))\n",
    "\n",
    "logits = np.dot(np.array([X0.ravel(), X1.ravel()]).T, w) + w0\n",
    "Z = tf.math.sigmoid(logits)\n",
    "lr_contour = ax.contour(X0, X1, np.array(Z).T.reshape(*X0.shape), levels=10)\n",
    "ax.clabel(lr_contour, inline=True, fontsize=10)\n",
    "contour_plot((x0_min, x0_max), (x1_min, x1_max), \n",
    "             lambda x: predict_class(prior_binary, class_conditionals_binary, x), \n",
    "             1, label_colours_binary, levels=[-0.5, 0.5, 1.5],\n",
    "             num_points=300)\n",
    "plt.title(\"Training set with prediction contours\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations on completing this programming assignment! In the next week of the course we will look at Bayesian neural networks and uncertainty quantification."
   ]
  }
 ],
 "metadata": {
  "coursera": {
   "course_slug": "probabilistic-deep-learning-with-tensorflow2",
   "graded_item_id": "D6ITt",
   "launcher_item_id": "DExYG"
  },
  "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
