{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Jt5iY6Lylfs3"
   },
   "source": [
    "# Alternative More Detailed Solution for Lab01 - Task C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:28.252327Z",
     "iopub.status.busy": "2022-10-20T20:24:28.251888Z",
     "iopub.status.idle": "2022-10-20T20:24:29.759964Z",
     "shell.execute_reply": "2022-10-20T20:24:29.759266Z"
    }
   },
   "outputs": [],
   "source": [
    "# Import necessary libraries\n",
    "%matplotlib inline\n",
    "import numpy as np\n",
    "from numpy.random import rand, randn\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# Data generation\n",
    "n, d, k = 100, 2, 2\n",
    "np.random.seed(20)\n",
    "X = rand(n, d)\n",
    "means = [rand(d) * 0.5 + 0.5, -rand(d) * 0.5 + 0.5]  # for better plotting when k = 2\n",
    "S = np.diag(rand(d))\n",
    "sigmas = [S] * k"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kG4qTMF7lfs-"
   },
   "source": [
    "## First try - Computing the probability using for loops\n",
    "\n",
    "The following was not clear in the instructions: The function you had to implement was `compute_log_p`, so you should have transformed the expression by applying a log first.\n",
    "\n",
    "It wouldn't have changed the result of the graph though.\n",
    "\n",
    "If you skipped this step and completed this exercise using for loops, you probably did something like the next cell.\n",
    "\n",
    "---\n",
    "\n",
    "The function to implement:\n",
    "$$p(x | \\mu, \\Sigma) = \\frac{1}{(2\\pi)^{d/2} |\\Sigma|^{1/2}} \\exp\\left(-\\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu)\\right)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:29.762904Z",
     "iopub.status.busy": "2022-10-20T20:24:29.762690Z",
     "iopub.status.idle": "2022-10-20T20:24:29.781440Z",
     "shell.execute_reply": "2022-10-20T20:24:29.780811Z"
    }
   },
   "outputs": [],
   "source": [
    "def compute_p_forloop(X, mean, sigma):\n",
    "    [n, d] = np.shape(X)\n",
    "\n",
    "    # The constant (same for all samples) term\n",
    "    c = 2 * np.power(np.pi, d / 2) * np.power(np.linalg.det(sigma), 0.5)\n",
    "    invSigma = np.linalg.inv(sigma)\n",
    "\n",
    "    result = np.zeros((n,))\n",
    "    for i in range(n):\n",
    "        xmu = X[i] - mean\n",
    "        result[i] = 1 / c * np.exp(-0.5 * (xmu).T.dot(invSigma).dot(xmu))\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pjN6SUvelftB"
   },
   "source": [
    "And it works pretty well. Comparing with our solution, the difference is very small, around $10^{-15}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:29.784533Z",
     "iopub.status.busy": "2022-10-20T20:24:29.784220Z",
     "iopub.status.idle": "2022-10-20T20:24:30.030174Z",
     "shell.execute_reply": "2022-10-20T20:24:30.029615Z"
    }
   },
   "outputs": [],
   "source": [
    "### -----\n",
    "### Applying log to our function + Solution\n",
    "\n",
    "\n",
    "def compute_log_p_forloop(X, mean, sigma):\n",
    "    return np.log(compute_p_forloop(X, mean, sigma))\n",
    "\n",
    "\n",
    "def compute_log_p_solution(X, mean, sigma):\n",
    "    dxm = X - mean\n",
    "    return (\n",
    "        -0.5 * np.sum(dxm * np.dot(dxm, np.linalg.inv(sigma)), axis=1)\n",
    "        - np.log(2 * np.pi) * (d / 2)\n",
    "        - 0.5 * np.log(np.linalg.det(sigma))\n",
    "    )\n",
    "\n",
    "\n",
    "### -----\n",
    "### Difference between solution and this implementation\n",
    "\n",
    "a = compute_log_p_forloop(X, means[0], sigmas[0])\n",
    "b = compute_log_p_solution(X, means[0], sigmas[0])\n",
    "\n",
    "print(\"|a-b|_2 =\", np.linalg.norm(a - b))\n",
    "\n",
    "### -----\n",
    "### Print the graphs\n",
    "\n",
    "\n",
    "def makeGraph(function, X, means, sigmas):\n",
    "    log_ps = [function(X, m, s) for m, s in zip(means, sigmas)]\n",
    "\n",
    "    assignments = np.argmax(log_ps, axis=0)\n",
    "\n",
    "    colors = np.array([\"red\", \"green\"])[assignments]\n",
    "    plt.title(function.__name__)\n",
    "    plt.scatter(X[:, 0], X[:, 1], c=colors, s=100)\n",
    "    plt.scatter(np.array(means)[:, 0], np.array(means)[:, 1], marker=\"*\", s=200)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "makeGraph(compute_log_p_forloop, X, means, sigmas)\n",
    "makeGraph(compute_log_p_solution, X, means, sigmas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IwM_E99OlftD"
   },
   "source": [
    "## Computing the log-probability directly\n",
    "\n",
    "**Why log?** - Our goal is to compare probabilities to see to which of the two stars a point belongs.\n",
    "\n",
    "But the formula for the probability is a bit heavy, with multiplications and exponents.\n",
    "\n",
    "By applying a log transform, we get additions and multiplications, which is easer to handle, and does not impact the comparison - if `a > b`, `log(a) > log(b)`.\n",
    "\n",
    "(If it does not make sense - don't worry - you'll see this in the coming lectures)\n",
    "\n",
    "### Simplifying the equation\n",
    "\n",
    "Notation: $x$ is a sample, so $[D \\times 1]$, $X$ is the matrix, so $[N \\times D]$\n",
    "\n",
    "$$p(x | \\mu, \\Sigma) = \\frac{1}{(2\\pi)^{d/2} |\\Sigma|^{1/2}} \\exp\\left(-\\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu)\\right)$$\n",
    "\n",
    "$$\\log p(x | \\mu, \\Sigma) = \\log \\left[ \\frac{1}{(2\\pi)^{d/2} |\\Sigma|^{1/2}} \\exp\\left(-\\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu)\\right) \\right]$$\n",
    "\n",
    "$$\\log p(x | \\mu, \\Sigma) = \\log \\left[ \\frac{1}{(2\\pi)^{d/2} |\\Sigma|^{1/2}} \\right] + \\log \\left[\\exp\\left(-\\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu)\\right) \\right]$$\n",
    "\n",
    "$$\\log p(x | \\mu, \\Sigma) = -\\log \\left[ (2\\pi)^{d/2} |\\Sigma|^{1/2} \\right] - \\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu)$$\n",
    "\n",
    "This gives us the following expression,\n",
    "$$\\log p(x | \\mu, \\Sigma) = - \\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu) + c(\\mu,\\Sigma),$$\n",
    "Where $c(\\mu, \\Sigma) = -\\log \\left[ (2\\pi)^{d/2} |\\Sigma|^{1/2} \\right] = - \\frac{1}{2}\\left[d\\log(2 \\pi) + \\log(|\\Sigma|)\\right]$\n",
    "\n",
    "Steps used:\n",
    "- $\\log(ab) = \\log(a) + \\log(b)$\n",
    "- $\\log(1/a) = -\\log(a)$\n",
    "- $\\log(\\exp(a)) = a$\n",
    "\n",
    "### Implementing this function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:30.033322Z",
     "iopub.status.busy": "2022-10-20T20:24:30.032911Z",
     "iopub.status.idle": "2022-10-20T20:24:30.211629Z",
     "shell.execute_reply": "2022-10-20T20:24:30.210758Z"
    }
   },
   "outputs": [],
   "source": [
    "def compute_log_p_forloop(X, mean, sigma):\n",
    "    [n, d] = np.shape(X)\n",
    "    result = np.zeros((n,))\n",
    "\n",
    "    constant = -0.5 * (d * np.log(2 * np.pi) + np.log(np.linalg.det(sigma)))\n",
    "    invSigma = np.linalg.inv(sigma)\n",
    "\n",
    "    for i in range(n):\n",
    "        xmu = X[i] - mean\n",
    "        result[i] = -(1 / 2) * (xmu).T.dot(invSigma).dot(xmu) + constant\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "### -----\n",
    "### Difference between solution and this implementation\n",
    "\n",
    "a = compute_log_p_forloop(X, means[0], sigmas[0])\n",
    "b = compute_log_p_solution(X, means[0], sigmas[0])\n",
    "print(\"|a-b|_2 =\", np.linalg.norm(a - b))\n",
    "\n",
    "### -----\n",
    "### Print the graphs\n",
    "\n",
    "makeGraph(compute_log_p_forloop, X, means, sigmas)\n",
    "makeGraph(compute_log_p_solution, X, means, sigmas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mhMdtdKDlftF"
   },
   "source": [
    "### Matrix formulation - using no for loops\n",
    "\n",
    "Now, how can we use numpy to avoid using for loops?\n",
    "\n",
    "This is firstly an algebra problem, not a programming one - the programming is just translation. Our main tools are\n",
    "- Traditional matrix multiplication\n",
    "- Addition and multiplication by scalars\n",
    "- Elementwise addition and multiplication by matrices\n",
    "- Summation and product over rows or columns\n",
    "- ...\n",
    "\n",
    "But first, we need to know what we are after. The expression for the log is\n",
    "$$\\log p(x | \\mu, \\Sigma) = - \\frac{1}{2}(x - \\mu)^T\\Sigma^{-1}(x-\\mu) + c(\\mu,\\Sigma),$$\n",
    "Where $c(\\mu, \\Sigma) = -\\log \\left[ (2\\pi)^{d/2} |\\Sigma|^{1/2} \\right] = - \\frac{1}{2}\\left[d\\log(2 \\pi) + \\log(|\\Sigma|)\\right]$\n",
    "\n",
    "And our function `compute_log_p` should return a vector of `N` elements, which is\n",
    "\n",
    "$$\n",
    "\\texttt{compute\\_log\\_p}(X, \\mu, \\Sigma) = \n",
    "\\begin{pmatrix}\n",
    "\\log p(x_1 | \\mu, \\Sigma)\\\\\n",
    "...\\\\\n",
    "\\log p(x_n | \\mu, \\Sigma)\\\\\n",
    "\\end{pmatrix}\n",
    "=\n",
    "-\\frac{1}{2}\n",
    "\\begin{pmatrix}\n",
    "(x_1 - \\mu)^T\\Sigma^{-1}(x_1-\\mu)\\\\\n",
    "...\\\\\n",
    "(x_N - \\mu)^T\\Sigma^{-1}(x_N-\\mu)\\\\\n",
    "\\end{pmatrix}\n",
    "+c(\\mu, \\Sigma)\n",
    "$$\n",
    "\n",
    "Let us focus on the matrix part of the formula, which we'll call $M$, after some simplification:\n",
    "- substitue $A = X - \\mu$ (a is a `[N x D]` matrix).\n",
    "- $A_i$ is the first row, a `D`-elements vector\n",
    "- $A_{ij}$ is the element at cell $i,j$. \n",
    "\n",
    "\n",
    "We have\n",
    "\n",
    "$$\n",
    "M=\n",
    "\\begin{pmatrix}\n",
    "A_1^T\\Sigma^{-1}A_1\\\\\n",
    "...\\\\\n",
    "A_N^T\\Sigma^{-1}A_N\\\\\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Can we simplify the expression? For a single row, we have\n",
    "\n",
    "\n",
    "$$M_i = A_i^T \\Sigma^{-1} A_i = \n",
    "\\begin{pmatrix}\n",
    "A_{i1} & ... & A_{iD}\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "    \\Sigma_{11} & ... & \\Sigma_{1D} \\\\\n",
    "    \\vdots      & \\ddots & \\vdots   \\\\\n",
    "    \\Sigma_{D1} & ... & \\Sigma_{DD} \\\\\n",
    "\\end{pmatrix}^{-1}\n",
    "\\begin{pmatrix}\n",
    "A_{i1} \\\\ ... \\\\ A_{iD}\n",
    "\\end{pmatrix}\n",
    "\\text{  -  Dimensions:  } [1 x D] [D x D] [D x 1] \\text{  }\n",
    "$$\n",
    "\n",
    "From here, there are two path - simplification using the properties of the matrices we have, and a more brute force approach.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "l_bG4nEPlftG"
   },
   "source": [
    "#### Simplifying using the properties of the matrices\n",
    "\n",
    "The thing to note is that $\\Sigma$ is a diagonal matrix (Data loading code: `S = np.diag(rand(d))`). Therefore, we have \n",
    "$$M_i = A_i^T \\Sigma^{-1} A_i = \n",
    "\\begin{pmatrix}\n",
    "A_{i1} & ... & A_{iD}\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "    1/\\Sigma_{11} & 0      & ...    & 0 \\\\\n",
    "    0           & \\ddots & \\ddots & \\vdots   \\\\\n",
    "    \\vdots      & \\ddots & \\ddots & \\vdots   \\\\\n",
    "    0           & ...    & ...    & 1/\\Sigma_{DD} \\\\\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "A_{i1} \\\\ ... \\\\ A_{iD}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "$$\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "A_{i1} & ... & A_{iD}\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "A_{i1}/\\Sigma_{11} \\\\ ... \\\\ A_{iD}/\\Sigma_{DD} \n",
    "\\end{pmatrix}\n",
    "=\n",
    "\\sum_{j=1}^D A_{ij}^2 /\\Sigma_{jj}\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "A_{i1}^2 & ... & A_{iD}^2\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "1/\\Sigma_{11} \\\\ ... \\\\ 1/\\Sigma_{DD} \n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Notice that on those last formulations, we have a `[1 x D] [D x 1]` system. The `[D x 1]` matrix is a transformation we apply to the `[1 x D]` input, and we can apply it to all samples by providing a `[N x D]` input, as follow.\n",
    "$$\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "A_{11}^2 & ...    & A_{1D}^2 \\\\\n",
    "\\vdots   & \\ddots & \\vdots   \\\\\n",
    "A_{N1}^2 & ...    & A_{ND}^2 \\\\\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "1/\\Sigma_{11} \\\\ ... \\\\ 1/\\Sigma_{DD} \n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Or, in (pseudo) code,\n",
    "\n",
    "    A = X - mu\n",
    "    A2 = A * A # element-wise multiplication\n",
    "    invSigma = np.linalg.inverse(Sigma)\n",
    "    diagInvSigma = np.diag(invSigma)\n",
    "    M = A2.dot(diagInvSigma)\n",
    "    \n",
    "    compute_log_p(X, mu, sigma) = - 0.5 * M + c(mu, sigma)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:30.214187Z",
     "iopub.status.busy": "2022-10-20T20:24:30.213964Z",
     "iopub.status.idle": "2022-10-20T20:24:30.425844Z",
     "shell.execute_reply": "2022-10-20T20:24:30.424857Z"
    }
   },
   "outputs": [],
   "source": [
    "def compute_log_p_sol1(X, mean, sigma):\n",
    "    [n, d] = np.shape(X)\n",
    "\n",
    "    constant = -0.5 * (d * np.log(2 * np.pi) + np.log(np.linalg.det(sigma)))\n",
    "\n",
    "    diagInvSigma = np.diag(np.linalg.inv(sigma))\n",
    "    xmu = X - mean\n",
    "    xmu2 = xmu * xmu\n",
    "    return -0.5 * xmu2.dot(diagInvSigma) + constant\n",
    "\n",
    "    # return result\n",
    "\n",
    "\n",
    "### -----\n",
    "### Difference between solution and this implementation\n",
    "\n",
    "a = compute_log_p_sol1(X, means[0], sigmas[0])\n",
    "b = compute_log_p_solution(X, means[0], sigmas[0])\n",
    "print(\"|a-b|_2 =\", np.linalg.norm(a - b))\n",
    "\n",
    "### -----\n",
    "### Print the graphs\n",
    "\n",
    "makeGraph(compute_log_p_sol1, X, means, sigmas)\n",
    "makeGraph(compute_log_p_solution, X, means, sigmas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SZj3FdhIlftH"
   },
   "source": [
    "#### More bruteforce approach\n",
    "\n",
    "If your $\\Sigma$ matrix is less nice, you might need to take some more steps. Again, looking at a single sample, we have\n",
    "\n",
    "$$M_i = A_i^T \\Sigma^{-1} A_i = \n",
    "\\begin{pmatrix}\n",
    "A_{i1} & ... & A_{iD}\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "    \\Sigma_{11} & ... & \\Sigma_{1D} \\\\\n",
    "    \\vdots      & \\ddots & \\vdots   \\\\\n",
    "    \\Sigma_{D1} & ... & \\Sigma_{DD} \\\\\n",
    "\\end{pmatrix}^{-1}\n",
    "\\begin{pmatrix}\n",
    "A_{i1} \\\\ ... \\\\ A_{iD}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Using $\\Sigma' = \\Sigma^{-1}$ to avoid confusion between the matrix inverse and the elementwise inverse, this gives\n",
    "\n",
    "$$M_n = \n",
    "\\begin{pmatrix}\n",
    "A_{n1} & ... & A_{nD}\n",
    "\\end{pmatrix}\n",
    "\\begin{pmatrix}\n",
    "\\sum_{i=1}^D a_{ni} \\Sigma'_{1i} \\\\\n",
    "\\vdots \\\\\n",
    "\\sum_{i=1}^D a_{ni} \\Sigma'_{Di} \\\\\n",
    "\\end{pmatrix}\n",
    "= \\sum_{i=1}^D \\sum_{j=1}^D a_{ni} a_{nj} \\Sigma'_{ij}\n",
    "$$\n",
    "\n",
    "In last section we found a nice way to split the one sample the expression to a `[1 x D][D x 1]` matrix multiplication, with the inputs on the left `[1 x D]` matrix and the transformation being the `[D x 1]` matrix; extanding to a `[N x D]` input is easy. Here, we are not assuming that $\\Sigma$ is diagonal, which complicates a little bit the system. We will need to bring two tools out of the box:\n",
    "\n",
    "* Column summation - transforms a `[A x B]` matrix into a `[A x 1]` matrix by summing all columns, for each row, as follow:\n",
    "$$\n",
    "\\texttt{column summation of }\n",
    "\\begin{pmatrix}\n",
    "y_{11} & y_{12} & ... & y_{1B} \\\\\n",
    "y_{21} & y_{22} & ... & y_{2B}\n",
    "\\end{pmatrix}\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "\\sum_{i=1}^B y_{1i} \\\\\n",
    "\\sum_{i=1}^B y_{2i}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "* Element-wise matrix multiplication - What the `*` operator does on numpy matrices - written $\\odot$ here, \n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "y_{1} & y_{2} & ... & y_{B} \n",
    "\\end{pmatrix}\n",
    "\\odot\n",
    "\\begin{pmatrix}\n",
    "z_{1} & z_{2} & ... & z_{B} \n",
    "\\end{pmatrix}\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "y_{1}z_{1} & y_{2}z_{2} & ... & y_{B}z_{B}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Using those tools, we can work out a better representation for our formula. We first expand one of the summation by doing a reverse-column-summation to go from our scalar result to a `[D x 1]` matrix.\n",
    "$$\n",
    "\\sum_{i=1}^D \\sum_{j=1}^D a_{ni} a_{nj} \\Sigma'_{ij}\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "A_{n1} \\sum_{j=1}^D A_{nj} \\Sigma'_{1j} & \n",
    "A_{n2} \\sum_{j=1}^D A_{nj} \\Sigma'_{2j} & \n",
    "...\n",
    "A_{nD} \\sum_{j=1}^D A_{nj} \\Sigma'_{Dj}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "We now can separate the matrix into two parts using element-wise multiplication.\n",
    "$$\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "A_{n1} & \n",
    "A_{n2} & \n",
    "...\n",
    "A_{nD}\n",
    "\\end{pmatrix}\n",
    "\\odot\n",
    "\\begin{pmatrix}\n",
    "\\sum_{j=1}^D A_{nj} \\Sigma'_{1j} & \n",
    "\\sum_{j=1}^D A_{nj} \\Sigma'_{2j} & \n",
    "...\n",
    "\\sum_{j=1}^D A_{nj} \\Sigma'_{Dj}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "$$\n",
    "=\n",
    "A_n\n",
    "\\odot\n",
    "\\begin{pmatrix}\n",
    "\\sum_{j=1}^D A_{nj} \\Sigma'_{1j} & \n",
    "\\sum_{j=1}^D A_{nj} \\Sigma'_{2j} & \n",
    "...\n",
    "\\sum_{j=1}^D A_{nj} \\Sigma'_{Dj}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "Now, we can see that the `[1 x D]` matrix on the right is the result of $A_n\\Sigma'$ (`[1 x D][D x D]`).\n",
    "$$\n",
    "=\n",
    "A_n \\odot (A_n \\Sigma')\n",
    "$$\n",
    "\n",
    "\n",
    "And this is it. If instead of $A_n$, we plug $A$ into the system, we get a `[N x D]` $\\odot$ `[N x D][D x D]` system, the solution for every sample in one line,\n",
    "$$A \\odot (A \\Sigma')$$\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:30.429844Z",
     "iopub.status.busy": "2022-10-20T20:24:30.429616Z",
     "iopub.status.idle": "2022-10-20T20:24:30.542573Z",
     "shell.execute_reply": "2022-10-20T20:24:30.541954Z"
    }
   },
   "outputs": [],
   "source": [
    "def compute_log_p_solution(X, mean, sigma):\n",
    "    d = X.shape[1]\n",
    "    c = -np.log(2 * np.pi) * (d / 2) - 0.5 * np.log(np.linalg.det(sigma))\n",
    "    A = X - mean\n",
    "    invSigma = np.linalg.inv(sigma)\n",
    "\n",
    "    return -0.5 * np.sum(A * (A.dot(invSigma)), axis=1) + c\n",
    "\n",
    "\n",
    "makeGraph(compute_log_p_solution, X, means, sigmas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Kx6AxuXglftJ"
   },
   "source": [
    "### Why do we even care?\n",
    "\n",
    "Time.\n",
    "\n",
    "For loops are much more expensive than lingear algebra - for which we have specialized libraries, Code that correctly uses linear algebra can run 10x to 100x faster than a for-loop program for the same output.\n",
    "\n",
    "Here is a comparison of the different solutions,\n",
    "* Using for loops\n",
    "* Using numpy to do the heavy lifting\n",
    "* Using the fact that $\\Sigma$ is a diagonal\n",
    "\n",
    "Most of the benefit comes from using numpy correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:30.546140Z",
     "iopub.status.busy": "2022-10-20T20:24:30.545742Z",
     "iopub.status.idle": "2022-10-20T20:24:58.000795Z",
     "shell.execute_reply": "2022-10-20T20:24:58.000105Z"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "def generateData(n, d, k):\n",
    "    X = rand(n, d)\n",
    "    means = [\n",
    "        rand(d) * 0.5 + 0.5,\n",
    "        -rand(d) * 0.5 + 0.5,\n",
    "    ]  # for better plotting when k = 2\n",
    "    S = np.diag(rand(d))\n",
    "    sigmas = [S] * k\n",
    "\n",
    "    return X, means, sigmas\n",
    "\n",
    "\n",
    "matrix_time = np.zeros(\n",
    "    10,\n",
    ")\n",
    "forloop_time = np.zeros(\n",
    "    10,\n",
    ")\n",
    "i = 0\n",
    "Ns = np.logspace(0, 7, num=10)\n",
    "for N in Ns:\n",
    "    X_n, means_n, sigmas_n = generateData(int(np.floor(N)), 2, 2)\n",
    "\n",
    "    start_time = time.time()\n",
    "    compute_log_p_solution(X_n, means_n[0], sigmas_n[0])\n",
    "    matrix_time[i] = time.time() - start_time\n",
    "\n",
    "    start_time = time.time()\n",
    "    compute_log_p_forloop(X_n, means_n[0], sigmas_n[0])\n",
    "    forloop_time[i] = time.time() - start_time\n",
    "\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:58.003785Z",
     "iopub.status.busy": "2022-10-20T20:24:58.003558Z",
     "iopub.status.idle": "2022-10-20T20:24:58.778324Z",
     "shell.execute_reply": "2022-10-20T20:24:58.777843Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.title(\"Computation time comparison\")\n",
    "h1 = plt.plot(np.floor(Ns), forloop_time, label=\"For loop\")\n",
    "h2 = plt.plot(np.floor(Ns), matrix_time, label=\"Matrix\")\n",
    "plt.xscale(\"log\")\n",
    "plt.yscale(\"log\")\n",
    "plt.ylabel(\"Time (s)\")\n",
    "plt.legend([\"For loop\", \"Matrix\"])\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:58.781079Z",
     "iopub.status.busy": "2022-10-20T20:24:58.780887Z",
     "iopub.status.idle": "2022-10-20T20:24:59.797938Z",
     "shell.execute_reply": "2022-10-20T20:24:59.797271Z"
    }
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import timeit\n",
    "\n",
    "\n",
    "def generateData(n, d, k):\n",
    "    X = rand(n, d)\n",
    "    means = [\n",
    "        rand(d) * 0.5 + 0.5,\n",
    "        -rand(d) * 0.5 + 0.5,\n",
    "    ]  # for better plotting when k = 2\n",
    "    S = np.diag(rand(d))\n",
    "    sigmas = [S] * k\n",
    "\n",
    "    return X, means, sigmas\n",
    "\n",
    "\n",
    "matrix_time = np.zeros(\n",
    "    10,\n",
    ")\n",
    "diag_time = np.zeros(\n",
    "    10,\n",
    ")\n",
    "i = 0\n",
    "Ns = np.logspace(0, 7, num=10)\n",
    "for N in Ns:\n",
    "    X_n, means_n, sigmas_n = generateData(int(np.floor(N)), 2, 2)\n",
    "\n",
    "    start_time = time.time()\n",
    "    compute_log_p_solution(X_n, means_n[0], sigmas_n[0])\n",
    "    matrix_time[i] = time.time() - start_time\n",
    "\n",
    "    start_time = time.time()\n",
    "    compute_log_p_sol1(X_n, means_n[0], sigmas_n[0])\n",
    "    diag_time[i] = time.time() - start_time\n",
    "\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-10-20T20:24:59.801428Z",
     "iopub.status.busy": "2022-10-20T20:24:59.801244Z",
     "iopub.status.idle": "2022-10-20T20:25:00.448928Z",
     "shell.execute_reply": "2022-10-20T20:25:00.448381Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.title(\"Computation time comparison\\n Assuming diagonal matrix vs. no assumption\")\n",
    "h1 = plt.plot(np.floor(Ns), diag_time, label=\"Assuming diagonal covariance\")\n",
    "h2 = plt.plot(np.floor(Ns), matrix_time, label=\"Matrix\")\n",
    "plt.xscale(\"log\")\n",
    "plt.yscale(\"log\")\n",
    "plt.ylabel(\"Time (s)\")\n",
    "plt.legend([\"Assuming diagonal covariance\", \"No assumption\"])\n",
    "plt.grid()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}