{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "73fdabd6",
   "metadata": {},
   "source": [
    "# 18.065 Problem Set 4\n",
    "\n",
    "Due Friday 4/7 at 1pm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb69bf7",
   "metadata": {},
   "source": [
    "## Problem 1\n",
    "\n",
    "A \"convex set\" $S$ is one such that for any two points $x,y$ in $S$, the straight line connecting them is also in $S$.  That is, $\\alpha x + (1-\\alpha)y \\in S$ for all $\\alpha \\in [0,1]$.\n",
    "\n",
    "**(a)** If $f_i(x)$ is a [convex function](https://en.wikipedia.org/wiki/Convex_function) (as defined in e.g. lecture 16), explain why the constraint $f_i(x) \\le 0$ defines a convex set of feasible points.\n",
    "\n",
    "**(b)** Explain why the intersection of two convex sets is a convex set.  (Hence, the feasible set for many convex constraints is a convex set.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "188e8b83",
   "metadata": {},
   "source": [
    "## Problem 2\n",
    "\n",
    "Suppose we are solving the $\\ell^1$-regularized least-square problem:\n",
    "$$\n",
    "\\min_{x\\in \\mathbb{R}^n} \\left( \\Vert b - Ax \\Vert_2^2 + \\lambda \\Vert x \\Vert_1 \\right)\n",
    "$$\n",
    "where $\\lambda > 0$ is some regularization parameter and $ \\Vert x \\Vert_1 = \\sum_i |x_i|$ is the $\\ell^1$ norm.\n",
    "\n",
    "Similar to the examples in 16, show that this can be converted into a an equivalent [quadratic programming (QP)](https://en.wikipedia.org/wiki/Quadratic_programming) problem — a convex quadratic objective with affine constraints — by introducing one or more dummy variables (an \"epigraph\" formulation).   This replaces the non-differentiable $\\ell^1$ problem with an equivalent *differentiable* problem with *differentiable constraints*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73427817",
   "metadata": {},
   "source": [
    "## Problem 3\n",
    "\n",
    "Let\n",
    "$$\n",
    "A(p) = A_0 + \\underbrace{\\begin{pmatrix} p_1 & & & \\\\ & p_2 & & \\\\ & & \\ddots & \\\\ & & & p_m \\end{pmatrix}}_{\\text{diagm(p)}}\n",
    "$$\n",
    "where $A_0$ is some $m \\times m$ matrix and $p \\in \\mathbb{R}^m$ are $m$ parameters.\n",
    "\n",
    "Now, suppose we compute\n",
    "$$\n",
    "f(p) = g(\\underbrace{A(p)^{-1} b}_{x(p)})\n",
    "$$\n",
    "where $g(x) = x^T G x$, $b \\in \\mathbb{R}^m$ is some vector, and $G = G^T$ is some symmetric $m \\times m$ matrix.\n",
    "\n",
    "**(a)** What is $\\nabla f$?  Explain how you can compute *all m* components of $\\nabla f$ by solving only *two* $m \\times m$ systems followed by $\\Theta(m)$ additional work in *total*.\n",
    "\n",
    "**(b)** Implement your solution from (a) by filling in the function `∇f(p)` below, and check that it correctly predicts $df = f(p + dp) - f(p)$ (*approximately*) for a random small $dp$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac419822",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "\n",
    "# some A₀, b, G for m=5\n",
    "m = 5\n",
    "A₀ = randn(m,m)\n",
    "b = randn(m)\n",
    "G = randn(m,m); G = G' + G\n",
    "\n",
    "# our functions\n",
    "g(x) = x' * G * x\n",
    "f(p) = g((A₀ + Diagonal(p)) \\ b)\n",
    "\n",
    "p = randn(m) # some random parameters\n",
    "f(p) # make sure it gives a number out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12609b9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# your solution to (b):\n",
    "\n",
    "function ∇f(p)\n",
    "    ????\n",
    "end\n",
    "\n",
    "dp = randn(m) * 1e-8 # a random small dp\n",
    "df = f(p + dp) - f(p)\n",
    "\n",
    "# check: ∇f approximately predicts df\n",
    "# ???"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84c4326e",
   "metadata": {},
   "source": [
    "## Problem 4\n",
    "\n",
    "In class, we considered steepest descent for $f(x) = \\kappa x_1^2 + x_2^2$ in $\\mathbb{R}^2$, and argued that for an arbitrary $x = [x_1,x_2]$ starting point and $\\kappa \\gg 1$, the steepest-descent step $x \\leftarrow x - sz$ is *approximately* $sz \\approx [x_1, x_2/\\kappa]$.\n",
    "\n",
    "**(a)** If $sz = [x_1, x_2/\\kappa]$ exactly, then on the next step we would have the new $x \\leftarrow x - sz = [0, (1-\\frac{1}{\\kappa}) x_2)]$.   However, explain why a more careful calculation shows that the new $x - sz \\approx [O(1/\\kappa^2), x_2-\\frac{1}{\\kappa} x_2 + O(1/\\kappa^3)]$, i.e. the first component is proportional to $1/\\kappa^2$ to leading order in $1/\\kappa$.\n",
    "\n",
    "**(b)** If you start with an $x = [\\#/\\kappa, x_2]$, i.e. where the first component is proportional to $1/\\kappa$ and $\\#$ is some number of the same order of magnitude as $x_2$, show that after one steepest-descent step (for $\\kappa \\gg 1$) the $x_1$ component is *still* roughly order $1/\\kappa$ but of the opposite sign, and $x_2$ again subtracts a term roughly proportional to $x_2/\\kappa$.\n",
    "\n",
    "**(c)** Implement this steepest-descent process numerically for $\\kappa = 100$ and a starting point $x_1 = 0.01234567$ and $x_2 = 0.8910$.  Plot $100x_1$ and $x_2$ for 100 iterations of steepest descent.  A more careful analysis would show a convergence proportional to $\\left( \\frac{\\kappa - 1}{\\kappa + 1} \\right)^k$, where $k$ is the iteration number, following equation (4) in the Strang book — include this function for comparison on your plot."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e9798d4",
   "metadata": {},
   "source": [
    "## Problem 5\n",
    "\n",
    "In the Adam algorithm for stochastic optimization, we keep a \"running average\" (\"1st moment\") $m^{(k)}$ of the estimated gradients.  If $\\tilde{\\nabla}f^{(k)}$ is our estimated (mini-batch) gradient at $x^{(k)}$ (the k-th step, for $k=1,2,3,\\ldots$), then the running average is computed as:\n",
    "$$\n",
    "m^{(k)} = \\beta_1 m^{(k-1)} + (1-\\beta_1) \\tilde{\\nabla}f^{(k)} \\, ,\n",
    "$$\n",
    "where $0 \\le \\beta_1 < 1$ (typically $\\beta_1 = 0.9$) is a given weight, and we define $m^{(0)} = \\vec{0}$.\n",
    "\n",
    "However, before we use the gradient to take a step, we \"de-bias\" it to:\n",
    "$$\n",
    "\\hat{m}^{(k)} = \\frac{m^{(k)}}{1 - \\beta_1^k}\n",
    "$$\n",
    "**Explain** why $\\hat{m}^{(k)}$ is exactly a weighted average of the estimates $\\tilde{\\nabla}f^{(j)}$ for $j = 1,\\ldots,k$ — that is, explain why dividing by $1 - \\beta_1^k$ exactly corresponds to dividing by the sum of the weights of each $\\tilde{\\nabla}f^{(j)}$ term.  (The formulas for the [sum of a geometric series](https://en.wikipedia.org/wiki/Geometric_series) will be helpful.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7f174a6",
   "metadata": {},
   "source": [
    "## Problem 6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3e6a871",
   "metadata": {},
   "source": [
    "The Julia notebook from class (lecture 21) implemented the Adam algorithm for stochastic optimization.  This code is reproduced below.\n",
    "\n",
    "A more recent algorithm, that has better convergence guarantees than Adam, is called [YOGI (Zaheer et al., 2018)](https://papers.nips.cc/paper_files/paper/2018/hash/90365351ccc7437a1309dc64e4db32a3-Abstract.html).   Make a copy of the code below and **modify it to implement the YOGI algorithm** as described in algorithm 2 of the linked paper.  (This algorithm omits the \"de-biasing\" step of problem 5, but you should include that in your implementation; the de-biasing is identical to that of Adam.)   **Plot the convergence** of YOGI for this problem, with the same hyper-parameters, along with the original Adam algorithms (using different colors superimposed on a single plot)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2181180f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 93664 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 1000x400 with 2 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "PyObject Text(0.5, 1.0, '|error| in x')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using LinearAlgebra, PyPlot\n",
    "\n",
    "# typical Adam hyper-parameters\n",
    "ϵ = 1e-8\n",
    "β₁ = 0.9\n",
    "β₂ = 0.999\n",
    "\n",
    "# more parameters\n",
    "log₁₀α = -1\n",
    "α = exp10(log₁₀α) # approx. maximum stepsize |Δx|\n",
    "M = 1 # mini-batch size\n",
    "iterations = 300\n",
    "\n",
    "x̂₀ = [1,2] # ground-truth model\n",
    "x = [0.0,0.0]\n",
    "X = [x]\n",
    "ΔX = [norm(x - x̂₀)]\n",
    "m = zero(x) # initial 1st-moment (mean gradient)\n",
    "v = zero(x) # initial 2nd-moment (mean gradient.^2)\n",
    "for t = 1:iterations\n",
    "    # generate a mini-batch of size M:\n",
    "    ã = randn(M)\n",
    "    b̃ = @. 1 + 2ã + randn()\n",
    "    Ã = [one.(ã) ã]\n",
    "\n",
    "    # approximate gradient from this mini-batch\n",
    "    ∇f̃ = (2/M)*(Ã'*(Ã*x - b̃))\n",
    "\n",
    "    # Adam update\n",
    "    m = β₁*m + (1-β₁)*∇f̃\n",
    "    v = β₂*v + (1-β₂)*(∇f̃.^2)\n",
    "    m̂ = m / (1 - β₁^t) # de-bias: normalize by total weights\n",
    "    v̂ = v / (1 - β₂^t) # de-bias: normalize by total weights\n",
    "    x = @. x - α * m̂ / (√v̂ + ϵ)\n",
    "\n",
    "    push!(X, x)\n",
    "    push!(ΔX, norm(x - x̂₀))\n",
    "end\n",
    "\n",
    "# objective function for contour plot:\n",
    "a = randn(1000)\n",
    "b = @. 1 + 2a + randn()\n",
    "A = [one.(a) a]\n",
    "f(x) = norm(b-A*x)^2 / length(b)\n",
    "f(x1,x2) = f([x1,x2])\n",
    "x1 = range(-2,3, length=100)\n",
    "x2 = range(0,4,length=100)\n",
    "fvals = f.(x1',x2)\n",
    "\n",
    "# plot progress\n",
    "\n",
    "figure(figsize=(10,4))\n",
    "\n",
    "subplot(1,2,1)\n",
    "plt.subplots_adjust(wspace=0.4)\n",
    "contour(x1', x2, fvals)\n",
    "plot(x̂₀[1], x̂₀[2], \"r*\")\n",
    "plot(getindex.(X,1), getindex.(X,2), \"ko-\")\n",
    "xlabel(L\"x_1\")\n",
    "ylabel(L\"x_2\")\n",
    "title(\"gradient-descent iterates x\")\n",
    "\n",
    "subplot(1,2,2)\n",
    "semilogy(ΔX, \"ko-\")\n",
    "ylim(1e-2, 2)\n",
    "xlabel(\"iteration\")\n",
    "ylabel(L\"\\Vert x - \\hat{x} \\Vert\")\n",
    "title(\"|error| in x\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "167a6e22",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.8.3",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
