{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Logistic and linear regression with deterministic and stochastic first order methods\n",
    "\n",
    "    Lab 2 : Optimisation - DataScience Master\n",
    "    Authors : Alexandre Gramfort, Hicham Janati\n",
    "   \n",
    "The aim of this lab is to implement and compare various batch and stochastic algorithms for linear and logistic regression with ridge penalization. \n",
    "\n",
    "The following methods are compared in this notebook.\n",
    "\n",
    "**Batch (deterministic) methods**\n",
    "\n",
    "- gradient descent (GD)\n",
    "- accelerated gradient descent (AGD)\n",
    "\n",
    "**Stochastic algorithms**\n",
    "\n",
    "- stochastic gradient descent (SGD)\n",
    "- stochastic averaged gradient (SAG)\n",
    "- stochastic variance reduced gradient (SVRG)\n",
    "\n",
    "Note that we consider as use-cases logistic and linear regression with ridge penalization only, although most of the algorithms below can be used with many other models, and other types of penalization, even non-smooth ones, such as the $\\ell_1$ penalization.\n",
    "\n",
    "\n",
    "## VERY IMPORTANT\n",
    "\n",
    "- This work **must be done by pairs of students**.\n",
    "- **Each** student must send their work using the **moodle platform**, please\n",
    "  check the deadline on moodle.\n",
    "- This means that **each student in the pair sends the same file**\n",
    "- The **name of the file must be** constructed as in the next cell\n",
    "\n",
    "### How to construct the name of your file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change here using YOUR first and last names\n",
    "fn1 = \"hicham\"\n",
    "ln1 = \"janati\"\n",
    "fn2 = \"alexandre\"\n",
    "ln2 = \"gramfort\"\n",
    "\n",
    "filename = \"_\".join(map(lambda s: s.strip().lower(), \n",
    "                        [\"lab2\", ln1, fn1, \"and\", ln2, fn2])) + \".ipynb\"\n",
    "print(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gentle reminder: no evaluation if you don't respect this EXACTLY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Table of content\n",
    "\n",
    "[1. Loss functions, gradients and step-sizes](#loss)<br>\n",
    "[2. Generate a dataset](#data)<br>\n",
    "[3. Deterministic methods](#batch)<br>\n",
    "[4. Stochastic methods](#stoc)<br>\n",
    "[5. Numerical comparison](#comp)<br>\n",
    "[6. Conclusion](#conc)<br>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",
    "import numba\n",
    "import numpy as np\n",
    "from scipy.linalg import norm\n",
    "import matplotlib.pyplot as plt\n",
    "from numba import njit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose a large font size by default and use tex for math\n",
    "usetex = False  # change this to True if you have a working LaTeX install\n",
    "\n",
    "fontsize = 16\n",
    "params = {\n",
    "    \"axes.labelsize\": fontsize + 2,\n",
    "    \"font.size\": fontsize + 2,\n",
    "    \"legend.fontsize\": fontsize + 2,\n",
    "    \"xtick.labelsize\": fontsize,\n",
    "    \"ytick.labelsize\": fontsize,\n",
    "    \"text.usetex\": usetex,\n",
    "}\n",
    "plt.rcParams.update(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='numba'></a>\n",
    "## 0. Make sure Numba is correctly installed\n",
    "Instructions available at: http://numba.pydata.org/numba-doc/latest/user/installing.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(numba.__version__) # Your version of Numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def go_fast(a): # Function is compiled to machine code when called the first time\n",
    "    trace = 0.0\n",
    "    # assuming square input matrix\n",
    "    for i in range(a.shape[0]):   # Numba likes loops\n",
    "        trace += np.tanh(a[i, i]) # Numba likes NumPy functions\n",
    "    return a + trace              # Numba likes NumPy broadcasting\n",
    "\n",
    "x = np.arange(100).reshape(10, 10)\n",
    "x = x.astype(float)\n",
    "go_fast(x) # Running a first time the function to compile it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.tanh(0)+ np.tanh(10) + np.tanh(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Iteration time of the Numba-compiled function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit go_fast(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Numba-compiled function have a special `.py_func` attribute which is the original uncompiled Python function.  Iteration time of the uncompiled function should be greater."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit go_fast.py_func(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='loss'></a>\n",
    "## 1. Loss functions, gradients and step-sizes\n",
    "\n",
    "\n",
    "We want to minimize\n",
    "$$\n",
    "\\frac 1n \\sum_{i=1}^n \\ell(a_i^\\top x, b_i) + \\frac \\lambda 2 \\|x\\|_2^2\n",
    "$$\n",
    "where\n",
    "- $\\ell(z, b) = \\frac 12 (b - z)^2$ (least-squares regression)\n",
    "- $\\ell(z, b) = \\log(1 + \\exp(-bz))$ (logistic regression).\n",
    "\n",
    "We write it as a minimization problem of the form\n",
    "$$\n",
    "\\frac 1n \\sum_{i=1}^n f_i(x)\n",
    "$$\n",
    "where\n",
    "$$\n",
    "f_i(x) = \\ell(a_i^\\top x, b_i) + \\frac \\lambda 2 \\|x\\|_2^2.\n",
    "$$\n",
    "\n",
    "For each cases, the gradients are\n",
    "$$\n",
    "\\nabla f_i(x) = (a_i^\\top x - b_i) a_i + \\lambda x\n",
    "$$\n",
    "and\n",
    "$$\n",
    "\\nabla f_i(x) = - \\frac{b_i}{1 + \\exp(b_i a_i^\\top x)} a_i + \\lambda x.\n",
    "$$\n",
    "\n",
    "Denote by $L$ (resp. $L_i$) the Lipschitz constant of $f$ (resp. $f_i$) and $\\mathbf A^\\top = [a_1, \\ldots, a_n].$\n",
    "One can easily see (using $\\|\\cdot\\|_{2}$ for the matrix spectrale norm) that for least-squares regression\n",
    "$$\n",
    "L = \\frac{ \\|\\mathbf A^\\top \\mathbf A \\|_{2}}{n} + \\lambda \\quad \\text{ and } L_i = \\| a_i \\|_2^2 + \\lambda\n",
    "$$\n",
    "while for logistic regression it is\n",
    "$$\n",
    "L = \\frac{ \\|\\mathbf A^\\top \\mathbf A \\|_{2}}{4 n} + \\lambda \\quad \\text{ and } L_i = \\frac 14 \\| a_i \\|_2^2 + \\lambda.\n",
    "$$\n",
    "For full-gradient methods, the theoretical step-size is $1 / L$, while for SAG and SVRG (see below) it can be taken as\n",
    "$1 / (\\max_{i=1,\\ldots,n} L_i)$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now introduce functions that will be used for the solvers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit    \n",
    "def grad_i_linreg(i, x, A, b, lbda):\n",
    "    \"\"\"Gradient with respect to a sample\"\"\"\n",
    "    a_i = A[i]\n",
    "    return (a_i.dot(x) - b[i]) * a_i + lbda * x\n",
    "\n",
    "\n",
    "@njit\n",
    "def grad_linreg(x, A, b, lbda):\n",
    "    \"\"\"Full gradient\"\"\"\n",
    "    g = np.zeros_like(x)\n",
    "    for i in range(n):\n",
    "        g += grad_i_linreg(i, x, A, b, lbda)\n",
    "    return g / n\n",
    "\n",
    "\n",
    "def loss_linreg(x, A, b, lbda):\n",
    "    return norm(A.dot(x) - b) ** 2 / (2. * n) + lbda * norm(x) ** 2 / 2.\n",
    "\n",
    "\n",
    "def lipschitz_linreg(A, b, lbda):\n",
    "    return norm(A, ord=2) ** 2 / n + lbda"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit    \n",
    "def grad_i_logreg(i, x, A, b, lbda):\n",
    "    \"\"\"Gradient with respect to a sample\"\"\"\n",
    "    a_i = A[i]\n",
    "    b_i = b[i]\n",
    "    return - a_i * b_i / (1. + np.exp(b_i * np.dot(a_i, x))) + lbda * x\n",
    "\n",
    "\n",
    "@njit\n",
    "def grad_logreg(x, A, b, lbda):\n",
    "    \"\"\"Full gradient\"\"\"\n",
    "    g = np.zeros_like(x)\n",
    "    for i in range(n):\n",
    "        g += grad_i_logreg(i, x, A, b, lbda)\n",
    "    return g / n\n",
    "\n",
    "\n",
    "def loss_logreg(x, A, b, lbda):\n",
    "    bAx = b * np.dot(A, x)\n",
    "    return np.mean(np.log(1. + np.exp(- bAx))) + lbda * norm(x) ** 2 / 2.\n",
    "\n",
    "\n",
    "def lipschitz_logreg(A, b, lbda):\n",
    "    return norm(A, ord=2) ** 2 / (4. * n) + lbda"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='data'></a>\n",
    "## 2. Generate a dataset\n",
    "\n",
    "We generate datasets for the least-squares and the logistic cases. First we define a function for the least-squares case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy.random import multivariate_normal, randn\n",
    "from scipy.linalg import toeplitz\n",
    "\n",
    "\n",
    "def simu_linreg(x, n, std=1., corr=0.5):\n",
    "    \"\"\"Simulation for the least-squares problem.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : ndarray, shape (d,)\n",
    "        The coefficients of the model\n",
    "    n : int\n",
    "        Sample size\n",
    "    std : float, default=1.\n",
    "        Standard-deviation of the noise\n",
    "    corr : float, default=0.5\n",
    "        Correlation of the features matrix\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    A : ndarray, shape (n, d)\n",
    "        The design matrix.\n",
    "    b : ndarray, shape (n,)\n",
    "        The targets.\n",
    "    \"\"\"\n",
    "    d = x.shape[0]\n",
    "    cov = toeplitz(corr ** np.arange(0, d))\n",
    "    A = multivariate_normal(np.zeros(d), cov, size=n)\n",
    "    noise = std * randn(n)\n",
    "    b = A.dot(x) + noise\n",
    "    return A, b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simu_logreg(x, n, std=1., corr=0.5):\n",
    "    \"\"\"Simulation for the logistic regression problem.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    x : ndarray, shape (d,)\n",
    "        The coefficients of the model\n",
    "    n : int\n",
    "        Sample size    \n",
    "    std : float, default=1.\n",
    "        Standard-deviation of the noise\n",
    "    corr : float, default=0.5\n",
    "        Correlation of the features matrix\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    A : ndarray, shape (n, d)\n",
    "        The design matrix.\n",
    "    b : ndarray, shape (n,)\n",
    "        The targets.\n",
    "    \"\"\"    \n",
    "    A, b = simu_linreg(x, n, std=1., corr=corr)\n",
    "    return A, np.sign(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = 50\n",
    "n = 10000\n",
    "idx = np.arange(d)\n",
    "\n",
    "# Ground truth coefficients of the model\n",
    "x_model_truth = (-1)**idx * np.exp(-idx / 10.)\n",
    "\n",
    "_A, _b = simu_linreg(x_model_truth, n, std=1., corr=0.1)\n",
    "#_A, _b = simu_logreg(x_model_truth, n, std=1., corr=0.7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.stem(x_model_truth);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numerically check loss and gradient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import check_grad\n",
    "\n",
    "lbda = 1. / n ** (0.5)\n",
    "\n",
    "A, b = simu_linreg(x_model_truth, n, std=1., corr=0.1)\n",
    "# Check that the gradient and the loss numerically match\n",
    "check_grad(loss_linreg, grad_linreg, np.random.randn(d), A, b, lbda)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lbda = 1. / n ** (0.5)\n",
    "\n",
    "A, b = simu_logreg(x_model_truth, n, std=1., corr=0.1)\n",
    "# Check that the gradient and the loss numerically match\n",
    "check_grad(loss_logreg, grad_logreg, np.random.randn(d), A, b, lbda)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choice of the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A, b = simu_linreg(x_model_truth, n, std=1., corr=0.9)\n",
    "loss = loss_linreg\n",
    "grad = grad_linreg\n",
    "grad_i = grad_i_linreg\n",
    "lipschitz_constant = lipschitz_linreg\n",
    "lbda = 1. / n ** (0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the theoretical step-size for gradient descent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "step = 1. / lipschitz_constant(A, b, lbda)\n",
    "\n",
    "print(\"step = %s\" % step)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get a very precise minimum to compute distances to minimum\n",
    "\n",
    "Here we use L-BFGS that you have not seen yet. Just take the estimated `x_min` and move on.\n",
    "You will see L-BFGS later in the course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import fmin_l_bfgs_b\n",
    "\n",
    "x_init = np.zeros(d)\n",
    "x_min, f_min, _ = fmin_l_bfgs_b(loss, x_init, grad,\n",
    "                                args=(A, b, lbda), pgtol=1e-30, factr=1e-30)\n",
    "\n",
    "\n",
    "print(f_min)\n",
    "print(norm(grad_linreg(x_min, A, b, lbda)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='batch'></a> \n",
    "\n",
    "## 3. Deterministic/Batch methods (GD, AGD)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a class to monitor iterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class monitor:\n",
    "    def __init__(self, algo, loss, x_min, args=()):\n",
    "        self.x_min = x_min\n",
    "        self.algo = algo\n",
    "        self.loss = loss\n",
    "        self.args = args\n",
    "        self.f_min = loss(x_min, *args)\n",
    "    \n",
    "    def run(self, *algo_args, **algo_kwargs):\n",
    "        t0 = time()\n",
    "        _, x_list = self.algo(*algo_args, **algo_kwargs)\n",
    "        self.total_time = time() - t0\n",
    "        self.x_list = x_list\n",
    "        self.err = [norm(x - self.x_min) for x in x_list]\n",
    "        self.obj = [self.loss(x, *self.args) - self.f_min for x in x_list]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of iterations\n",
    "n_iter = 200"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gradient descent (GD)\n",
    "\n",
    "We recall that an iteration of batch gradient writes\n",
    "\n",
    "$$\n",
    "x_{k+1} \\gets x_k - \\eta \\nabla f(x_k)\n",
    "$$\n",
    "\n",
    "where $\\eta$ is the step-size (that can be chosen in theory as $\\eta = 1 / L$, with $L$ the Lipshitz constant of $\\nabla f$, see above)\n",
    "\n",
    "*QUESTION*:\n",
    "- Fill in the iteration of the GD solver in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def gd(x_init, grad, n_iter=100, step=1., store_every=1, args=()):\n",
    "    \"\"\"Gradient descent algorithm.\"\"\"\n",
    "    x = x_init.copy()\n",
    "    x_list = []\n",
    "    for i in range(n_iter):\n",
    "        ### TODO\n",
    "        if i % store_every == 0:\n",
    "            x_list.append(x.copy())\n",
    "    return x, x_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "step = 1. / lipschitz_linreg(A, b, lbda)\n",
    "x_init = np.zeros(d)\n",
    "monitor_gd = monitor(gd, loss, x_min, (A, b ,lbda))\n",
    "monitor_gd.run(x_init, grad, n_iter, step, args=(A, b, lbda))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Accelerated Gradient Descent (AGD)\n",
    "\n",
    "We recall that an iteration of AGD (see FISTA) writes:\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "x_{k+1} &\\gets y_k - \\eta \\nabla f(y_k) \\\\\n",
    "t_{k+1} &\\gets \\frac{1 + \\sqrt{1 + 4 t_k^2}}{2} \\\\\n",
    "y_{k+1} &\\gets x_{k+1} + \\frac{t_k-1}{t_{k+1}} (x_{k+1} - x_k)\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "where $\\eta$ is the step-size (that can be chosen in theory as $\\eta = 1 / L$, with $L$ the Lipschitz constant of $\\nabla f$, see above)\n",
    "\n",
    "*QUESTION*:\n",
    "- Fill in the iteration of the AGD solver in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def agd(x_init, grad, n_iter=100, step=1., args=(), store_every=1):\n",
    "    \"\"\"Accelerated Gradient Descent algorithm.\"\"\"\n",
    "    x = x_init.copy()\n",
    "    y = x_init.copy()\n",
    "    t = 1.\n",
    "    x_list = []\n",
    "    for i in range(n_iter):\n",
    "        if i % store_every == 0:\n",
    "            x_list.append(x.copy())\n",
    "        ### TODO\n",
    "    return x, x_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "step = 1. / lipschitz_linreg(A, b, lbda)\n",
    "x_init = np.zeros(d)\n",
    "monitor_agd = monitor(agd, loss, x_min, (A, b ,lbda))\n",
    "monitor_agd.run(x_init, grad, n_iter, step, args=(A, b, lbda))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A first numerical comparison of deterministic solvers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, define some plotting functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_epochs(monitors, solvers):\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    for monit in monitors:\n",
    "        plt.semilogy(monit.obj, lw=2)\n",
    "        plt.title(\"Loss\")\n",
    "        plt.xlabel(\"Epoch\")\n",
    "        plt.ylabel(\"objective\")\n",
    "\n",
    "    plt.legend(solvers)\n",
    "\n",
    "    plt.subplot(1, 2, 2)\n",
    "\n",
    "    for monit in monitors:\n",
    "        plt.semilogy(monit.err, lw=2)\n",
    "        plt.title(\"Distance to optimum\")\n",
    "        plt.xlabel(\"Epoch\")\n",
    "        plt.ylabel(\"$\\|x_k - x^*\\|_2$\")\n",
    "\n",
    "    plt.legend(solvers)\n",
    "    \n",
    "\n",
    "def plot_time(monitors, solvers, xscale='linear', yscale='log'):\n",
    "    for monit in monitors:\n",
    "        objs = monit.obj\n",
    "        plt.semilogy(np.linspace(0, monit.total_time, len(objs)), objs, lw=2)\n",
    "        plt.title(\"Loss\")\n",
    "        plt.xlabel(\"Timing\")\n",
    "        plt.ylabel(\"$f(x_k) - f(x^*)$\")\n",
    "        plt.yscale(yscale)\n",
    "        plt.xscale(xscale)\n",
    "\n",
    "    plt.legend(solvers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monitors = [monitor_gd, monitor_agd]\n",
    "\n",
    "solvers = [\"GD\", \"AGD\"]\n",
    "\n",
    "plot_epochs(monitors, solvers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_time(monitors, solvers, xscale='log')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First conclusions\n",
    "\n",
    "*QUESTIONS*:\n",
    "\n",
    "- Give some first conclusions about the batch solver studied here\n",
    "- What do you observe about AGD? is it suprising ?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='stoc'></a> \n",
    "## 4. Stochastic methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_iter = 50\n",
    "\n",
    "# generate indices of random samples\n",
    "random_indices = np.random.randint(0, n, n * n_iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SGD\n",
    "\n",
    "We recall that an iteration of SGD writes\n",
    "\n",
    "- Pick $i$ uniformly at random in $\\{1, \\ldots, n\\}$\n",
    "- Apply\n",
    "$$\n",
    "x_{t+1} \\gets x_t - \\frac{\\eta_0}{\\sqrt{t+1}} \\nabla f_i(x_t)\n",
    "$$\n",
    "\n",
    "where $\\eta_0$ is a step-size to be tuned by hand.\n",
    "\n",
    "*QUESTION*:\n",
    "- Fill in the iteration of the SGD solver in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def sgd(x_init, random_indices, grad_i, n_iter=100, step=1., store_every=n, args=()):\n",
    "    \"\"\"Stochastic gradient descent algorithm.\"\"\"\n",
    "    x = x_init.copy()\n",
    "    x_list = []\n",
    "    for idx in range(n_iter):\n",
    "        i = random_indices[idx]\n",
    "        \n",
    "        ### TODO\n",
    "        \n",
    "        # Update metrics after each iteration.\n",
    "        if idx % store_every == 0:\n",
    "            x_list.append(x.copy())\n",
    "    return x, x_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "step0 = 1e-1\n",
    "x_init = np.zeros(d)\n",
    "\n",
    "monitor_sgd = monitor(sgd, loss, x_min, (A, b ,lbda))\n",
    "monitor_sgd.run(x_init, random_indices, grad_i, n_iter * n, step0, args=(A, b, lbda))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SAG\n",
    "\n",
    "We recall that an iteration of SAG writes\n",
    "\n",
    "For $t=1, \\ldots, $ until convergence\n",
    "\n",
    "1. Pick $i_t$ uniformly at random in $\\{1, \\ldots, n\\}$\n",
    "\n",
    "2. Update the average of gradients\n",
    "$$\n",
    "G_t \\gets \\frac 1n \\sum_{i=1}^n g_i^t\n",
    "$$\n",
    "where \n",
    "$$\n",
    "g_i^t =\n",
    "\\begin{cases}\n",
    "    \\nabla f_{i}(x_t) &\\text{ if } i = i_t \\\\\n",
    "    g_i^{t-1} & \\text{ otherwise.}\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "3. Apply the step \n",
    "$$x_{t+1} \\gets x_t - \\eta G_t$$\n",
    "where $\\eta$ is the step-size (see code below).\n",
    "\n",
    "*QUESTION*:\n",
    "- Fill in the iteration of the SAG solver in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def sag(x_init, random_indices, grad_i, n_iter=100, step=1., store_every=n, args=()):\n",
    "    \"\"\"Stochastic average gradient algorithm.\"\"\"\n",
    "    x = x_init.copy()\n",
    "    # Old gradients\n",
    "    gradient_memory = np.zeros((n, d))\n",
    "    averaged_gradient = np.zeros(d)\n",
    "    x_list = []\n",
    "    for idx in range(n_iter):\n",
    "        i = random_indices[idx]\n",
    "        \n",
    "        ### TODO\n",
    "        \n",
    "        # Update metrics after each iteration.\n",
    "        if idx % store_every == 0:\n",
    "            x_list.append(x.copy())\n",
    "    return x, x_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_squared_sum = np.max(np.sum(A ** 2, axis=1))\n",
    "step = 1.0 / (max_squared_sum + lbda)\n",
    "\n",
    "x_init = np.zeros(d)\n",
    "monitor_sag = monitor(sag, loss, x_min, (A, b ,lbda))\n",
    "monitor_sag.run(x_init, random_indices, grad_i, n_iter * n, step, args=(A, b, lbda))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SVRG\n",
    "\n",
    "We recall that an iteration of SVRG writes\n",
    "\n",
    "For $k=1, \\ldots, $ until convergence\n",
    "\n",
    "1. Set $\\tilde x \\gets \\tilde x^{(k)}$ and $x_1^{(k)} \\gets \\tilde x$\n",
    "2. Compute $\\mu_k \\gets \\nabla f(\\tilde x)$\n",
    "3. For $t=1, \\ldots, n$\n",
    "    4. Pick $i$ uniformly at random in $\\{1, \\ldots, n\\}$\n",
    "    5. Apply the step \n",
    "$$\n",
    "x_{t+1}^{(k)} \\gets x_t^{(k)} - \\eta \\big(\\nabla f_{i}(x_t^{(k)}) - \\nabla f_{i}(\\tilde x) + \\mu_k \\big) \n",
    "$$\n",
    "\n",
    "6. Set $\\tilde x^{(k+1)} \\gets x_{n+1}^{(k)}$\n",
    "\n",
    "where $\\eta$ is the step-size (see code below).\n",
    "\n",
    "*QUESTION*:\n",
    "- Fill in the iteration of the SVRG solver in the cell below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@njit\n",
    "def svrg(x_init, random_indices, grad, grad_i, n_iter=100, step=1., store_every=n, args=()):\n",
    "    \"\"\"Stochastic variance reduction gradient algorithm.\"\"\"\n",
    "    x = x_init.copy()\n",
    "    x_old = x.copy()\n",
    "    x_list = []\n",
    "    for idx in range(n_iter):\n",
    "        \n",
    "        ### TODO     \n",
    "        \n",
    "        # Update metrics after each iteration.\n",
    "        if idx % store_every == 0:\n",
    "            x_list.append(x.copy())\n",
    "    return x, x_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_init = np.zeros(d)\n",
    "monitor_svrg = monitor(svrg, loss, x_min, (A, b ,lbda))\n",
    "monitor_svrg.run(x_init, random_indices, grad, grad_i, n_iter * n, step, args=(A, b, lbda))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "monitors = [monitor_sgd, monitor_sag, monitor_svrg]\n",
    "\n",
    "solvers = [\"SGD\", \"SAG\", \"SVRG\"]\n",
    "\n",
    "plot_epochs(monitors, solvers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_time(monitors, solvers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='comp'></a> \n",
    "## 5. Numerical comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "monitors = [monitor_gd, monitor_agd,\n",
    "            monitor_sgd, monitor_sag, monitor_svrg]\n",
    "\n",
    "solvers = [\"GD\", \"AGD\", \"SGD\", \"SAG\", \"SVRG\"]\n",
    "\n",
    "plot_epochs(monitors, solvers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_time(monitors, solvers, xscale='log')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='conc'></a>\n",
    "## 6. Conclusion\n",
    "\n",
    "*QUESTIONS*:\n",
    "- Compare and comment your results\n",
    "- Change the value of the ridge regularization (the ``lbda`` parameter) to low ridge $\\lambda = 1 / n$ and high ridge regularization $\\lambda = 1$ and compare your results. Comment.\n",
    "- Play also with the level of correlation between features (parameter ``corr`` above), and compare results with low and high correlation.\n",
    "- Conclude"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
