just_the_code/roch_mmids_chap_opt_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "48ad6e8a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "remove-cell",
     "colab-keep"
    ]
   },
   "source": [
    "***\n",
    "\n",
    "*Course:* [Math 535](https://people.math.wisc.edu/~roch/mmids/) - Mathematical Methods in Data Science (MMiDS)  \n",
    "*Chapter:* 6-Optimization theory and algorithms   \n",
    "*Author:* [Sebastien Roch](https://people.math.wisc.edu/~roch/), Department of Mathematics, University of Wisconsin-Madison  \n",
    "*Updated:* April 12, 2024   \n",
    "*Copyright:* &copy; 2024 Sebastien Roch\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34a31479",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# You will need the files:\n",
    "#     * mmids.py\n",
    "#     * advertising.csv \n",
    "#     * lebron.csv \n",
    "#     * SAHeart.csv \n",
    "# from https://github.com/MMiDS-textbook/MMiDS-textbook.github.io/tree/main/utils\n",
    "#\n",
    "# IF RUNNING ON GOOGLE COLAB (RECOMMENDED):\n",
    "# \"Upload to session storage\" from the Files tab on the left\n",
    "# Alternative instructions: https://colab.research.google.com/notebooks/io.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c51f0d4",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# PYTHON 3\n",
    "import numpy as np\n",
    "from numpy import linalg as LA\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import mmids\n",
    "seed = 535\n",
    "rng = np.random.default_rng(seed)\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2db1f431",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Motivating example:  deciphering handwriting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d766cbc0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Figure:** Helpful map of ML by scitkit-learn ([Source](https://scikit-learn.org/stable/tutorial/machine_learning_map/index.html))\n",
    "\n",
    "![ml-cheat-sheet](https://scikit-learn.org/stable/_static/ml_map.png)\n",
    "\n",
    "$\\bowtie$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a616c35",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We now turn to classification.\n",
    "\n",
    "Quoting [Wikipedia](https://en.wikipedia.org/wiki/Statistical_classification):\n",
    "\n",
    "> In machine learning and statistics, classification is the problem of identifying to which of a set of categories (sub-populations) a new observation belongs, on the basis of a training set of data containing observations (or instances) whose category membership is known. Examples are assigning a given email to the \"spam\" or \"non-spam\" class, and assigning a diagnosis to a given patient based on observed characteristics of the patient (sex, blood pressure, presence or absence of certain symptoms, etc.). Classification is an example of pattern recognition. In the terminology of machine learning, classification is considered an instance of supervised learning, i.e., learning where a training set of correctly identified observations is available."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc5dbdfc",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We will illustrate this problem on the [MNIST](https://en.wikipedia.org/wiki/MNIST_database) dataset. Quoting [Wikipedia](https://en.wikipedia.org/wiki/MNIST_database) again:\n",
    "\n",
    "> The MNIST database (Modified National Institute of Standards and Technology database) is a large database of handwritten digits that is commonly used for training various image processing systems. The database is also widely used for training and testing in the field of machine learning. It was created by \"re-mixing\" the samples from NIST's original datasets. The creators felt that since NIST's training dataset was taken from American Census Bureau employees, while the testing dataset was taken from American high school students, it was not well-suited for machine learning experiments. Furthermore, the black and white images from NIST were normalized to fit into a 28x28 pixel bounding box and anti-aliased, which introduced grayscale levels. The MNIST database contains 60,000 training images and 10,000 testing images. Half of the training set and half of the test set were taken from NIST's training dataset, while the other half of the training set and the other half of the test set were taken from NIST's testing dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "badd213f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Figure:** MNIST sample images ([Source](https://commons.wikimedia.org/wiki/File:MnistExamples.png))\n",
    "\n",
    "![MNIST sample images](https://upload.wikimedia.org/wikipedia/commons/2/27/MnistExamples.png)\n",
    "\n",
    "$\\bowtie$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deccd8dd",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We first load the data and convert it to an appropriate matrix representation. The data can be accessed with [`torchvision.datasets.MNIST`](https://pytorch.org/vision/stable/generated/torchvision.datasets.MNIST.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f755e7",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import datasets, transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ae9b572",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "# Download and load the MNIST dataset\n",
    "mnist = datasets.MNIST(root='./data', \n",
    "                       train=True, \n",
    "                       download=True, \n",
    "                       transform=transforms.ToTensor())\n",
    "\n",
    "# Convert the dataset to a PyTorch DataLoader\n",
    "train_loader = torch.utils.data.DataLoader(mnist, \n",
    "                                           batch_size=len(mnist), \n",
    "                                           shuffle=False)\n",
    "\n",
    "# Extract images and labels from the DataLoader\n",
    "imgs, labels = next(iter(train_loader))\n",
    "imgs = imgs.squeeze().numpy()\n",
    "labels = labels.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68eb8371",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The [`squeeze()`](https://pytorch.org/docs/stable/generated/torch.Tensor.squeeze.html) above removes the color dimension in the image, which is grayscale. The [`numpy()`](https://pytorch.org/docs/stable/generated/torch.Tensor.numpy.html) converts the PyTorch tensors into Numpy arrays. See [`torch.utils.data.DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) for details on the data loading. We will say more about PyTorch later in this chapter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "297daa2f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "For example, the first image and its label are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0692013",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(imgs[0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65e7d851",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "outputs": [],
   "source": [
    "labels[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c29e8df1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "For now, we look at a subset of the samples: the 0's and 1's."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9aef1e8",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Filter out images with labels 0 and 1\n",
    "mask = (labels == 0) | (labels == 1)\n",
    "imgs01 = imgs[mask]\n",
    "labels01 = labels[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23a68965",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In this new dataset, the first sample is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce9969ef",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(imgs01[0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ec9ac66",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "labels01[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd783f4f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Next, we transform the images into vectors. For this we use the [`reshape()`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.reshape.html) function, which changes the dimensions of an array without changing its data. Here the first dimension, which runs across the samples, remains of the same length `len(imgs01)`. The `-1` is understood to be whatever is needed to \"fit\" the remaining dimensions, here $28 \\times 28 = 784$. In other words, we are effectively \"flattening\" each $28 \\times 28$ image into a single vector of length $784$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23ff4b66",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "imgs01.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba7f1d3c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "X = imgs01.reshape(len(imgs01), -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3076a7a1",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4efc97b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "y = labels01"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebcb79b3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The input data is now of the form $\\{(\\mathbf{x}_i, y_i) : i=1,\\ldots, n\\}$ where $\\mathbf{x}_i \\in \\mathbb{R}^d$ are the features and $y_i \\in \\{0,1\\}$ is the label. Above we use the matrix representation $X \\in \\mathbb{R}^{d \\times n}$ with columns $\\mathbf{x}_i$, $i = 1,\\ldots, n$ and $\\mathbf{y} = (y_1, \\ldots, y_n)^T \\in \\{0,1\\}^n$. \n",
    "\n",
    "Our goal: \n",
    "\n",
    "> learn a classifier from the examples $\\{(\\mathbf{x}_i, y_i) : i=1,\\ldots, n\\}$, that is, a function $\\hat{f} : \\mathbb{R}^d \\to \\mathbb{R}$ such that $\\hat{f}(\\mathbf{x}_i) \\approx y_i$.\n",
    "\n",
    "We may want to enforce that the output is in $\\{0,1\\}$ as well. This problem is referred to as [binary classification](https://en.wikipedia.org/wiki/Binary_classification)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec241024",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "A natural approach to this type of [supervised learning](https://en.wikipedia.org/wiki/Supervised_learning) problem is to define two objects:\n",
    "\n",
    "1. **Family of classifiers:** A class $\\widehat{\\mathcal{F}}$ of classifiers from which to pick $\\hat{f}$.\n",
    "\n",
    "2. **Loss function:** A loss function $\\ell(\\hat{f}, (\\mathbf{x},y))$ which quantifies how good of a fit $\\hat{f}(\\mathbf{x})$ is to $y$.\n",
    "\n",
    "Our goal is then to solve\n",
    "\n",
    "$$\n",
    "\\min_{\\hat{f} \\in \\widehat{\\mathcal{F}}} \\frac{1}{n} \\sum_{i=1}^n \\ell(\\hat{f}, (\\mathbf{x}_i, y_i)),\n",
    "$$\n",
    "\n",
    "that is, we seek to find a classifier among $\\widehat{\\mathcal{F}}$ that minimizes the average loss over the examples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf4fb17",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "For instance, in [logistic regression](https://en.wikipedia.org/wiki/Logistic_regression), we consider linear classifiers of the form\n",
    "\n",
    "$$\n",
    "\\hat{f}(\\mathbf{x})\n",
    "= \\sigma(\\mathbf{x}^T \\boldsymbol{\\theta})\n",
    "\\qquad\n",
    "\\text{with}\n",
    "\\qquad\n",
    "\\sigma(t) = \\frac{1}{1 + e^{-t}}\n",
    "$$\n",
    "\n",
    "where $\\boldsymbol{\\theta} \\in \\mathbb{R}^d$ is a parameter vector. And we use the [cross-entropy loss](https://en.wikipedia.org/wiki/Cross_entropy#Cross-entropy_loss_function_and_logistic_regression)\n",
    "\n",
    "$$\n",
    "\\ell(\\hat{f}, (\\mathbf{x}, y))\n",
    "= -  y \\log(\\sigma(\\mathbf{x}^T \\boldsymbol{\\theta}))\n",
    "- (1-y) \\log(1- \\sigma(\\mathbf{x}^T \\boldsymbol{\\theta})).\n",
    "$$\n",
    "\n",
    "In parametric form, the problem boils down to\n",
    "\n",
    "$$\n",
    "\\min_{\\boldsymbol{\\theta} \\in \\mathbb{R}^d}\n",
    "- \\frac{1}{n} \\sum_{i=1}^n y_i \\log(\\sigma(\\mathbf{x}_i^T \\boldsymbol{\\theta}))\n",
    "- \\frac{1}{n} \\sum_{i=1}^n (1-y_i) \\log(1- \\sigma(\\mathbf{x}_i^T \\boldsymbol{\\theta})).\n",
    "$$\n",
    "\n",
    "To obtain a prediction in $\\{0,1\\}$ here, we could cutoff $\\hat{f}(\\mathbf{x})$ at a threshold $\\tau \\in [0,1]$, that is, return $\\mathbf{1}\\{\\hat{f}(\\mathbf{x}) > \\tau\\}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34928a14",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We will explain in a later chapter where this choice comes from."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c49291c0",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The purpose of this chapter is to develop some of the mathematical theory and algorithms needed to solve this type of optimization formulation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2fcb932",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Background: review of differentiable functions of several variables and introduction to automatic differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fba1fb8a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Brief introduction to automatic differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5cdb4ce",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We illustrate the use of [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation) to compute gradients. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8db8f43e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Quoting [Wikipedia](https://en.wikipedia.org/wiki/Automatic_differentiation):\n",
    "\n",
    "> In mathematics and computer algebra, automatic differentiation (AD), also called algorithmic differentiation or computational differentiation, is a set of techniques to numerically evaluate the derivative of a function specified by a computer program. AD exploits the fact that every computer program, no matter how complicated, executes a sequence of elementary arithmetic operations (addition, subtraction, multiplication, division, etc.) and elementary functions (exp, log, sin, cos, etc.). By applying the chain rule repeatedly to these operations, derivatives of arbitrary order can be computed automatically, accurately to working precision, and using at most a small constant factor more arithmetic operations than the original program. Automatic differentiation is distinct from symbolic differentiation and numerical differentiation (the method of finite differences). Symbolic differentiation can lead to inefficient code and faces the difficulty of converting a computer program into a single expression, while numerical differentiation can introduce round-off errors in the discretization process and cancellation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c895647a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We will use [PyTorch](https://pytorch.org/tutorials/). It uses [tensors](https://pytorch.org/tutorials/beginner/basics/tensorqs_tutorial.html), which in many ways behave similarly to Numpy arrays. See [here](https://pytorch.org/tutorials/beginner/basics/autogradqs_tutorial.html) for a quick introduction. Here is an example. We first initialize the tensors. Here each corresponds to a single real variable. With the option [`requires_grad=True`](https://pytorch.org/docs/stable/generated/torch.Tensor.requires_grad.html#torch.Tensor.requires_grad), we indicate that these are variables with respect to which a gradient will be taken later. We initialize the tensors where the derivatives will be computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4acd06e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Initialize variables\n",
    "x = torch.tensor(1.0, requires_grad=True)\n",
    "y = torch.tensor(2.0, requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef74f662",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The function [`.backward()`](https://pytorch.org/docs/stable/generated/torch.Tensor.backward.html) computes the gradient using backpropagation, to which we will return later. The partial derivatives are accessed with [`.grad`](https://pytorch.org/docs/stable/generated/torch.Tensor.grad.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c2ae210",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Perform automatic differentiation\n",
    "f = 3 * x**2 + torch.exp(x) + y\n",
    "f.backward()  # Compute gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "023cb814",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Print gradients\n",
    "print(x.grad)  # df/dx\n",
    "print(y.grad)  # df/dy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbb2282e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The input parameters can also be vectors, which allows to consider function of large numbers of variables. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5debd5c8",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# New variables for the second example\n",
    "z = torch.tensor([1., 2., 3.], requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d13ee38a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Perform automatic differentiation\n",
    "g = torch.sum(z**2)\n",
    "g.backward()  # Compute gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16add8ec",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Print gradient\n",
    "print(z.grad)  # gradient is (2 z_1, 2 z_2, 2 z_3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d14521e7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is another typical example in a data science context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c82a494",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Variables for the third example\n",
    "X = torch.randn(3, 2)  # Random dataset (features)\n",
    "y = torch.tensor([[1., 0., 1.]])  # Dataset (labels)\n",
    "theta = torch.ones(2, 1, requires_grad=True)  # Parameter assignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90c819b2",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Perform automatic differentiation\n",
    "predict = X @ theta  # Classifier with parameter vector theta\n",
    "loss = torch.sum((predict - y)**2)  # Loss function\n",
    "loss.backward()  # Compute gradients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e21df72a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Print gradient\n",
    "print(theta.grad)  # gradient of loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8063e021",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\blambda}{\\boldsymbol{\\lambda}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5b7c065",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Optimality conditions and convexity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "654347a4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** Consider $f(x) = e^x$. Then $f'(x) = f''(x) = e^x$. Suppose we are interested in approximating $f$ in the interval $[0,1]$. We take $a=0$ and $b=1$ in *Taylor's Theorem*. The linear term is \n",
    "\n",
    "$$\n",
    "f(a) + (x-a) f'(a) = 1 + x e^0 = 1 + x.\n",
    "$$\n",
    "\n",
    "Then for any $x \\in [0,1]$\n",
    "\n",
    "$$\n",
    "f(x) = 1 + x + \\frac{1}{2}x^2 e^{\\xi_x}\n",
    "$$\n",
    "\n",
    "where $\\xi_x \\in (0,1)$ depends on $x$. We get a uniform bound on the error over $[0,1]$ by replacing $\\xi_x$ with its worst possible value over $[0,1]$ \n",
    "\n",
    "$$\n",
    "|f(x) - (1+x)| \\leq \\frac{1}{2}x^2 e^{\\xi_x} \\leq \\frac{e}{2} x^2.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f964d4a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "x = np.linspace(0,1,100)\n",
    "y = np.exp(x)\n",
    "taylor = 1 + x\n",
    "err = (np.exp(1)/2) * x**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dadc838",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x,y,label='f')\n",
    "plt.plot(x,taylor,label='taylor')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d55b03df",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "If we plot the upper and lower bounds, we see that $f$ indeed falls within them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9171b77a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x,y,label='f')\n",
    "plt.plot(x,taylor,label='taylor')\n",
    "plt.plot(x,taylor-err,linestyle=':',color='green',label='lower')\n",
    "plt.plot(x,taylor+err,linestyle='--',color='green',label='upper')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "111b08a2",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d6c4547",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** Let $f(x) = x^3$. Then $f'(x) = 3 x^2$ and $f''(x) = 6 x$ so that $f'(0) = 0$ and $f''(0) \\geq 0$. Hence $x=0$ is a stationary point. But $x=0$ is not a local minimizer. Indeed $f(0) = 0$ but, for any $\\delta > 0$, $f(-\\delta) < 0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd16471f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-2,2,100)\n",
    "y = x**3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f7f10b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x,y)\n",
    "plt.ylim(-5,5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b73ee87",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee1b6c86",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** If we want to minimize $2 x_1^2 + 3 x_2^2$ over all two-dimensional unit vectors $\\mathbf{x} = (x_1, x_2)$, then we can let\n",
    "\n",
    "$$\n",
    "f(\\mathbf{x}) = 2 x_1^2 + 3 x_2^2\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "h_1(\\mathbf{x}) = 1 - x_1^2 - x_2^2 = 1 - \\|\\mathbf{x}\\|^2.\n",
    "$$\n",
    "\n",
    "Observe that we could have chosen a different equality constraint to express the same minimization problem. $\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "823852f8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** **(continued)** Returning to the previous example,\n",
    "\n",
    "$$\n",
    "\\nabla f(\\mathbf{x})\n",
    "= \\left(\n",
    "\\frac{\\partial f(\\mathbf{x})}{\\partial x_1},\n",
    "\\frac{\\partial f(\\mathbf{x})}{\\partial x_2}\n",
    "\\right)\n",
    "= (4 x_1, 6 x_2)\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "\\nabla h_1(\\mathbf{x})\n",
    "= \\left(\n",
    "\\frac{\\partial h_1(\\mathbf{x})}{\\partial x_1},\n",
    "\\frac{\\partial h_1(\\mathbf{x})}{\\partial x_2}\n",
    "\\right)\n",
    "= (- 2 x_1, - 2 x_2).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0615c5b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The conditions in the theorem read\n",
    "\n",
    "\\begin{align*}\n",
    "&4 x_1 - 2 \\lambda_1 x_1  = 0\\\\\n",
    "&6 x_2 - 2 \\lambda_1 x_2  = 0.\n",
    "\\end{align*}\n",
    "\n",
    "The constraint $x_1^2 + x_2^2 = 1$ must also be satisfied. Observe that the linear independence condition is automatically satisfied since there is only one constraint."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3549e1a4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "There are several cases to consider. \n",
    "\n",
    "1- If neither $x_1$ nor $x_2$ is $0$, then the first equation gives $\\lambda_1 = 2$ while the second one gives $\\lambda_1 = 3$. So that case cannot happen.\n",
    "\n",
    "2- If $x_1 = 0$, then $x_2 = 1$ or $x_2 = -1$ by the constraint and the second equation gives $\\lambda_1 = 3$ in either case.\n",
    "\n",
    "3- If $x_2 = 0$, then $x_1 = 1$ or $x_1 = -1$ by the constraint and the first equation gives $\\lambda_1 = 2$ in either case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ef414c3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Does any of these last four solutions, i.e., $(x_1,x_2,\\lambda_1) = (0,1,3)$, $(x_1,x_2,\\lambda_1) = (0,-1,3)$, $(x_1,x_2,\\lambda_1) = (1,0,2)$ and $(x_1,x_2,\\lambda_1) = (-1,0,2)$, actually correspond to a local minimizer?\n",
    "\n",
    "This problem can be solved manually. Indeed, replace $x_2^2 = 1 - x_1^2$ into the objective function to obtain \n",
    "\n",
    "$$\n",
    "2 x_1^2 + 3(1 - x_1^2)\n",
    "= -x_1^2 + 3.\n",
    "$$\n",
    "\n",
    "This is minimized for the largest value that $x_1^2$ can take, namely when $x_1 = 1$ or $x_1 = -1$. Indeed, we must have $0 \\leq x_1^2 \\leq x_1^2 + x_2^2 = 1$. So both $(x_1, x_2) = (1,0)$ and $(x_1, x_2) = (-1,0)$ are global minimizers. A fortiori, they must be local minimizers. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a323fb2",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "What about $(x_1,x_2) = (0,1)$ and $(x_1,x_2) = (0,-1)$? Arguing as above, they in fact correspond to global *maximizers* of the objective function. $\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e68ee84",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** **(continued)** Returning to the previous example, the points satisfying $h_1(\\mathbf{x}) = 0$ sit on the circle of radius $1$ around the origin. We have already seen that \n",
    "\n",
    "$$\n",
    "\\nabla h_1(\\mathbf{x})\n",
    "= \\left(\n",
    "\\frac{\\partial h_1(\\mathbf{x})}{\\partial x_1},\n",
    "\\frac{\\partial h_1(\\mathbf{x})}{\\partial x_2}\n",
    "\\right)\n",
    "= (- 2 x_1, - 2 x_2).\n",
    "$$\n",
    "\n",
    "Here is code plotting these (courtesy of ChatGPT 4). It uses [`numpy.meshgrid`](https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html) to generate a grid of points for $x_1$ and $x_2$, and [`matplotlib.pyplot.contour`](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.contour.html) to plot the constraint set as a [contour line](https://en.wikipedia.org/wiki/Contour_line) (for the constant value $0$) of $h_1$. The gradients are plotted with the [`matplotlib.pyplot.quiver`](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.quiver.html) function, which is used for plotting vectors as arrows. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08b18f75",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "# Define the constraint function\n",
    "def h1(x1, x2):\n",
    "    return 1 - x1**2 - x2**2\n",
    "\n",
    "# Generate a grid of points for x1 and x2\n",
    "x1 = np.linspace(-1.5, 1.5, 400)\n",
    "x2 = np.linspace(-1.5, 1.5, 400)\n",
    "X1, X2 = np.meshgrid(x1, x2)\n",
    "\n",
    "# Compute constraint function on grid\n",
    "H1 = h1(X1, X2)\n",
    "\n",
    "# Points on the constraint where the gradients will be plotted\n",
    "points = [\n",
    "    (0.5, np.sqrt(3)/2),\n",
    "    (-0.5, np.sqrt(3)/2),\n",
    "    (0.5, -np.sqrt(3)/2),\n",
    "    (-0.5, -np.sqrt(3)/2),\n",
    "    (1, 0),\n",
    "    (-1, 0),\n",
    "    (0, 1),\n",
    "    (0, -1)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ca3204a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 6))\n",
    "plt.grid(True)\n",
    "plt.axis('equal')\n",
    "\n",
    "# Plot the constraint set where h1(x1, x2) = 0\n",
    "plt.contour(X1, X2, H1, levels=[0], colors='blue')\n",
    "\n",
    "# Plot gradients of h1 (red) at specified points\n",
    "for x1, x2 in points:\n",
    "    plt.quiver(x1, x2, -2*x1, -2*x2, scale=10, color='red')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e8349b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 6))\n",
    "plt.grid(True)\n",
    "plt.axis('equal')\n",
    "plt.contour(X1, X2, H1, levels=[0], colors='blue')\n",
    "for x1, x2 in points:\n",
    "    plt.quiver(x1, x2, -x1/np.sqrt(x1**2 + x2**2), \n",
    "               -x2/np.sqrt(x1**2 + x2**2), \n",
    "               scale=10, color='red')\n",
    "    plt.quiver(x1, x2, 4*x1/np.sqrt(16 * x1**2 + 36 * x2**2), \n",
    "               6*x2/np.sqrt(16 * x1**2 + 36 * x2**2), \n",
    "               scale=10, color='green')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e4ca33c",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We see that, at $(-1,0)$ and $(1,0)$, the gradient is indeed orthogonal to the first-order feasible directions. $\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bfaffd0",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\bsigma}{\\boldsymbol{\\sigma}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d9ae6a6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Gradient descent and its convergence analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45789b50",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We implement gradient descent in Python. We assume that a function `f` and its gradient `grad_f` are provided. We first code the basic steepest descent step with a step size $\\alpha =$ `alpha`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "435a464f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def desc_update(grad_f, x, alpha):\n",
    "    return x - alpha*grad_f(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e9f4264",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def gd(f, grad_f, x0, alpha=1e-3, niters=int(1e6)):\n",
    "    \n",
    "    xk = x0\n",
    "    for _ in range(niters):\n",
    "        xk = desc_update(grad_f, xk, alpha)\n",
    "\n",
    "    return xk, f(xk)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc46ed90",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We illustrate on a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d642e831",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def f(x): \n",
    "    return (x-1)**2 + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bafe3d0",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-5,5,100)\n",
    "plt.plot(xgrid, f(xgrid))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ae42768",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def grad_f(x):\n",
    "    return 2*(x-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f71f27",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c24a8be",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We found a global minmizer in this case."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad09a3c2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The next example shows that a different local minimizer may be reached depending on the starting point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bbd3eac",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def f(x): \n",
    "    return 4 * (x-1)**2 * (x+1)**2 - 2*(x-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ca5e159",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-2,2,100)\n",
    "plt.plot(xgrid, f(xgrid), label='f')\n",
    "plt.ylim((-1,10))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70ae5bb1",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**CLICK ON TARGET:** If we start gradient descent from $-2$, where will it converge? $\\ddagger$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20660f36",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def grad_f(x): \n",
    "    return 8 * (x-1) * (x+1)**2 + 8 * (x-1)**2 * (x+1) - 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc63a886",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-2,2,100)\n",
    "plt.plot(xgrid, f(xgrid), label='f')\n",
    "plt.plot(xgrid, grad_f(xgrid), label='grad_f')\n",
    "plt.ylim((-10,10))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "302c3e6e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a3fc108",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, -2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6660108",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In the final example, we end up at a stationary point that is not a local minimizer. Here both the first and second derivatives are zero. This is known as a [saddle point](https://en.wikipedia.org/wiki/Saddle_point)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60cf8957",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return x**3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d2343ac",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-2,2,100)\n",
    "plt.plot(xgrid, f(xgrid), label='f')\n",
    "plt.ylim((-10,10))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1e4c95d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def grad_f(x):\n",
    "    return 3 * x**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c649f86e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-2,2,100)\n",
    "plt.plot(xgrid, f(xgrid), label='f')\n",
    "plt.plot(xgrid, grad_f(xgrid), label='grad_f')\n",
    "plt.ylim((-10,10))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fb97e80",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9660a95",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, -2, niters=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ba9720a",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb4096fa",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We revisit our first simple single-variable example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b436561",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def f(x): \n",
    "    return (x-1)**2 + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cff1723",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-5,5,100)\n",
    "plt.plot(xgrid, f(xgrid))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b85923e6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Recall that the first derivative is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5526e5d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def grad_f(x):\n",
    "    return 2*(x-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b90063",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "So the second derivative is $f''(x) = 2$. Hence, this $f$ is $L$-smooth and $m$-strongly convex with $L = m = 2$. The theory we developed suggests taking step size $\\alpha_t = \\alpha = 1/L = 1/2$. It also implies that\n",
    "\n",
    "$$\n",
    "f(x^1) - f(x^*)\n",
    "\\leq \\left(1 - \\frac{m}{L}\\right) [f(x^0) - f(x^*)]\n",
    "= 0.\n",
    "$$\n",
    "\n",
    "We converge in one step! And that holds for any starting point $x^0$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ad7616",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Let's try this!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ab4d26",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, 0, alpha=0.5, niters=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "008490a4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Let's try a different starting point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f201a314",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "gd(f, grad_f, 100, alpha=0.5, niters=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a711fbe7",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07f8014d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Application to logistic regression"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46ba7ad1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We return to logistic regression, which we alluded to in the motivating example of this chapter.\n",
    "\n",
    "The input data is of the form $\\{(\\boldsymbol{\\alpha}_i, b_i) : i=1,\\ldots, n\\}$ where $\\boldsymbol{\\alpha}_i = (\\alpha_{i,1}, \\ldots, \\alpha_{i,d}) \\in \\mathbb{R}^d$ are the features and $b_i \\in \\{0,1\\}$ is the label. As before we use a matrix representation: $A \\in \\mathbb{R}^{n \\times d}$ has rows $\\boldsymbol{\\alpha}_i^T$, $i = 1,\\ldots, n$ and $\\mathbf{b} = (b_1, \\ldots, b_n) \\in \\{0,1\\}^n$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d6b1709",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Logistic model.** We summarize the logistic regression approach. Our goal is to find a function of the features that approximates the probability of the label $1$. For this purpose, we model the [log-odds](https://en.wikipedia.org/wiki/Logit) (or logit function) of the probability of label $1$ as a linear function of the features $\\boldsymbol{\\alpha}  \\in \\mathbb{R}^d$\n",
    "\n",
    "$$\n",
    "\\log \\frac{p(\\mathbf{x}; \\boldsymbol{\\alpha})}{1-p(\\mathbf{x}; \\boldsymbol{\\alpha})}\n",
    "= \\boldsymbol{\\alpha}^T \\mathbf{x}\n",
    "$$\n",
    "\n",
    "where $\\mathbf{x} \\in \\mathbb{R}^d$ is the vector of coefficients (i.e., parameters). Inverting this expression gives\n",
    "\n",
    "$$\n",
    "p(\\mathbf{x}; \\boldsymbol{\\alpha})\n",
    "= \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x})\n",
    "$$\n",
    "\n",
    "where the [sigmoid](https://en.wikipedia.org/wiki/Logistic_function) function is\n",
    "\n",
    "$$\n",
    "\\sigma(z)\n",
    "= \\frac{1}{1 + e^{-z}}\n",
    "$$\n",
    "\n",
    "for $z \\in \\mathbb{R}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be7b6353",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We plot the sigmoid function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4221193",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "outputs": [],
   "source": [
    "def sigmoid(z): \n",
    "    return 1/(1+np.exp(-z))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f2a273e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "grid = np.linspace(-5, 5, 100)\n",
    "plt.plot(grid,sigmoid(grid),'r')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82c77718",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We seek to maximize the probability of observing the data (also known as [likelihood function](https://en.wikipedia.org/wiki/Likelihood_function)) assuming the labels are independent given the features, which is given by\n",
    "\n",
    "$$\n",
    "\\mathcal{L}(\\mathbf{x}; A, \\mathbf{b})\n",
    "= \\prod_{i=1}^n p(\\boldsymbol{\\alpha}_i; \\mathbf{x})^{b_i} \n",
    "(1- p(\\boldsymbol{\\alpha}_i; \\mathbf{x}))^{1-b_i}\n",
    "$$\n",
    "\n",
    "Taking a logarithm, multiplying by $-1/n$ and substituting the sigmoid function, we want to minimize the [cross-entropy loss](https://en.wikipedia.org/wiki/Cross_entropy#Cross-entropy_loss_function_and_logistic_regression)\n",
    "\n",
    "$$\n",
    "\\ell(\\mathbf{x}; A, \\mathbf{b})\n",
    "= \\frac{1}{n} \\sum_{i=1}^n \\left\\{- b_i \\log(\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}))\n",
    "- (1-b_i) \\log(1- \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}))\\right\\}.\n",
    "$$\n",
    "\n",
    "We used standard properties of the logarithm: for $x, y > 0$, $\\log(xy) = \\log x + \\log y$ and $\\log(x^y) = y \\log x$. \n",
    "\n",
    "Hence, we want to solve the minimization problem\n",
    "\n",
    "$$\n",
    "\\min_{\\mathbf{x} \\in \\mathbb{R}^d} \\ell(\\mathbf{x}; A, \\mathbf{b}).\n",
    "$$\n",
    "\n",
    "We are implicitly using here that the logarithm is a strictly increasing function and therefore does not change the global maximum of a function. Multiplying by $-1$ changes the global maximum into a global minimum."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02fe9781",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To use gradient descent, we need the gradient of $\\ell$. We use the *Chain Rule* and first compute the derivative of $\\sigma$ which is\n",
    "\n",
    "$$\n",
    "\\sigma'(z)\n",
    "= \\frac{e^{-z}}{(1 + e^{-z})^2}\n",
    "= \\frac{1}{1 + e^{-z}}\\left(1 - \\frac{1}{1 + e^{-z}}\\right)\n",
    "= \\sigma(z) (1 - \\sigma(z)).\n",
    "$$\n",
    "\n",
    "The latter expression is known as the [logistic differential equation](https://en.wikipedia.org/wiki/Logistic_function#Logistic_differential_equation). It arises in a variety of applications, including the modeling of [population dynamics](https://towardsdatascience.com/covid-19-infection-in-italy-mathematical-models-and-predictions-7784b4d7dd8d). Here it will be a convenient way to compute the gradient. \n",
    "\n",
    "Observe that, for $\\boldsymbol{\\alpha} = (\\alpha_{1}, \\ldots, \\alpha_{d})  \\in \\mathbb{R}^d$, by the *Chain Rule*\n",
    "\n",
    "$$\n",
    "\\nabla\\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x})\n",
    "= \\sigma'(\\boldsymbol{\\alpha}^T \\mathbf{x}) \\nabla (\\boldsymbol{\\alpha}^T \\mathbf{x})\n",
    "= \\sigma'(\\boldsymbol{\\alpha}^T \\mathbf{x}) \\boldsymbol{\\alpha}\n",
    "$$\n",
    "\n",
    "where, throughout, the gradient is with respect to $\\mathbf{x}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae687082",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Alternatively, we can obtain the same formula by applying the single-variable *Chain Rule*\n",
    "\n",
    "\\begin{align*}\n",
    "\\frac{\\partial}{\\partial x_j} \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x})\n",
    "&= \\sigma'(\\boldsymbol{\\alpha}^T \\mathbf{x}) \\frac{\\partial}{\\partial x_j}(\\boldsymbol{\\alpha}^T \\mathbf{x})\\\\\n",
    "&= \\sigma'(\\boldsymbol{\\alpha}^T \\mathbf{x}) \\frac{\\partial}{\\partial x_j}\\left(\\alpha_{j} x_{j} + \\sum_{\\ell=1, \\ell \\neq j}^d \\alpha_{\\ell} x_{\\ell}\\right)\\\\\n",
    "&= \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, \\alpha_{j}\n",
    "\\end{align*}\n",
    "\n",
    "so that\n",
    "\n",
    "\\begin{align*}\n",
    "\\nabla\\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x})\n",
    "&= \\left(\\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, \\alpha_{1}, \\ldots, \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, \\alpha_{d}\\right)\\\\\n",
    "&= \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, (\\alpha_{1}, \\ldots, \\alpha_{d})\\\\\n",
    "&= \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, \\boldsymbol{\\alpha}.\n",
    "\\end{align*}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ade07b9e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "By another application of the *Chain Rule*, since $\\frac{\\mathrm{d}}{\\mathrm{d} z} \\log z = \\frac{1}{z}$,\n",
    "\n",
    "\\begin{align*}\n",
    "&\\nabla\\ell(\\mathbf{x}; A, \\mathbf{b})\\\\\n",
    "&= \\nabla\\left[\\frac{1}{n} \\sum_{i=1}^n \\left\\{- b_i \\log(\\sigma(\\boldsymbol{\\alpha_i}^T \\mathbf{x}))\n",
    "- (1-b_i) \\log(1- \\sigma(\\boldsymbol{\\alpha_i}^T \\mathbf{x}))\\right\\}\\right]\\\\\n",
    "&= - \\frac{1}{n} \\sum_{i=1}^n \\frac{b_i}{\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})} \\nabla\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})\n",
    "- \\frac{1}{n} \\sum_{i=1}^n \\frac{1-b_i}{1- \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})} \\nabla(1 - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}))\\\\\n",
    "&= - \\frac{1}{n} \\sum_{i=1}^n \\frac{b_i}{\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})} \\nabla\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})\n",
    "+ \\frac{1}{n} \\sum_{i=1}^n \\frac{1-b_i}{1- \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})} \\nabla\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}).\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58fc4d36",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Using the expression for the gradient of the sigmoid functions, this is\n",
    "\n",
    "\\begin{align*}\n",
    "&= - \\frac{1}{n} \\sum_{i=1}^n \\frac{b_i}{\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})} \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})) \\,\\boldsymbol{\\alpha}_i\\\\\n",
    "&\\quad\\quad + \\frac{1}{n} \\sum_{i=1}^n \\frac{1-b_i}{1- \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})} \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})) \\,\\boldsymbol{\\alpha}_i\\\\\n",
    "&= - \\frac{1}{n} \\sum_{i=1}^n \\left(\n",
    "b_i (1 - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})) - (1-b_i)\\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}) \n",
    "\\right)\\,\\boldsymbol{\\alpha}_i\\\\\n",
    "&= - \\frac{1}{n} \\sum_{i=1}^n (\n",
    "b_i - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}) \n",
    ") \\,\\boldsymbol{\\alpha}_i.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab28189",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To implement this formula below, it will be useful to re-write it in terms of the matrix representation $A \\in \\mathbb{R}^{n \\times d}$ (which has rows $\\boldsymbol{\\alpha}_i^T$, $i = 1,\\ldots, n$) and $\\mathbf{b} = (b_1, \\ldots, b_n) \\in \\{0,1\\}^n$. Let $\\bsigma : \\mathbb{R}^n \\to \\mathbb{R}$ be the vector-valued function that applies the sigmoid $\\sigma$ entry-wise, i.e., $\\bsigma(\\mathbf{z}) = (\\sigma(z_1),\\ldots,\\sigma(z_n))$ where $\\mathbf{z} = (z_1,\\ldots,z_n)$. Thinking of $\\sum_{i=1}^n (b_i - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x})\\,\\boldsymbol{\\alpha}_i$ as a linear combination of the columns of $A^T$ with coefficients being the entries of the vector $\\mathbf{b} - \\bsigma(A \\mathbf{x})$, we that \n",
    "\n",
    "$$\n",
    "\\nabla\\ell(\\mathbf{x}; A, \\mathbf{b})\n",
    "= - \\frac{1}{n} \\sum_{i=1}^n (\n",
    "b_i - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}) \n",
    ") \\,\\boldsymbol{\\alpha}_i\n",
    "= -\\frac{1}{n} A^T [\\mathbf{b} - \\bsigma(A \\mathbf{x})].\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e58cb2ad",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We turn to the Hessian. By symmetry, we can think of the $j$-th column of the Hessian as the gradient of the partial derivative with respect to $x_j$. Hence we start by computing the gradient of the $j$-th entry of the summands in the gradient of $\\ell$. We note that, for $\\boldsymbol{\\alpha} = (\\alpha_{1}, \\ldots, \\alpha_{d})  \\in \\mathbb{R}^d$,\n",
    "\n",
    "$$\n",
    "\\nabla [(b - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, \\alpha_{j}] \n",
    "= - \\nabla [\\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x})] \\, \\alpha_{j} \n",
    "=  - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}))\\, \\boldsymbol{\\alpha}\\alpha_{j}.\n",
    "$$\n",
    "\n",
    "Thus, using the fact that $\\boldsymbol{\\alpha} \\alpha_{j}$ is the $j$-th column of the matrix $\\boldsymbol{\\alpha} \\boldsymbol{\\alpha}^T$, we get\n",
    "\n",
    "$$\n",
    "\\mathbf{H}_{\\ell}(\\mathbf{x}; A, \\mathbf{b})\n",
    "= \\frac{1}{n} \\sum_{i=1}^n \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}) (1 - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}))\\, \\boldsymbol{\\alpha}_i \\boldsymbol{\\alpha}_i^T\n",
    "$$\n",
    "\n",
    "where $\\mathbf{H}_{\\ell}(\\mathbf{x}; A, \\mathbf{b})$ indicates the Hessian with respect to the $\\mathbf{x}$ variables, for fixed $A, \\mathbf{b}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc58b6dd",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "For step size $\\beta$, one step of gradient descent is therefore\n",
    "\n",
    "$$\n",
    "\\mathbf{x}^{t+1}\n",
    "= \\mathbf{x}^{t} +\\beta \\frac{1}{n} \\sum_{i=1}^n (\n",
    "b_i - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}^t) \n",
    ") \\,\\boldsymbol{\\alpha}_i.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e0ff12d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** Before implementing GD for logistic regression, we return to our proof of convergence for smooth functions using a special case. We illustrate it on a random dataset. The functions $\\hat{f}$, $\\mathcal{L}$ and $\\frac{\\partial}{\\partial x}\\mathcal{L}$ are defined next."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fccf0cd2",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def fhat(x,a):\n",
    "    return 1 / ( 1 + np.exp(-np.outer(x,a)) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1080f07",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def loss(x,a,b): \n",
    "    return np.mean(-b*np.log(fhat(x,a)) - (1 - b)*np.log(1 - fhat(x,a)), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8443062e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def grad(x,a,b):\n",
    "    return -np.mean((b - fhat(x,a))*a, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d51abbf6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "n = 10000\n",
    "a = 2*rng.uniform(0,1,n) - 1\n",
    "b = rng.integers(2, size=n)\n",
    "x = np.linspace(-1,1,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c0dfa41",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x, loss(x,a,b), label='loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "def8061d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We plot next the upper and lower bounds in the *Quadratic Bound for Smooth Functions* around $x = x_0$. It turns out we can take $L=1$ because all features are uniformly random between $-1$ and $1$. Observe that minimizing the upper quadratic bound leads to a decrease in $\\mathcal{L}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4835cacf",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "x0 = -0.3\n",
    "x = np.linspace(x0-0.05,x0+0.05,100)\n",
    "upper = loss(x0,a,b) + (x - x0)*grad(x0,a,b) + (1/2)*(x - x0)**2 # upper approximation\n",
    "lower = loss(x0,a,b) + (x - x0)*grad(x0,a,b) - (1/2)*(x - x0)**2 # lower approximation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "032e88b3",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x, loss(x,a,b), label='loss')\n",
    "plt.plot(x, upper, label='upper')\n",
    "plt.plot(x, lower, label='lower')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faf67d65",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49fdca9c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We modify our implementation of gradient descent to take a dataset as input. (That will also be useful to generalize to so-called stochastic gradient descent; see below.) Recall that to run gradient descent, we first implement a function computing a descent update. It takes as input a function `grad_fn` computing the gradient itself, as well as a current iterate and a step size. We now also feed a dataset as additional input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b21b147a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def desc_update_for_logreg(grad_fn, A, b, curr_x, beta):\n",
    "    gradient = grad_fn(curr_x, A, b)\n",
    "    return curr_x - beta*gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec474036",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We are ready to implement GD. Our function takes as input a function `loss_fn` computing the objective, a function `grad_fn` computing the gradient, the dataset `A` and `b`, and an initial guess `init_x`. Optional parameters are the step size and the number of iterations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b16d3938",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def gd_for_logreg(loss_fn, grad_fn, A, b, init_x, beta=1e-3, niters=int(1e5)):\n",
    "    \n",
    "    # initialization\n",
    "    curr_x = init_x\n",
    "    \n",
    "    # until the maximum number of iterations\n",
    "    for iter in range(niters):\n",
    "        curr_x = desc_update_for_logreg(grad_fn, A, b, curr_x, beta)\n",
    "    \n",
    "    return curr_x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7521882f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To implement `loss_fn` and `grad_fn`, we define the sigmoid as above. Below, `pred_fn` is $\\bsigma(A \\mathbf{x})$. Here we write the loss function as\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\ell(\\mathbf{x}; A, \\mathbf{b})\n",
    "&= \\frac{1}{n} \\sum_{i=1}^n \\left\\{- b_i \\log(\\sigma(\\boldsymbol{\\alpha_i}^T \\mathbf{x}))\n",
    "- (1-b_i) \\log(1- \\sigma(\\boldsymbol{\\alpha_i}^T \\mathbf{x}))\\right\\}\\\\\n",
    "&= \\mathrm{mean}\\left(-\\mathbf{b} \\odot \\mathbf{log}(\\bsigma(A \\mathbf{x})) - (\\mathbf{1} - \\mathbf{b}) \\odot \\mathbf{log}(\\mathbf{1} - \\bsigma(A \\mathbf{x}))\\right),\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $\\odot$ is the Hadamard product, or element-wise product (for example $\\mathbf{u} \\odot \\mathbf{v} = (u_1 v_1, \\ldots,u_n v_n)^T$), the logarithm (denoted in bold) is applied element-wise and $\\mathrm{mean}(\\mathbf{z})$ is the mean of the entries of $\\mathbf{z}$ (i.e., $\\mathrm{mean}(\\mathbf{z}) = n^{-1} \\sum_{i=1}^n z_i$). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7e99957",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def pred_fn(x, A): \n",
    "    return sigmoid(A @ x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14ce001b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def loss_fn(x, A, b): \n",
    "    return np.mean(-b*np.log(pred_fn(x, A)) - (1 - b)*np.log(1 - pred_fn(x, A)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f31d1d84",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def grad_fn(x, A, b):\n",
    "    return -A.T @ (b - pred_fn(x, A))/len(b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53396a7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We can choosed a step size based on the smoothness of the objective as above. Recall that [`numpy.linalg.norm`](https://numpy.org/doc/stable/reference/generated/numpy.linalg.norm.html) computes the Frobenius norm by default. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e957b2c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def stepsize_for_logreg(A, b):\n",
    "    L = LA.norm(A)**2 /len(b)\n",
    "    return 1/L"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa33f7aa",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Lebron James 2017 NBA Playoffs dataset** We start with a simple dataset from UC Berkeley's [DS100](http://www.ds100.org) course. The file `lebron.csv` is available [here](https://github.com/MMiDS-textbook/MMiDS-textbook.github.io/tree/main/utils/datasets). Quoting a previous version of the course's textbook:\n",
    "\n",
    "> In basketball, players score by shooting a ball through a hoop. One such player, LeBron James, is widely considered one of the best basketball players ever for his incredible ability to score. LeBron plays in the National Basketball Association (NBA), the United States's premier basketball league. We've collected a dataset of all of LeBron's attempts in the 2017 NBA Playoff Games using the NBA statistics website (https://stats.nba.com/).\n",
    "\n",
    "We first load the data and look at its summary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3a1d02c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('lebron.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a39dfef9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "df.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f1e089e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The two columns we will be interested in are `shot_distance` (LeBron's distance from the basket when the shot was attempted (ft)) and `shot_made` (0 if the shot missed, 1 if the shot went in). As the summary table above indicates, the average distance was `10.6953` and the frequency of shots made was `0.565104`. We extract those two columns and display them on a scatter plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14daded1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "feature = df['shot_distance']\n",
    "label = df['shot_made']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b3a22d9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(feature, label, alpha=0.2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b983060e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "As you can see, this kind of data is hard to vizualize because of the superposition of points with the same $x$ and $y$-values. One trick is to jiggle the $y$'s a little bit by adding Gaussian noise. We do this next and plot again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b5fefcb",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "label_jitter = label + 0.05*rng.normal(0,1,len(label))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f247da",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(feature, label_jitter, alpha=0.2)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81cf373f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We apply GD to logistic regression. We first construct the data matrices $A$ and $\\mathbf{b}$. To allow an affine function of the features, we add a column of $1$'s as we have done before.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67db400e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "A = np.stack((np.ones(len(label)),feature),axis=-1)\n",
    "b = label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9809c0b3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We run GD starting from $(0,0)$ with a step size computed from the smoothness of the objective as above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cf69d3f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "stepsize = stepsize_for_logreg(A, b)\n",
    "print(stepsize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "416ff9cd",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "init_x = np.zeros(A.shape[1])\n",
    "best_x = gd_for_logreg(loss_fn, grad_fn, A, b, init_x, beta=stepsize)\n",
    "print(best_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a313cb7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Finally we plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8825f57e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "grid = np.linspace(np.min(feature), np.max(feature), 100)\n",
    "feature_grid = np.stack((np.ones(len(grid)),grid),axis=-1)\n",
    "predict_grid = sigmoid(feature_grid @ best_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4a7cfe",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(feature, label_jitter, alpha=0.2)\n",
    "plt.plot(grid,predict_grid,'r')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa7e4c73",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Stochastic gradient descent** In [stochastic gradient descent](https://en.wikipedia.org/wiki/Stochastic_gradient_descent) (SGD), a variant of gradient descent, we pick a sample $I_t$ uniformly at random in $\\{1,\\ldots,n\\}$ and update as follows\n",
    "\n",
    "$$\n",
    "\\mathbf{x}^{t+1}\n",
    "= \\mathbf{x}^{t} +\\beta \\, (\n",
    "b_{I_t} - \\sigma(\\boldsymbol{\\alpha}_{I_t}^T \\mathbf{x}^t) \n",
    ") \\, \\boldsymbol{\\alpha}_{I_t}.\n",
    "$$\n",
    "\n",
    "For the mini-batch version of SGD, we pick a random sub-sample $\\mathcal{B}_t \\subseteq \\{1,\\ldots,n\\}$ of size $B$\n",
    "\n",
    "$$\n",
    "\\mathbf{x}^{t+1}\n",
    "= \\mathbf{x}^{t} +\\beta \\frac{1}{B} \\sum_{i\\in \\mathcal{B}_t} (\n",
    "b_i - \\sigma(\\boldsymbol{\\alpha}_i^T \\mathbf{x}^t) \n",
    ") \\,\\boldsymbol{\\alpha}_i.\n",
    "$$\n",
    "\n",
    "The key observation about the two stochastic updates above is that, in expectation, they perform a step of gradient descent. That turns out to be enough and it has computational advantages.\n",
    "\n",
    "The only modification needed to the code is to pick a random mini-batch which can be fed to the descent update sub-routine as dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83336dc",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def sgd_for_logreg(loss_fn, grad_fn, A, b, \n",
    "                   init_x, beta=1e-3, niters=int(1e5), batch=40):\n",
    "    \n",
    "    # initialization\n",
    "    curr_x = init_x\n",
    "    \n",
    "    # until the maximum number of iterations\n",
    "    nsamples = len(b)\n",
    "    for _ in range(niters):\n",
    "        I = rng.integers(nsamples, size=batch)\n",
    "        curr_x = desc_update_for_logreg(grad_fn, A[I,:], b[I], curr_x, beta)\n",
    "    \n",
    "    return curr_x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "837ca3b9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**South African Heart Disease dataset** We analyze a dataset from [[ESL](https://web.stanford.edu/~hastie/ElemStatLearn/)], which can be downloaded [here](https://web.stanford.edu/~hastie/ElemStatLearn/data.html). Quoting [[ESL](https://web.stanford.edu/~hastie/ElemStatLearn/), Section 4.4.2] \n",
    "\n",
    "> The data [...] are a subset of the Coronary Risk-Factor Study (CORIS) baseline survey, carried out in three rural areas of the Western Cape, South Africa (Rousseauw et al., 1983). The aim of the study was to establish the intensity of ischemic heart disease risk factors in that high-incidence region. The data represent white males between 15 and 64, and the response variable is the presence or absence of myocardial infarction (MI) at the time of the survey (the overall prevalence of MI was 5.1% in this region). There are 160 cases in our data set, and a sample of 302 controls. These data are described in more detail in Hastie and Tibshirani (1987).\n",
    "\n",
    "We load the data, which we slightly reformatted and look at a summary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb2cb862",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('SAHeart.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1f4d9c",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "df.describe()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a990d427",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Our goal to predict `chd`, which stands for coronary heart disease, based on the other variables (which are briefly described [here](https://web.stanford.edu/~hastie/ElemStatLearn/datasets/SAheart.info.txt)). We use logistic regression again. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6143b9c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We first construct the data matrices. We only use three of the predictors, as the convergence is quite slow. Try it for yourself! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ced4d41",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "feature = df[['tobacco', 'ldl', 'age']].to_numpy()\n",
    "print(feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "490e0ce8",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "label = df['chd'].to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b471fa7b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "A = np.concatenate((np.ones((len(label),1)),feature),axis=1)\n",
    "print(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b698fc73",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "b = label"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a991b7c8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We use the same functions `loss_fn` and `grad_fn`, which were written for general logistic regression problems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "968299ac",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "init_x = np.zeros(A.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d903f67",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "stepsize = stepsize_for_logreg(A, b)\n",
    "print(stepsize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ece836ba",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "best_x = gd_for_logreg(loss_fn, grad_fn, A, b, \n",
    "                       init_x, beta=stepsize, niters=1000000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54322622",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(best_x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a79c0d7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The outcome is harder to vizualize. To get a sense of how accurate the result is, we compare our predictions to the true labels. By prediction, let us say that we mean that we predict label $1$ whenever $\\sigma(\\boldsymbol{\\alpha}^T \\mathbf{x}) > 1/2$. We try this on the training set. (A better approach would be to split the data into training and testing sets, but we will not do this here.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25067a11",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def logis_acc(x, A, b):\n",
    "    return np.sum((pred_fn(x, A) > 0.5) == b)/len(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e722161a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "logis_acc(best_x, A, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2855a08f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We also try mini-batch stochastic gradient descent (SGD). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "575176ca",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "init_x = np.zeros(A.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b7eea5f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "best_x = sgd_for_logreg(loss_fn, grad_fn, A, b, \n",
    "                        init_x, beta=stepsize, niters=1000000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93fffef7",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(best_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35fd2b9b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "logis_acc(best_x, A, b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48b3264c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Implementing gradient descent in PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab8739a8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Rather than explicitly specifying the gradient function, we could use PyTorch to compute it automatically. This is done next. Note that the descent update is done within [`with torch.no_grad()`](https://pytorch.org/docs/stable/generated/torch.no_grad.html), which ensures that the update operation itself is not tracked for gradient computation. Here the input `x0` as well as the output `xk.numpy(force=True)` are Numpy arrays. The function [`numpy()`](https://pytorch.org/docs/stable/generated/torch.Tensor.numpy.html) converts a PyTorch tensor to a Numpy array (see the documentation for an explanation of the `force=True` option)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf2e56d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def gd_with_ad(f, x0, alpha=1e-3, niters=int(1e6)):\n",
    "    xk = torch.tensor(x0, \n",
    "                      requires_grad=True, \n",
    "                      dtype=torch.float)\n",
    "    \n",
    "    for _ in range(niters):\n",
    "        # Compute the function value and its gradient\n",
    "        value = f(xk)\n",
    "        value.backward()\n",
    "\n",
    "        # Perform a gradient descent step\n",
    "        with torch.no_grad():  # Temporarily set all requires_grad flags to False\n",
    "            xk -= alpha * xk.grad\n",
    "\n",
    "        # Zero the gradients for the next iteration\n",
    "        xk.grad.zero_()\n",
    "\n",
    "    return xk.numpy(force=True), f(xk).item()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2b85139",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We revisit a previous example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fefb24d4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return x**3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e62facc5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "xgrid = np.linspace(-2,2,100)\n",
    "plt.plot(xgrid, f(xgrid), label='f')\n",
    "plt.ylim((-10,10))\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a36abcd",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "gd_with_ad(f, 2, niters=int(1e4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f577dd",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "gd_with_ad(f, -2, niters=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61e28e62",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aaf5a4ea",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Backpropagation and neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "324a67b2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**The `Advertising` dataset and the least-squares solution** We return to the `Advertising` dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79b80247",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "df = pd.read_csv('advertising.csv')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ce0bb25",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "n = len(df.index)\n",
    "print(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5515218",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We first compute the solution using the least-squares approach we detailed previously. We use [`numpy.column_stack`](https://numpy.org/doc/stable/reference/generated/numpy.column_stack.html#numpy.column_stack) to add a column of ones to the feature vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a6a5eb5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "TV = df['TV'].to_numpy()\n",
    "radio = df['radio'].to_numpy()\n",
    "newspaper = df['newspaper'].to_numpy()\n",
    "sales = df['sales'].to_numpy()\n",
    "features = np.stack((TV, radio, newspaper), axis=-1)\n",
    "A = np.column_stack((np.ones(n), features))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d27718f2",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "coeff = mmids.ls_by_qr(A, sales)\n",
    "print(coeff)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcab4629",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "np.mean((A @ coeff - sales)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84d362b2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Solving the problem using PyTorch** We will be using PyTorch to implement the previous method. We first convert the data into PyTorch tensors. We then use [`torch.utils.data.TensorDataset`](https://pytorch.org/docs/stable/data.html#torch.utils.data.TensorDataset) to create the dataset. Finally, [`torch.utils.data.DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader) provides the utilities to load the data in batches for training. We take mini-batches of size `BATCH_SIZE = 64` and we apply a random permutation of the samples on every pass (with the option `shuffle=True`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1632bc3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Convert data to PyTorch tensors\n",
    "features_tensor = torch.tensor(features, \n",
    "                               dtype=torch.float32)\n",
    "sales_tensor = torch.tensor(sales, \n",
    "                            dtype=torch.float32).view(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f7ab7ec",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Create a dataset and dataloader for training\n",
    "BATCH_SIZE = 64\n",
    "train_dataset = TensorDataset(features_tensor, sales_tensor)\n",
    "train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8677329",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Now we construct our model. It is simply an affine map from $\\mathbb{R}^3$ to $\\mathbb{R}$. Note that there is no need to pre-process the inputs by adding $1$s. A constant term (or \"bias variable\") is automatically added by PyTorch (unless one chooses the option [`bias=False`](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f400cf",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Define the model using nn.Sequential\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(3, 1)  # 3 input features, 1 output value\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52816aba",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Finally, we are ready to run an optimization method of our choice on the loss function, which are specified next. There are many [optimizers](https://pytorch.org/docs/stable/optim.html#algorithms) available. (See this [post](https://hackernoon.com/demystifying-different-variants-of-gradient-descent-optimization-algorithm-19ae9ba2e9bc) for a brief explanation of many common optimizers.) Here we use SGD as the optimizer. And the loss function is the MSE. A quick tutorial is [here](https://pytorch.org/tutorials/beginner/basics/optimization_tutorial.html).\n",
    "\n",
    "Choosing the right number of passes (i.e. epochs) through the data requires some experimenting. Here $10^4$ suffices. But in the interest of time, we will run it only for $100$ epochs. As you will see from the results, this is not quite enough. On each pass, we compute the output of the current model, use `backward()` to obtain the gradient, and then perform a descent update with `step()`. We also have to reset the gradients first (otherwise they add up by default). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f98a1881",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Compile the model: define loss function and optimizer\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4184af13",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Train the model\n",
    "epochs = 100\n",
    "for epoch in range(epochs):\n",
    "    for inputs, targets in train_loader:\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, targets)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    if (epoch+1) % 10 == 0:\n",
    "        print(f\"Epoch {epoch+1}/{epochs}, Loss: {loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97156a22",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The final parameters and loss are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08499b03",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Get and print the model weights and bias\n",
    "weights = model[0].weight.detach().numpy()\n",
    "bias = model[0].bias.detach().numpy()\n",
    "print(\"Weights:\", weights)\n",
    "print(\"Bias:\", bias)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea649765",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Evaluate the model\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    total_loss = 0\n",
    "    for inputs, targets in train_loader:\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, targets)\n",
    "        total_loss += loss.item()\n",
    "        \n",
    "    print(f\"Mean Squared Error on Training Set: {total_loss / len(train_loader)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f7d9dd0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**MNIST dataset** We will use the [MNIST](https://en.wikipedia.org/wiki/MNIST_database) dataset introduced earlier in the chapter. This example is inspired by [these](https://pytorch.org/tutorials/beginner/basics/quickstart_tutorial.html) [tutorial](https://www.tensorflow.org/tutorials/keras/classification)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcb16cf1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Figure:** MNIST sample images ([Source](https://commons.wikimedia.org/wiki/File:MnistExamples.png))\n",
    "\n",
    "![MNIST sample images](https://upload.wikimedia.org/wikipedia/commons/2/27/MnistExamples.png)\n",
    "\n",
    "$\\bowtie$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e01db5",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We first load the data. As before, the training dataset is a tensor -- think matrix with $3$ indices. One index runs through the $60,000$ training images, while the other two indices run through the horizontal and vertical pixel axes of each image. Here each image is $28 \\times 28$. The training labels are between $0$ and $9$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa05d9ee",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "# Load and normalize the MNIST dataset\n",
    "train_dataset = datasets.MNIST(root='./data', \n",
    "                               train=True, \n",
    "                               download=True, \n",
    "                               transform=transforms.ToTensor())\n",
    "\n",
    "test_dataset = datasets.MNIST(root='./data', \n",
    "                              train=False, \n",
    "                              download=True, \n",
    "                              transform=transforms.ToTensor())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "372fccbe",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 32\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, \n",
    "                                           batch_size=BATCH_SIZE, \n",
    "                                           shuffle=True)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, \n",
    "                                          batch_size=BATCH_SIZE, \n",
    "                                          shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d24fb19",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Implementation** We implement multinomial logistic regression to learn a classifier for the MNIST data. We first check for the availability of GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0caab6e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Check for GPU availability\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(\"Using device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3670382",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In PyTorch, composition of functions can be achieved with [`torch.nn.Sequential`](https://pytorch.org/docs/stable/generated/torch.nn.Sequential.html). Our model is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72aac70d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Define the model using nn.Sequential and move it to the device (GPU if available)\n",
    "model = nn.Sequential(\n",
    "    nn.Flatten(),\n",
    "    nn.Linear(28 * 28, 10)\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f286715",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The [`torch.nn.Flatten`](https://pytorch.org/docs/stable/generated/torch.nn.Flatten.html) layer turns each input image into a vector of size $784$ (where $784 = 28^2$ is the number of pixels in each image). The final output is $10$-dimensional."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6624fffa",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here we use the [`torch.optim.Adam`](https://pytorch.org/docs/stable/generated/torch.optim.Adam.html) optimizer (you can try SGD, but it is slow). The loss function is the [cross-entropy](https://en.wikipedia.org/wiki/Cross_entropy), as implemented by [`torch.nn.CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52ed894a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Compile the model: define loss function and optimizer\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e279320",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In the interest of time, we train for 3 epochs only. An epoch is one training iteration where all samples are iterated once (in a randomly shuffled order)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "410b8560",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Training function\n",
    "def train(dataloader, model, loss_fn, optimizer, device):\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "        \n",
    "        # Compute prediction error\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "        \n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54ea3d51",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Training loop\n",
    "epochs = 3  # Adjust the number of epochs as needed\n",
    "for epoch in range(epochs):\n",
    "    train(train_loader, model, loss_fn, optimizer, device)\n",
    "    \n",
    "    if (epoch+1) % 1 == 0:\n",
    "        print(f\"Epoch {epoch+1}/{epochs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "470fa959",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Because of the issue of [overfitting](https://en.wikipedia.org/wiki/Overfitting), we use the *test* images to assess the performance of the final classifier. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5596c27c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Evaluation function\n",
    "def test(dataloader, model, device):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()\n",
    "    test_loss, correct = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(dim=1) == y).type(torch.float).sum().item()\n",
    "    test_loss /= num_batches\n",
    "    accuracy = correct / size\n",
    "    print(f\"Test error: {(100*accuracy):>0.1f}% accuracy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ef32ae",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "# Evaluate the model\n",
    "test(test_loader, model, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc3a9e3e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To make a prediction, we take a [`torch.nn.functional.softmax`](https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html) of the output of our model. It transforms the output into a probability for each label. It is implicitly included in the cross-entropy loss, but is not actually part of the trained model. (Note that the softmax itself has no parameter.) \n",
    "\n",
    "As an illustration, we do this for each test image. We use [`torch.cat`](https://pytorch.org/docs/stable/generated/torch.cat.html) to concatenate a sequence of tensors into a single tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c04d497b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def predict_softmax(dataloader, model, device):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    model.eval()  # Set the model to evaluation mode\n",
    "    predictions = []\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            probabilities = F.softmax(pred, dim=1)\n",
    "            predictions.append(probabilities.cpu())  # Move predictions to CPU\n",
    "\n",
    "    return torch.cat(predictions, dim=0)\n",
    "\n",
    "predictions = predict_softmax(test_loader, model, device).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19fd2756",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The result for the first test image is shown below. To make a prediction, we choose the label with the highest probability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e171f07",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(predictions[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5281abf5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "predictions[0].argmax(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d553d17a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The truth is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7ec2935",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "images, labels = next(iter(test_loader))\n",
    "images = images.squeeze().numpy()\n",
    "labels = labels.numpy()\n",
    "labels[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4771639c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Above, `next(iter(test_loader))` loads the first batch of test images. (See [here](https://docs.python.org/3/tutorial/classes.html#iterators) for background on iterators in Python.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "965950e3",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The following code, adapted from [here](https://www.tensorflow.org/tutorials/keras/classification), provides a neat vizualization of the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8e6650a",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "class_names = ['0', '1', '2', '3', '4',\n",
    "               '5', '6', '7', '8', '9']\n",
    "\n",
    "def plot_image(predictions_array, true_label, img):\n",
    "    \n",
    "    plt.grid(False)\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "\n",
    "    plt.imshow(img, cmap=plt.cm.binary)\n",
    "\n",
    "    predicted_label = np.argmax(predictions_array)\n",
    "    if predicted_label == true_label:\n",
    "        color = 'blue'\n",
    "    else:\n",
    "        color = 'red'\n",
    "\n",
    "    plt.xlabel(f\"{class_names[predicted_label]} {100*np.max(predictions_array):2.0f}% ({class_names[true_label]})\", \n",
    "               color=color)\n",
    "\n",
    "def plot_value_array(predictions_array, true_label):\n",
    "    \n",
    "    plt.grid(False)\n",
    "    plt.xticks(range(10))\n",
    "    plt.yticks([])\n",
    "    thisplot = plt.bar(range(10), predictions_array, color=\"#777777\")\n",
    "    plt.ylim([0, 1])\n",
    "    predicted_label = np.argmax(predictions_array)\n",
    " \n",
    "    thisplot[predicted_label].set_color('red')\n",
    "    thisplot[true_label].set_color('blue')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dc97896",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here's the first one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a566016",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Visualization code for individual image\n",
    "i = 0\n",
    "plt.figure(figsize=(6,3))\n",
    "plt.subplot(1,2,1)\n",
    "plot_image(predictions[i], labels[i], images[i])  \n",
    "plt.subplot(1,2,2)\n",
    "plot_value_array(predictions[i], labels[i]) \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "326a4b7e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "This one is a little less clear. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5720ad47",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Visualization code for individual and multiple images\n",
    "i = 11\n",
    "plt.figure(figsize=(6,3))\n",
    "plt.subplot(1,2,1)\n",
    "plot_image(predictions[i], labels[i], images[i])  \n",
    "plt.subplot(1,2,2)\n",
    "plot_value_array(predictions[i], labels[i]) \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7689687b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "This one is wrong."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "432043f6",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Visualization code for individual and multiple images\n",
    "i = 8\n",
    "plt.figure(figsize=(6,3))\n",
    "plt.subplot(1,2,1)\n",
    "plot_image(predictions[i], labels[i], images[i])  \n",
    "plt.subplot(1,2,2)\n",
    "plot_value_array(predictions[i], labels[i]) \n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "334f9de3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Implementation** We implement a neural network in PyTorch. We use the MNIST dataset again. We have already loaded it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ea22cf9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We construct a three-layer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe14f59b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Define the model using nn.Sequential\n",
    "model = nn.Sequential(\n",
    "    nn.Flatten(),                      # Flatten the input\n",
    "    nn.Linear(28 * 28, 32),            # First Linear layer with 32 nodes\n",
    "    nn.Sigmoid(),                      # Sigmoid activation function\n",
    "    nn.Linear(32, 10)                  # Second Linear layer with 10 nodes (output layer)\n",
    ").to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeb7545c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "As we did for multinomial logistic regression, we use the Adam optimizer and the cross-entropy loss. We also monitor progress by keeping track of the accuracy on the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7901d55e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Define the loss function and the optimizer\n",
    "loss_fn = nn.CrossEntropyLoss()  \n",
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caff08a0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Again, we train for 3 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "259d667d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Training loop\n",
    "epochs = 3  # Adjust the number of epochs as needed\n",
    "for epoch in range(epochs):\n",
    "    train(train_loader, model, loss_fn, optimizer, device)\n",
    "    \n",
    "    if (epoch+1) % 1 == 0:\n",
    "        print(f\"Epoch {epoch+1}/{epochs}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25939825",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "On the test data, we get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a3d43ec",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Evaluate the model\n",
    "test(test_loader, model, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6863745",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "This is a significantly more accurate model than what we obtained using multinomial logistic regression. One can do even better using a neural network tailored for images, known as [convolutional neural networks](https://cs231n.github.io/convolutional-networks/)."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}