day2/notebooks/CAVI-linreg.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/PGM-Lab/2023-probai-private/blob/main/Day2-BeforeLunch/notebooks/CAVI-linreg.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-11T11:49:17.876357Z",
     "start_time": "2023-05-11T11:49:13.846684Z"
    },
    "id": "suGRNn-Pq3Se",
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install -q scipy pandas seaborn matplotlib numpy\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "%matplotlib inline\n",
    "import matplotlib.colors as mcolors\n",
    "from scipy.stats import multivariate_normal, norm\n",
    "import time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "#from IPython.display import display, clear_output\n",
    "import ssl\n",
    "ssl._create_default_https_context = ssl._create_unverified_context\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LzKwg-lUq3Sg"
   },
   "source": [
    "# 1. Dataset and miscellaneous\n",
    "\n",
    "The following example is taken from \\[1\\].  We would like to explore the relationship between topographic heterogeneity of a nation as measured by the Terrain Ruggedness Index (variable *rugged* in the dataset) and its GDP per capita. In particular, it was noted by the authors in \\[1\\] that terrain ruggedness or bad geography is related to poorer economic performance outside of Africa, but rugged terrains have had a reverse effect on income for African nations. Let us look at the data \\[2\\] and investigate this relationship.  We will be focusing on three features from the dataset:\n",
    "  - `rugged`: quantifies the Terrain Ruggedness Index\n",
    "  - `cont_africa`: whether the given nation is in Africa\n",
    "  - `rgdppc_2000`: Real GDP per capita for the year 2000\n",
    " \n",
    "  \n",
    "We will take the logarithm for the response variable GDP as it tends to vary exponentially. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "-P2oriwKq3Sh",
    "outputId": "ca32be2e-de45-463e-83b2-36b3c770604a",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of african countries: 49\n",
      "Number of non-african countries: 121\n"
     ]
    }
   ],
   "source": [
    "np.random.seed(42)\n",
    "DATA_URL = \"https://d2hg8soec8ck9v.cloudfront.net/datasets/rugged_data.csv\"\n",
    "data = pd.read_csv(DATA_URL, encoding=\"ISO-8859-1\")\n",
    "df = data[[\"cont_africa\", \"rugged\", \"rgdppc_2000\"]]\n",
    "df = df.sample(frac=1)\n",
    "df = df[np.isfinite(df.rgdppc_2000)]\n",
    "df[\"rgdppc_2000\"] = np.log(df[\"rgdppc_2000\"])\n",
    "df = df[[\"cont_africa\", \"rugged\", \"rgdppc_2000\"]]\n",
    "\n",
    "# Divide the data into predictors and response and store the data in numpy arrays organized in a dictionary with one \n",
    "# element for african and non-african nations, respectively.\n",
    "data = np.array(df)\n",
    "x_data = {'non-african': data[data[:, 0] == 0, 1].reshape(-1,1), 'african': data[data[:, 0] == 1, 1].reshape(-1,1)}\n",
    "y_data = {'non-african': data[data[:, 0] == 0, -1], 'african': data[data[:, 0] == 1, -1]}\n",
    "\n",
    "print(f\"Number of african countries: {x_data['african'].shape[0]}\")\n",
    "print(f\"Number of non-african countries: {x_data['non-african'].shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "y0gtwqk5q3Si",
    "outputId": "88d8daa0-6991-4b89-b5fe-a8d7fddc2b6c",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Display first 10 entries \n",
    "display(df[0:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u6ic0DZ4q3Si",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def prepare_figure(title='Scatter plot of data', x_data_: dict = None, y_data_: dict = None):\n",
    "    \"\"\"\n",
    "    Plot the data and return the figure axis for possible subsequent additional plotting.\n",
    "    :param title: Title of the plot\n",
    "    :param x_data_: dictionary with data for the prdictor variable\n",
    "    :param y_data_: dictionary with data for the response variable.\n",
    "    :return: Figure axis.\n",
    "    \"\"\"\n",
    "    if x_data_ is None and y_data_ is None:\n",
    "        x_data_ = x_data\n",
    "        y_data_ = y_data\n",
    "\n",
    "    fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 6), sharey=True)\n",
    "    fig.suptitle(title, fontsize=16)\n",
    "\n",
    "    for idx, cont in enumerate(x_data):\n",
    "        ax[idx].scatter(x_data_[cont], y_data_[cont], c='black')\n",
    "        ax[idx].set(xlabel=\"Terrain Ruggedness Index\", ylabel=\"log GDP (2000)\",\n",
    "                    title=f\"{cont} nations\")\n",
    "\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "otzsrc69q3Sj",
    "outputId": "633d8b44-43c4-440f-fbf5-db2b3076b27f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Let's display the data\n",
    "prepare_figure('Scatter plot of data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fv7iu1ZXq3Sj"
   },
   "source": [
    "## 1.1 Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SYRHU_xkq3Sk",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_data_subsets(x_data_: dict, y_data_: dict, frac: float = 1, verbose=True):\n",
    "    \"\"\"\n",
    "    Sample a subset of the data sets. Both data sets are again dictionaries with elements for\n",
    "    african and non-african nations.\n",
    "    :param x_data_: Dictionary with data for the predictor variable\n",
    "    :param y_data_: Dictionary with data for the response variable\n",
    "    :param frac: fraction of the data to be sampled\n",
    "    :param verbose: How much detail should be displayed\n",
    "    :return: Four dictionaries, two for each of the subsets.\n",
    "    \"\"\"\n",
    "\n",
    "    x_data_ic = {}\n",
    "    y_data_ic = {}\n",
    "    x_data_ec = {}\n",
    "    y_data_ec = {}\n",
    "\n",
    "    # Iterate over the two keys (african and non-african) and pick out the data\n",
    "    for cont in x_data_.keys():\n",
    "        indices = np.random.choice(x_data_[cont].shape[0], int(x_data_[cont].shape[0] * frac), replace=False)\n",
    "        mask = np.zeros(x_data_[cont].shape[0], dtype=bool)\n",
    "        mask[indices] = True\n",
    "        x_data_ic[cont] = x_data_[cont][mask, :]\n",
    "        y_data_ic[cont] = y_data_[cont][mask]\n",
    "        x_data_ec[cont] = x_data_[cont][~mask, :]\n",
    "        y_data_ec[cont] = y_data_[cont][~mask]\n",
    "\n",
    "    if verbose:\n",
    "        for cont in x_data_ic.keys():\n",
    "            print(f\"{cont} entries: {x_data_ic[cont].shape[0]}\")\n",
    "\n",
    "    return x_data_ic, y_data_ic, x_data_ec, y_data_ec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hJHoupAsq3Sk",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def plot_model(ax, x_data_: dict = None, model: dict = None, linewith: int = 2,\n",
    "               label: str = None, line_color: str = 'red', num_samples: int = 0):\n",
    "    \"\"\"\n",
    "    Plot the model using the figure axes in ax\n",
    "    :param ax: Figure axes\n",
    "    :param x_data_: dictionary with data for the predictor variable\n",
    "    :param model: The model (i.e., weights/parameters and possibly covariance) in a dictionary with one element for\n",
    "    african and non-african.\n",
    "    :param linewith: width of line to be plotted\n",
    "    :param label: the label to go with the model plot\n",
    "    :param line_color: color of line to be plotted\n",
    "    :param num_samples: If the model also contains a covariance matrix we treat the weights 'w' as mean values and\n",
    "    sample realizations of the model from the weight distribution.\n",
    "    :return: Figure axes.\n",
    "    \"\"\"\n",
    "    if x_data_ is None:\n",
    "        x_data_ = x_data\n",
    "\n",
    "    if model is None:\n",
    "        print(\"Ups, no model supplied\")\n",
    "    for idx, cont in enumerate(model):\n",
    "        # We treat the first element in the weight matrix as the bias hence we add a constant one column to the data\n",
    "        x_aug = np.c_[np.ones(x_data_[cont].shape[0]), x_data_[cont]]\n",
    "\n",
    "        # Sample from the distribution over the weights if a covariance matrix is given.\n",
    "        if num_samples > 0 and 'cov' in model[cont]:\n",
    "            for _ in range(num_samples):\n",
    "                w_sample = np.random.multivariate_normal(mean=model[cont]['w'],\n",
    "                                                         cov=model[cont]['cov'])\n",
    "                ax[idx].plot(x_data[cont], x_aug @ w_sample , 'g-', alpha=.05)\n",
    "\n",
    "        ax[idx].plot(x_data_[cont], x_aug @ model[cont]['w'],\n",
    "                     color=line_color,\n",
    "                     linewidth=linewith,\n",
    "                     label=label)\n",
    "        if label is not None:\n",
    "            ax[idx].legend()\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KVb7gSjJq3Sl"
   },
   "source": [
    "# 1. Linear regression\n",
    "\n",
    "Linear regression is one of the most commonly used supervised learning tasks in machine learning. Suppose we are given a dataset $\\mathcal{D}$ of the form\n",
    "\n",
    "$$ \\mathcal{D}  = \\{ (X_i, y_i) \\} \\qquad \\text{for}\\qquad i=1,2,...,N$$\n",
    "\n",
    "The goal of linear regression is to fit a function to the data of the form:\n",
    "\n",
    "$$ y = w X + b + \\epsilon, $$\n",
    "\n",
    "where $w$ and $b$ are learnable parameters and $\\epsilon$ represents observation noise. Specifically $w$ is a matrix of weights and $b$ is a bias vector.\n",
    "\n",
    "Let's first implement linear regression in PyTorch and learn point estimates for the parameters $w$ and $b$.  Then we'll see how to incorporate uncertainty into our estimates."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OOCxe4I0q3Sl"
   },
   "source": [
    "## 1.1 Model\n",
    "We would like to predict log GDP per capita of a nation as a function of whether the nation is in Africa and its Terrain Ruggedness Index. As indicated by the data partitioning above, we will make one linear regression model for african and non-african nations, respectively. \n",
    "\n",
    "To sprcify our regression model, we will define a specific object encapsulating the model.  Our input `x_data` is a tensor of size $N \\times 1$ and our output `y_data` is a tensor of size $N \\times 1$.  The method `predict(self,x_data)` defines a linear transformation of the form $Xw + b$ where $w$ is the weight matrix and $b$ is the additive bias.\n",
    "\n",
    "The parameters of the model are defined using ``torch.nn.Parameter``, and will be learned during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Krf1kLuwq3Sl",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class RegressionModel():\n",
    "    def __init__(self):\n",
    "        self.w = torch.nn.Parameter(torch.zeros(1, 1))\n",
    "        self.b = torch.nn.Parameter(torch.zeros(1, 1))\n",
    "\n",
    "    def params(self):\n",
    "        return {\"b\": self.b, \"w\": self.w}\n",
    "\n",
    "    def predict(self, x_data):\n",
    "        return (self.b + torch.mm(self.w, torch.t(x_data))).squeeze(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9jnKOmtaq3Sl"
   },
   "source": [
    "## 1.2 Training\n",
    "We will use the mean squared error (MSE) as our loss and Adam as our optimizer. We would like to optimize the parameters of the `regression_model` neural net above. We will use a somewhat large learning rate of `0.05` and run for 1000 iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sUnxf5oJq3Sl",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def least_squares_solution(x_data, y_data, verbose=True):\n",
    "    regression_model = RegressionModel()\n",
    "    loss_fn = torch.nn.MSELoss(reduction='sum')\n",
    "    optim = torch.optim.Adam(regression_model.params().values(), lr=0.05)\n",
    "    num_iterations = 1000\n",
    "\n",
    "    param = {}\n",
    "    for cont in x_data.keys():\n",
    "        param[cont] = {}\n",
    "        if verbose:\n",
    "            print(f\"Learning model for {cont} nations\")\n",
    "        for j in range(num_iterations):\n",
    "            # run the model forward on the data\n",
    "            y_pred = regression_model.predict(torch.tensor(x_data[cont], dtype=torch.float))\n",
    "            # calculate the mse loss\n",
    "            loss = loss_fn(y_pred, torch.tensor(y_data[cont], dtype=torch.float))\n",
    "            # initialize gradients to zero\n",
    "            optim.zero_grad()\n",
    "            # backpropagate\n",
    "            loss.backward()\n",
    "            # take a gradient step\n",
    "            optim.step()\n",
    "            if (j + 1) % 500 == 0 and verbose:\n",
    "                print(\"[iteration %04d] loss: %.4f\" % (j + 1, loss.item()))\n",
    "\n",
    "        # Store the learned parameters \n",
    "        param[cont]['w'] = np.r_[regression_model.params()['b'].detach().numpy().copy().flatten(),\n",
    "                      np.transpose(regression_model.params()['w'].detach().numpy()).copy().flatten()]\n",
    "    # Inspect learned parameters\n",
    "    if verbose:\n",
    "        print(\"Learned parameters:\")\n",
    "        for cont in param.keys():\n",
    "            print(f\"{cont}: weights = {param[cont]['w']}\")\n",
    "\n",
    "    return param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HCuJdncNq3Sm",
    "outputId": "9431fc76-bb86-4fd6-afd7-20e00df13863",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Learn the model\n",
    "model_lr = least_squares_solution(x_data, y_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9SlgGJSRq3Sm"
   },
   "source": [
    "## 1.3 Evaluating the model\n",
    "\n",
    "We now plot the regression line learned for african and non-african nations relating the rugeedness index with the GDP of the country."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JorVe430q3Sm",
    "outputId": "8545b36e-a7e9-4b34-ef41-c84e522438a9",
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = prepare_figure(title='Regression line')\n",
    "plot_model(ax, model=model_lr)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HdAGo3yEq3Sm"
   },
   "source": [
    "## 1.4 The relationship between ruggedness and log GPD\n",
    "\n",
    "Using this analysis, we can estimate the relationship between ruggedness and log GPD. As can be seen, this relationship is positive for African nations, but negative for Non African Nations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jABQAd3Uq3Sm",
    "outputId": "14d9cd42-d77d-47f9-fccc-8bd398436af4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"Slope for non-african nations: {model_lr['non-african']['w'][1]}\")\n",
    "print(f\"Slope for african nations: {model_lr['african']['w'][1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ibw169Snq3Sm"
   },
   "source": [
    "## 1.5 The effect of outliers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hdtrwsgEq3Sm",
    "outputId": "c15248ea-c688-4157-863d-b23838a28ed8",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add a couple of 'outliers' to the data\n",
    "x_data_with_outliers = x_data.copy()\n",
    "y_data_with_outliers = y_data.copy()\n",
    "x_data_with_outliers['african'] = np.concatenate((x_data['african'], [[6]], [[5.5]]), axis=0)\n",
    "x_data_with_outliers['non-african'] = np.concatenate((x_data['non-african'], [[4]], [[5.2]]), axis=0)\n",
    "y_data_with_outliers['african'] = np.concatenate((y_data['african'], [5.2], [6.3]), axis=0)\n",
    "y_data_with_outliers['non-african'] = np.concatenate((y_data['non-african'], [10], [11]), axis=0)\n",
    "\n",
    "# Learn a new model with the modified dataset\n",
    "model_outlier = least_squares_solution(x_data_with_outliers, y_data_with_outliers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "MN3dw3ujq3Sn",
    "outputId": "3f206a5b-0380-45f0-bfa6-9a775b71692d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = prepare_figure(\"Linear regression with outliers\")\n",
    "ax[0].scatter(x_data_with_outliers['non-african'][-2:],\n",
    "              y_data_with_outliers['non-african'][-2:], c='green', label='Outliers')\n",
    "ax[1].scatter(x_data_with_outliers['african'][-2:],\n",
    "              y_data_with_outliers['african'][-2:], c='green', label='Outliers')\n",
    "\n",
    "plot_model(ax, x_data_with_outliers, model=model_outlier, label='Model influenced by outliers', line_color='green')\n",
    "plot_model(ax, x_data_with_outliers, model=model_lr, label='Previous model')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "B1tx5Z-Rq3Sn",
    "outputId": "4d3d331e-59a2-4e5c-8318-52af5db5954a",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# What can we now say about the slope?\n",
    "print(f\"Slope for non-african nations: {model_outlier['non-african']['w'][1]}\")\n",
    "print(f\"Slope for african nations: {model_outlier['african']['w'][1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "r2eyN8q6q3Sn"
   },
   "source": [
    "# 2. Bayesian linear regressison\n",
    "\n",
    "Following the approach from earlier today, we will model the data using a Bayesian linear regression model:\n",
    "<img src=\"Bayesian_linear_regression.png\" style=\"width: 400px;\"/>\n",
    "\n",
    "The quantitative part of the model is specified as: \n",
    "- Number of data dim: $M$\n",
    "- Number of data inst: $N$\n",
    "- $Y_{i}|\\{{\\bf w}, {\\bf x}_i, \\theta \\} \\sim \\mathcal{N}({\\bf w}^T{\\bf x}_i, 1/\\theta)$    \n",
    "- ${\\bf W} \\sim {\\mathcal N}({\\bf 0}, \\gamma_w^{-1}{\\bf I}_{M\\times M})$\n",
    "- $b\\sim {\\mathcal N}(0,\\gamma_b^{-1})$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DzvuJhb_q3Sn"
   },
   "source": [
    "## 2.1 Helper-routine: Calculate ELBO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kamfCqt5q3Sn",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def calculate_ELBO(x_data, y_data, gamma_w, gamma_b, theta, q_w_mean, q_w_prec, q_b_mean, q_b_prec):\n",
    "    \"\"\"\n",
    "    Helper routine: Calculate ELBO. Data is the sampled x and y values, gamma is the prior precision over the \n",
    "    weights and theta is the prior precision associated with y. Everything prefixed a 'q' relates to the \n",
    "    variational posterior.\n",
    "    \n",
    "    Note: This function obviously only works for this particular model and is not a general solution.\n",
    "\n",
    "    :param x_data: The predictors\n",
    "    :param y_data: The response variable\n",
    "    :param gamma_w: prior precision for the weights\n",
    "    :param gamma_b: prior precision for the intercept\n",
    "    :param theta: prior precision for y\n",
    "    :param q_w_mean: VB posterior mean for the distribution of the weights w \n",
    "    :param q_w_prec: VB posterior precision for the distribution of the weights w \n",
    "    :param q_b_mean: VB posterior mean for the intercept b\n",
    "    :param q_b_prec: VB posterior precision for the intercept b\n",
    "    :return: the ELBO\n",
    "    \"\"\"\n",
    "    \n",
    "    # We calculate the ELBO as E_q log p(y,x,w,b) - E_q log q(w,b), where\n",
    "    # log p(y,x,w) = sum_i log p(y|x,w,b) + log p(w) + log p(b)\n",
    "    # log q(w,b) = log q(w) + log q(b)\n",
    "\n",
    "    M = x_data.shape[1]\n",
    "\n",
    "    # E_q log p(w)\n",
    "    E_log_p = -0.5 * M * np.log(2 * np.pi) + 0.5 * M * np.log(gamma_w) - 0.5 * gamma_w * np.sum(np.diagonal(np.linalg.inv(q_w_prec))\n",
    "                                                                                    + (q_w_mean*q_w_mean).flatten())\n",
    "    # E_q log p(b)\n",
    "    E_log_p += -0.5 * np.log(2 * np.pi) + 0.5 * np.log(gamma_b) - 0.5 * gamma_b * (1/q_b_prec + q_b_mean**2)\n",
    "\n",
    "    # sum_i E_q log p(y|x,w,b)\n",
    "    E_w_w = np.linalg.inv(q_w_prec) + q_w_mean @ q_w_mean.transpose()\n",
    "    E_b_b = 1/q_b_prec + q_b_mean**2\n",
    "    for i in range(x_data.shape[0]):\n",
    "        E_x_ww_x = np.matmul(x_data[i, :].transpose(), np.matmul(E_w_w, x_data[i, :]))\n",
    "        E_log_p += -0.5 * np.log(2 * np.pi) + 0.5 * np.log(theta) \\\n",
    "                   - 0.5 * theta * (y_data[i]**2 + E_x_ww_x + E_b_b\n",
    "                                    + 2 * q_b_mean * np.matmul(q_w_mean.transpose(), x_data[i, :])\n",
    "                                    - 2 * y_data[i] * np.matmul(q_w_mean.transpose(), x_data[i,:])\n",
    "                                    - 2 * y_data[i] * q_b_mean)\n",
    "\n",
    "    # Entropy of q_b\n",
    "    ent = 0.5 * np.log(1 * np.pi * np.exp(1) / q_b_prec)\n",
    "    ent += 0.5 * np.log(np.linalg.det(2 * np.pi * np.exp(1) * np.linalg.inv(q_w_prec)))\n",
    "\n",
    "    return E_log_p - ent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UqJnjtXgq3Sn"
   },
   "source": [
    "## 2.2 Variational solution\n",
    "First we consider a full mean filed approach, where the variational approximation factorizes as\n",
    "$$\n",
    "q({\\bf w}, b) = q(b)\\prod _{i=1}^Mq(w_i)\n",
    "$$\n",
    "Actually, in our specific case we only have one weight besides the bias, hence the approximation actually simplies to\n",
    "$$\n",
    "q(w, b) = q(b)q(w).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_fXw7Ck_q3Sn"
   },
   "source": [
    "### 2.2.1 Updating equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iGOsXAQJq3Sn",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The variational updating rule for weight component\n",
    "def update_w_comp(x_data, y_data, gamma_w, theta, q_w_mean, q_w_prec, q_b_mean, comp):\n",
    "    \"\"\"\n",
    "    The variational updating rule for a specific weight component. This updating function is defined for a general\n",
    "    Bayesian linear regression model, which can contain multiple predictor variables. 'comp' is used to index the\n",
    "    specific component to be updated.\n",
    "    :param x_data: A numpy array of size N x M with data for the predictor variable. In our case, M=1.\n",
    "    :param y_data: A numpy array of size N x 1.\n",
    "    :param gamma_w: precision for the weights.\n",
    "    :param theta: prior precision for y.\n",
    "    :param q_w_mean: VB posterior mean for the distribution of the weights w\n",
    "    :param q_w_prec: VB posterior precision for the distribution of the weights w\n",
    "    :param q_b_mean: VB posterior mean for the intercept b\n",
    "    :param comp: The weight component to be updated.\n",
    "    :return: The updated precision and mean value for the weight component.\n",
    "    \"\"\"\n",
    "    # Lenght of weight vector\n",
    "    M = x_data.shape[1]\n",
    "    # The precision (a scalar)\n",
    "    tau = gamma_w\n",
    "    # The mean (a scalar)\n",
    "    mu = 0.0\n",
    "    for i in range(x_data.shape[0]):\n",
    "        tau += theta * x_data[i, comp]**2\n",
    "        mu += (y_data[i] - q_b_mean - (np.sum(x_data[i, :] @ q_w_mean) - x_data[i, comp]*q_w_mean[comp])) \\\n",
    "                * x_data[i, comp]\n",
    "    mu = theta * 1/tau * mu\n",
    "\n",
    "    # Update the appropriate entries in the mean vector and precision matrix\n",
    "    q_w_prec[comp, comp] = tau\n",
    "    q_w_mean[comp] = mu.item()\n",
    "\n",
    "    return q_w_prec, q_w_mean\n",
    "\n",
    "# The variational updating rule for the intercept\n",
    "def update_b(x_data, y_data, gamma_b, theta, q_w_mean):\n",
    "    \"\"\"\n",
    "    The variational updating rule for  the bias term\n",
    "    :param x_data: A numpy array of size N x M with data for the predictor variable. In our case, M=1.\n",
    "    :param y_data: A numpy array of size N x 1.\n",
    "    :param gamma_b: precision for the bias term\n",
    "    :param theta: prior precision for y\n",
    "    :param q_w_mean: VB posterior mean for the distribution of the weights w\n",
    "    :return: The updated precision and mean value for the bias term.\n",
    "    \"\"\"\n",
    "    # The precision (a scalar)\n",
    "    tau = (gamma_b + theta * x_data.shape[0])\n",
    "    # The mean (a scalar)\n",
    "    mu = 0\n",
    "    for i in range(x_data.shape[0]):\n",
    "        mu += (y_data[i] - q_w_mean.transpose() @ x_data[i, :])\n",
    "    mu = 1/tau * theta * mu\n",
    "\n",
    "    return tau, mu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PjEYw22tq3So"
   },
   "source": [
    "### 2.2.2 Do the VB (full mean field)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nCgaOivgq3So",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def variational_solution(x_data, y_data, gamma_w: float = 1, gamma_b: float = 1, theta: float = 1, verbose=True):\n",
    "    \"\"\"\n",
    "    Perform variational inference using the coordinate ascent updating rules. The two models (for african and\n",
    "    non-african nations are learned independently.\n",
    "    :param x_data: dictionary with data for the predictor variable\n",
    "    :param y_data: dictionary with data for the response variable\n",
    "    :param gamma_w: precision for the weight/slope term\n",
    "    :param gamma_b: precision for the bias term\n",
    "    :param theta: prior precision for y\n",
    "    :param verbose: Do we want to display information from the updating process?\n",
    "    :return: The model (i.e., mean and covariance) in a dictionary with one element for african and non-african.\n",
    "    \"\"\"\n",
    "    param = {}\n",
    "\n",
    "    if verbose:\n",
    "        fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 6), sharey=True)\n",
    "        fig.suptitle(\"Evolution of ELBO\", fontsize=16)\n",
    "        \n",
    "     # Do the learning for african and non-african nations\n",
    "    for idx, cont in enumerate(x_data):\n",
    "        param[cont] = {}\n",
    "        M = x_data[cont].shape[1]\n",
    "        q_w_mean = np.random.normal(0, 1, (1, 1))\n",
    "        q_w_prec = np.array([[1]])\n",
    "        q_b_mean = np.random.normal(0, 1)\n",
    "        q_b_prec = 1\n",
    "        elbos = []\n",
    "\n",
    "        # Calculate ELBO\n",
    "        this_lb = calculate_ELBO(x_data[cont], y_data[cont], gamma_w, gamma_b, theta, q_w_mean, q_w_prec, q_b_mean, q_b_prec)\n",
    "        elbos.append(this_lb)\n",
    "        previous_lb = -np.inf\n",
    "        # Start iterating\n",
    "        if verbose:\n",
    "            print(\"\\n\" + 100 * \"=\" + \"\\n   VB iterations:\\n\" + 100 * \"=\")\n",
    "        for iteration in range(100):\n",
    "\n",
    "            # Update the variational distributions\n",
    "            for i in range(M):\n",
    "                q_w_prec, q_w_mean = update_w_comp(x_data[cont], y_data[cont], gamma_w, theta, q_w_mean, q_w_prec, q_b_mean, i)\n",
    "            q_b_prec, q_b_mean = update_b(x_data[cont], y_data[cont], gamma_b, theta, q_w_mean)\n",
    "\n",
    "            this_lb = calculate_ELBO(x_data[cont], y_data[cont], gamma_w, gamma_b, theta, q_w_mean, q_w_prec, q_b_mean, q_b_prec)\n",
    "            elbos.append(this_lb)\n",
    "            if verbose:\n",
    "                print(f\"Iteration {iteration:2d}. ELBO: {this_lb.item():13.7f}\")\n",
    "            if this_lb < previous_lb:\n",
    "                raise ValueError(\"ELBO is decreasing. Something is wrong! Goodbye...\")\n",
    "\n",
    "            if iteration > 0 and np.abs((this_lb - previous_lb) / previous_lb) < 1E-8:\n",
    "                # Very little improvement. We are done.\n",
    "                break\n",
    "\n",
    "            # If we didn't break we need to run again. Update the value for \"previous\"\n",
    "            previous_lb = this_lb\n",
    "        if verbose:\n",
    "            print(\"\\n\" + 100 * \"=\" + \"\\n\")\n",
    "\n",
    "        param[cont]['w'] = np.r_[q_b_mean, q_w_mean.flatten()]\n",
    "        param[cont]['cov'] = np.linalg.inv(np.diag(np.r_[q_b_prec, q_w_prec.flatten()]))\n",
    "\n",
    "        if verbose:\n",
    "            ax[idx].plot(range(len(elbos)), elbos, label='Mean field approximation: ELBO')\n",
    "            ax[idx].set(xlabel=\"Number of iterations\",\n",
    "                        ylabel=\"ELBO\",\n",
    "                        title=f\"{cont} nations\")\n",
    "            ax[idx].legend()\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "    # Inspect learned parameters\n",
    "    if verbose:\n",
    "        print(\"Learned parameters:\")\n",
    "        for cont in param.keys():\n",
    "            print(f\"For {cont} nations\")\n",
    "            for name, value in param[cont].items():\n",
    "                print(f\"{name} = \")\n",
    "                print(f\"{value}\")\n",
    "\n",
    "    return param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qTMPc3s-q3So",
    "outputId": "a520fd88-0d90-4b19-d275-fc8e2508fd5f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_mf = variational_solution(x_data, y_data, gamma_w=1, gamma_b=1, theta=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8xgreYnSq3So"
   },
   "source": [
    "### 2.2.3 Model evaluation\n",
    "\n",
    "To get a sense of the robustness of the model we draw samples from the posterior variational distributions over the weights and intercept; each sample correspond to a regression line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pU4b4OA6q3So",
    "outputId": "cf5c6899-3415-4579-b2f2-d8704ab5b01b",
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = prepare_figure(\"Sampling parameters from the mean field variational distribution\")\n",
    "plot_model(ax, model=model_mf, num_samples=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_mZfI1Dtq3So"
   },
   "source": [
    "## 2.3 Bayes solution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ctyAJXddq3So"
   },
   "source": [
    "### 2.3.1 Updating equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R4jXnTEgq3So",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Calculate the exact posterior over the model parameters\n",
    "def bayes_solution(x_data: dict, y_data: dict, gamma_w: float = 1, gamma_b: float = 1, theta: float = 1, verbose=True):\n",
    "    \"\"\"\n",
    "    Calculate the exact posterior over the model parameters\n",
    "    :param x_data: dictionary with data for the predictor variable\n",
    "    :param y_data: dictionary with data for the response variable\n",
    "    :param gamma_w: precision for the weight/slope term\n",
    "    :param gamma_b: precision for the bias term\n",
    "    :param theta: prior precision for y\n",
    "    :param verbose: Do we want to display information from the updating process?\n",
    "    :return: The learned model in the form of a dictionary\n",
    "    \"\"\"\n",
    "    param = {}\n",
    "    for cont in x_data.keys():\n",
    "        param[cont] = {}\n",
    "        x_aug = np.c_[np.ones(x_data[cont].shape[0]), x_data[cont]]\n",
    "        prior_cov = np.diag([1. / gamma_b] + [1. / gamma_w])\n",
    "        prior_mean = np.zeros(x_aug.shape[1])\n",
    "        param[cont]['cov'] = np.linalg.inv(np.linalg.inv(prior_cov) + theta * np.transpose(x_aug) @ x_aug)\n",
    "        param[cont]['w'] = param[cont]['cov'] @ (np.linalg.inv(prior_cov) @ prior_mean + theta * np.transpose(x_aug) @ y_data[cont])\n",
    "\n",
    "    # Inspect learned parameters\n",
    "    if verbose:\n",
    "        print(\"Learned parameters:\")\n",
    "        for cont in param.keys():\n",
    "            print(f\"For {cont} nations\")\n",
    "            for name, value in param[cont].items():\n",
    "                print(f\"{name} = \")\n",
    "                print(f\"{value}\")\n",
    "\n",
    "    return param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "git6UDrEq3So",
    "outputId": "ba121cca-0eac-4b4b-9032-a145a46b10b6",
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_bayes = bayes_solution(x_data, y_data, gamma_w=1, gamma_b=1, theta=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_lvuwMcIq3So"
   },
   "source": [
    "### 2.3.1 Model evaluation\n",
    "\n",
    "To get a sense of the robustness of the model we draw samples from the posterior distributions over the weights and intercept; each sample correspond to a regression line"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RvZ4c7W9q3So",
    "outputId": "6f130e64-8eb2-403e-a098-df9bc2c8651d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = prepare_figure(\"Sampling parameters from the Bayes solution\")\n",
    "plot_model(ax, model=model_bayes, num_samples=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yEqJr3wpq3So"
   },
   "source": [
    "Recreating the plots for the variational solution can make the comparison easier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Mu3_CIKSq3Sp",
    "outputId": "d2eb9339-4523-4e27-8625-3046c8c9aae8",
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = prepare_figure(\"Sampling parameters from the mean field variational distribution\")\n",
    "plot_model(ax, model=model_mf, num_samples=100)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "spCWxb_Hq3Sp"
   },
   "source": [
    "# 3 Model comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yE26zWIyq3Sp"
   },
   "source": [
    "## 3.1 Additional helper functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7ICLSuDDq3Sp"
   },
   "source": [
    "### 3.1.1 Visualizing bivariate densities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qbU2U1uYq3Sp",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def density_plt(model_list: list):\n",
    "    \"\"\"\n",
    "    Make a density plot of the models in the list. The first model is assumed to be the Bayes solution, based on which \n",
    "    the coordinate axes are defined. \n",
    "    \"\"\"\n",
    "\n",
    "    fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))\n",
    "    fig.suptitle(\"Density plots\", fontsize=16)\n",
    "\n",
    "    # We assume that the first model in the list is the Bayes model\n",
    "    model_bayes = model_list[0]\n",
    "    x_mesh = {}\n",
    "    y_mesh = {}\n",
    "    for cont in model_bayes.keys():\n",
    "        x = np.linspace(model_bayes[cont]['w'][0] - 5 * model_bayes[cont]['cov'][0, 0],\n",
    "                        model_bayes[cont]['w'][0] + 5 * model_bayes[cont]['cov'][0, 0], 100)\n",
    "        y = np.linspace(model_bayes[cont]['w'][1] - 5 * model_bayes[cont]['cov'][1, 1],\n",
    "                        model_bayes[cont]['w'][1] + 5 * model_bayes[cont]['cov'][1, 1], 100)\n",
    "        x_mesh[cont], y_mesh[cont] = np.meshgrid(x, y)\n",
    "\n",
    "    for model_idx, model in enumerate(model_list):\n",
    "        for idx, cont in enumerate(model):\n",
    "            dist = multivariate_normal(mean=model[cont]['w'], cov=model[cont]['cov'])\n",
    "            input_vector = np.vstack([x_mesh[cont].flatten(), y_mesh[cont].flatten()]).T\n",
    "            value = dist.pdf(input_vector).reshape(x_mesh[cont].shape)\n",
    "            ax[idx].contour(\n",
    "                x_mesh[cont], y_mesh[cont], value,\n",
    "                linewidths=1.,\n",
    "                linestyles='solid',\n",
    "                levels=3,\n",
    "                colors = list(mcolors.BASE_COLORS)[model_idx]\n",
    "            )\n",
    "            ax[idx].set(xlabel=\"$w_0$\",\n",
    "                        ylabel=\"$w_1$\",\n",
    "                        title=f\"{cont} nations\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RrwOJMCEq3Sp",
    "outputId": "638aaec6-becb-4bd2-da72-a95d131b6ea4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "density_plt([model_bayes, model_mf])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j4hedNzPq3Sp"
   },
   "source": [
    "### 3.1.2 Visualizing the marginals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "X2TqGHy7q3Sp",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def marginal_plot(model_bayes: dict, model_mf: dict, num_samples: int = 10000):\n",
    "    \"\"\"\n",
    "    Make a plot of the marginal density for the two models.\n",
    "    \"\"\"\n",
    "    weight_samples_bayes = {}\n",
    "    weight_samples_vi = {}\n",
    "\n",
    "    for cont in model_bayes.keys():\n",
    "        # Results for Bayes distribution\n",
    "        weight_samples_bayes[cont] = np.random.multivariate_normal(mean=model_bayes[cont]['w'],\n",
    "                                                                   cov=model_bayes[cont]['cov'], \n",
    "                                                                   size=num_samples)\n",
    "        # Results for mean field\n",
    "        weight_samples_vi[cont] = np.random.multivariate_normal(mean=model_mf[cont]['w'], \n",
    "                                                                cov=model_mf[cont]['cov'], \n",
    "                                                                size=num_samples)\n",
    "\n",
    "    fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(12, 6))\n",
    "    fig.suptitle(\"Marginal density for bias and slope\", fontsize=16)\n",
    "    fig.tight_layout(pad=2.0)\n",
    "    for idx, cont in enumerate(x_data):\n",
    "        ax[0, idx].set_title(f\"{cont} nations (bias)\")\n",
    "        sns.kdeplot(weight_samples_vi[cont][:, 0], ax=ax[0, idx],  label=\"Variational inference - bias\")\n",
    "        sns.kdeplot(weight_samples_bayes[cont][:, 0], ax=ax[0, idx], label=\"Bayes - bias\")\n",
    "        ax[1, idx].set_title(f\"{cont} nations (slope)\")\n",
    "        sns.kdeplot(weight_samples_vi[cont][:, 1], ax=ax[1, idx], label=\"Variational inference - slope\")\n",
    "        sns.kdeplot(weight_samples_bayes[cont][:, 1], ax=ax[1, idx],  label=\"Bayes - slope\")\n",
    "\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "    for cont in model_bayes.keys():\n",
    "        print(f\"Mean field std for bias ({cont}): {np.std(weight_samples_vi[cont][:, 0])}\")\n",
    "        print(f\"Mean field std for slope ({cont}): {np.std(weight_samples_vi[cont][:, 1])}\")\n",
    "        print(f\"Bayes std for bias ({cont}): {np.std(weight_samples_bayes[cont][:, 0])}\")\n",
    "        print(f\"Bayes std for slope ({cont}): {np.std(weight_samples_bayes[cont][:, 1])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-tWsyM2Iq3Sp",
    "outputId": "6acd3a27-9ca6-49b2-f9e7-4703fefd2375",
    "tags": []
   },
   "outputs": [],
   "source": [
    "marginal_plot(model_bayes, model_mf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j0qNSsBSq3Sp"
   },
   "source": [
    "## 3.2 Working with the posterior predictive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CAVTnEOuq3Sp"
   },
   "source": [
    "### 3.2.1 Get the posterior predictive \n",
    "Given the posterior distribution\n",
    "$$\n",
    "p({\\bf w}| {\\bf X}, {\\bf y}, 1/\\theta) =  {\\mathcal N}({\\bf w}| {\\bf w}_p, \\Sigma_p),\n",
    "$$\n",
    "\n",
    "where ${\\bf w}_p$ and $\\Sigma_p$ is the (approximate) posterior mean and (approximate) posterior covariance, respectively, we can calculate the posterior predictive distribution for a data point ${\\bf x}$:\n",
    "$$\n",
    "p(y | {\\bf x}, {\\mathcal D}, 1/\\theta) = \\int \\mathcal N(y|{\\bf x}^T{\\bf w}, 1/\\theta) {\\mathcal N}({\\bf w}| {\\bf w}_p, \\Sigma_p),\n",
    "$$\n",
    "\n",
    "\n",
    "See, e.g., Murphy (2022), *Probabilistic machine learning: an introduction*. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a_P51WTLq3Sp",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_posterior_predictive(model: dict, x_data_: np.array, theta: float = 1):\n",
    "    \"\"\"\n",
    "    Calculate the posterior predictive distribution for the model.\n",
    "    :param model: The model used for predictions. \n",
    "    :param x_data_: The input data for which the predictions should be made.\n",
    "    :param theta: The prior precision for y.\n",
    "    :return: A dictionary with the predictions containing also the variance terms.\n",
    "    \"\"\"\n",
    "    predictions = {}\n",
    "\n",
    "    for cont in model.keys():\n",
    "        predictions[cont] = {}\n",
    "        x_aug = np.c_[np.ones(x_data_[cont].shape[0]), x_data_[cont]]\n",
    "        predictions[cont]['pred'] = model[cont]['w'] @ x_aug.T\n",
    "        predictions[cont]['var'] = 1./theta + np.diagonal(x_aug @ model[cont]['cov'] @ x_aug.T)\n",
    "\n",
    "    return predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GDMJmRapq3Sp",
    "tags": []
   },
   "outputs": [],
   "source": [
    "preds = get_posterior_predictive(model_bayes, x_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8KN61I6rq3Sp"
   },
   "source": [
    "### 3.2.2 Visualize the predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IDgKDtxhq3Sp",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def predictive_plot(ax, x_data_: dict, predictions: dict, label: str = None):\n",
    "    \"\"\"\n",
    "    Plot the posterior predictive distribution.\n",
    "    \"\"\"\n",
    "    for idx, cont in enumerate(x_data_.keys()):\n",
    "        sort_idx = np.argsort(x_data_[cont].flatten())\n",
    "        sorted_input = x_data_[cont][sort_idx]\n",
    "        preds = predictions[cont]['pred'][sort_idx]\n",
    "        var = predictions[cont]['var'][sort_idx]\n",
    "        ax[idx].fill_between(sorted_input.flatten(), preds - 2*np.sqrt(var),\n",
    "                             preds + 2*np.sqrt(var), color='b', alpha=0.1)\n",
    "        ax[idx].plot(sorted_input, preds, color='red', linewidth=2, label=label)\n",
    "    return ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sSdau7rmq3Sp",
    "outputId": "93056903-9396-4561-c7c5-d21c6d64adcf",
    "tags": []
   },
   "outputs": [],
   "source": [
    "ax = prepare_figure('Posterior predictive')\n",
    "predictive_plot(ax, x_data, preds)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WomacVYmq3Sq"
   },
   "source": [
    "### 3.2.2 Calculate log-likelihoods"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sSuB0aN_q3Sq"
   },
   "source": [
    "We can evaluate the predictive qualities of our model using the log-likelihood:\n",
    "$$\n",
    "\\log p({\\bf y}|{\\bf X},\\text{posterior predictive mean and variance}) = \\sum_{i=1}^N \\log p(y_i|{\\bf x}_i,\\text{posterior predictive mean and variance})\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "w6lHibowq3Sq",
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_ll(x_data_,y_data_: dict, model: dict):\n",
    "\n",
    "    preds = get_posterior_predictive(model, x_data_)\n",
    "    ll = {}\n",
    "    for cont in model.keys():\n",
    "        ll[cont] = 1./y_data_[cont].shape[0]*np.sum(norm.logpdf(y_data_[cont], loc=preds[cont]['pred'], scale=np.sqrt(preds[cont]['var'])))\n",
    "\n",
    "    return ll"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wJ61iGm7q3Sq",
    "outputId": "ef33069f-6eea-44d3-83f9-1b1cb5379d75",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get the log-likelihood over the training set.\n",
    "ll = get_ll(x_data, y_data, model_bayes)\n",
    "print(ll)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9BmfshPdq3Sq"
   },
   "source": [
    "# 4. Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Td0F9QwSq3Sq"
   },
   "source": [
    "In these exercises you should first familiarize yourself with the notebook, in particular try to connect the variational implementation to the slides previously covered.\n",
    "\n",
    "Afterwards, explore and compare the models above by, e.g., varying the model hyperparameters and the data set sizes being used. Demonstrations are given below.\n",
    "\n",
    "**Optionally:** Derive the updating rule for $q(w)$ and/or $q(b)$ "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l7IDsQbAq3Sq"
   },
   "source": [
    "# 5. Demonstrations/Solutions\n",
    "\n",
    "The demonstrations below should only be considered inspiration for your own investigations into the properties of the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l1LPH85Wq3Sq"
   },
   "source": [
    "### Exploring the effect of changing the priors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "rCOD80dNq3Sq",
    "outputId": "19400973-0f6a-4e49-dad1-80002dd77d6c",
    "tags": []
   },
   "outputs": [],
   "source": [
    "theta = 1\n",
    "_model_ = bayes_solution(x_data, y_data, gamma_w=1, gamma_b=1, theta=theta, verbose=True)\n",
    "ax = prepare_figure(\"Sampling parameters from the posterior distribution\")\n",
    "plot_model(ax, model=_model_, num_samples=100)\n",
    "preds = get_posterior_predictive(_model_, x_data, theta=theta)\n",
    "ax = prepare_figure('Posterior predictive')\n",
    "predictive_plot(ax, x_data, preds)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZtM5l4lCq3Sq"
   },
   "source": [
    "### Experimenting with the priors and the amount of data used for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "66ZkuH9Uq3Sq",
    "outputId": "c7fe4700-122f-45bc-d04b-2db5994b899d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "x_train, y_train, x_test, y_test = get_data_subsets(x_data, y_data, 0.7, verbose=False)\n",
    "x_train, y_train, _, _ = get_data_subsets(x_train, y_train, 0.1, verbose=False)\n",
    "\n",
    "for cont in x_train.keys():\n",
    "    print(f\"{cont} entries: {x_train[cont].shape[0]}\")\n",
    "for cont in x_test.keys():\n",
    "    print(f\"{cont} entries: {x_test[cont].shape[0]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7Nave6lMq3Sq",
    "outputId": "96f89361-2668-46f9-d56d-e993c311f1da",
    "tags": []
   },
   "outputs": [],
   "source": [
    "theta = 1\n",
    "_model_vi = variational_solution(x_train, y_train, gamma_w=1, gamma_b=1, theta=theta, verbose=False)\n",
    "ax = prepare_figure(\"Sampling parameters from the mean field variational distribution\", x_train, y_train)\n",
    "plot_model(ax, model=_model_vi, num_samples=100)\n",
    "preds = get_posterior_predictive(_model_vi, x_train, theta=theta)\n",
    "ax = prepare_figure('Posterior predictive', x_train, y_train)\n",
    "predictive_plot(ax, x_train, preds)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qEA3uYvaq3Sq"
   },
   "source": [
    "### Varying the amount of data used for training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zKqp-qFHq3Sq",
    "outputId": "824994f8-939e-432a-85ca-bf83c1362e6c",
    "tags": []
   },
   "outputs": [],
   "source": [
    "x_train, y_train, x_test, y_test = get_data_subsets(x_data, y_data, 0.7, verbose=False)\n",
    "\n",
    "# We keep track of the log-likelihood calculations for future plotting\n",
    "ll = {}\n",
    "for cont in x_train.keys():\n",
    "    ll[cont] = []\n",
    "\n",
    "theta = 1     \n",
    "#Data fraction to consider\n",
    "data_fractions = np.linspace(0.05,1,5)\n",
    "for f in data_fractions:\n",
    "    _x_train, _y_train, _, _ = get_data_subsets(x_train, y_train, f, verbose=False)\n",
    "    _model = variational_solution(_x_train, _y_train, gamma_w=1, gamma_b=1, theta=theta, verbose=False)\n",
    "    _model_lr = least_squares_solution(_x_train, _y_train, verbose=False)\n",
    "    for cont in x_train.keys():\n",
    "        ll[cont].append(get_ll(x_test, y_test, _model)[cont])\n",
    "    ax = prepare_figure('Model comparison', _x_train, _y_train)\n",
    "    preds = get_posterior_predictive(_model, x_train, theta=theta)\n",
    "    ax = predictive_plot(ax, x_train, preds, label='Posterior predictive')    \n",
    "    plot_model(ax, model=_model_lr, line_color='green', label='OLS')\n",
    "    plt.show()       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ucymgJXaq3Sq",
    "outputId": "9687d3ca-187d-481e-d28b-b2ebb19b6d3f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot log-likelihood calculations\n",
    "plt.plot(data_fractions, ll['african'])\n",
    "plt.title('Log-likelihood vs. fraction of full training set used for learning')\n",
    "plt.xlabel('Fraction of full training set')\n",
    "plt.ylabel('Log-likelihood')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OHifWwDiq3Sq"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "name": "CAVI-linreg.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}