{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "701e5077",
   "metadata": {},
   "source": [
    "# 18.065 Problem Set 5 Solutions\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74f502b8",
   "metadata": {},
   "source": [
    "## Problem 1 (5+6 points)\n",
    "\n",
    "Consider the following optimization problem:\n",
    "$$\n",
    "\\min_{x \\in \\mathbb{R}^2} x_1 \\\\\n",
    "\\mbox{  subject to } x_2 \\le x_1^3 \\mbox{ and } x_2 \\ge 0 \\, .\n",
    "$$\n",
    "\n",
    "**(a)** Draw a sketch of the feasible set in the $x_1, x_2$ plane and indicate the optimum $x_*$.\n",
    "\n",
    "**(b)** Show that the optimum $x_*$ does *not* satisfy the KKT conditions, but explain why this is possible because the LICQ conditions are violated (see the last slide of lecture 22).\n",
    "\n",
    "(Most problems have local minima that satisfy KKT, but you can see from the picture in (a) that this is a weird case!)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb563847",
   "metadata": {},
   "source": [
    "### Solution:\n",
    "\n",
    "**(a)** We will use PyPlot to draw our feasible set as a light-blue shaded region, which lies below the cubic $x_2 = x_1^3$ and above the $x_2=0$ axis, and hence is a **cusp-like** region that **only includes** $x_1 \\ge 0$.\n",
    "\n",
    "We also the primal optimum $x_*$ as a red star, which lies at $\\boxed{x_1 = 0, x_2 = 0}$, since this is the smallest $x_1$ point in the feasible region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b77b6846",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30312 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 640x480 with 1 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "using PyPlot\n",
    "\n",
    "# plot constraints\n",
    "x₁ = range(-2,2,length=200)\n",
    "plot(x₁, x₁.^3, \"b-\", linewidth=3)\n",
    "text(1, 1.5, \"x₂ ≤ x₁³\", color=\"b\", rotation=45)\n",
    "plot(x₁, zero(x₁), \"c-\", linewidth=3)\n",
    "text(1.5, 0.4, \"x₂ ≥ 0\", color=\"c\")\n",
    "\n",
    "# plot feasible set\n",
    "x₁⁺ = range(0,2,length=100)\n",
    "fill_between(x₁⁺, zero(x₁⁺), x₁⁺.^3, color=(0.8,0.8,1))\n",
    "text(1.4, 1.8, \"feasible\")\n",
    "\n",
    "# plot optimum\n",
    "plot([0],[0],\"r*\",markersize=15)\n",
    "text(0.1,0.4,L\"x_*\",color=\"r\")\n",
    "\n",
    "title(\"problem 1(a): feasible set and primal optimum\")\n",
    "\n",
    "# center the axes:\n",
    "ax = gca()\n",
    "ax.spines[\"left\"].set_position(\"zero\")\n",
    "ax.spines[\"right\"].set_color(\"none\")\n",
    "ax.spines[\"bottom\"].set_position(\"zero\")\n",
    "ax.spines[\"top\"].set_color(\"none\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "376a7e9c",
   "metadata": {},
   "source": [
    "**(b)** Our objective is $f_0(x) = x_1$, and our two constraints $f_i(x) \\le 0$ are $f_1(x) = x_2 - x_1^3$ and $f_2(x) = -x_2$.  Therefore, at the optimum $x_* = (0,0)$, their gradients are:\n",
    "$$\n",
    "\\left. \\nabla f_0 \\right|_{x_*} = (1,0), \\qquad\n",
    "\\left. \\nabla f_1 \\right|_{x_*} = (0,1), \\qquad\n",
    "\\left. \\nabla f_2 \\right|_{x_*} = (0,-1) \\, .\n",
    "$$\n",
    "These are depicted graphically on the diagram below.\n",
    "\n",
    "But these **cannot possibly** satisfy the KKT condition $\\boxed{\\nabla L = 0}$ (for the Lagrangian $L = f_0 + \\lamba_1 f_1 + \\lambda_2 f_2$), for **any** Lagrange multipliers $\\lambda_1, \\lambda_2$, because $\\nabla f_0$ is **linearly independent of** (*orthogonal* to!) $\\nabla f_1$ and $\\nabla f_2$!\n",
    "\n",
    "It is possible to violate the KKT conditions at $x_*$ in this problem because the optimum **violates LICQ**: the constraint gradients $\\nabla f_1$ and $\\nabla f_2$ are linearly dependent (anti-parallel!) at $x_*$.\n",
    "\n",
    "This happens because the feasible set **has a cusp** (where two constraints become parallel) right at the primal optimum.   Obviously, this is an odd case that doesn't arise too often in practice — most optimization problems typically satisfy LICQ and KKT conditions at any local optimum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f2bfd986",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37360 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 640x480 with 1 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot constraints\n",
    "x₁ = range(-2,2,length=200)\n",
    "plot(x₁, x₁.^3, \"b-\", linewidth=3)\n",
    "text(1, 1.5, \"x₂ ≤ x₁³\", color=\"b\", rotation=45)\n",
    "plot(x₁, zero(x₁), \"c-\", linewidth=3)\n",
    "text(1.5, 0.4, \"x₂ ≥ 0\", color=\"c\")\n",
    "\n",
    "# plot feasible set\n",
    "x₁⁺ = range(0,2,length=100)\n",
    "fill_between(x₁⁺, zero(x₁⁺), x₁⁺.^3, color=(0.8,0.8,1))\n",
    "text(1.4, 1.8, \"feasible\")\n",
    "\n",
    "# plot optimum\n",
    "plot([0],[0],\"r*\",markersize=15)\n",
    "text(0.1,0.4,L\"x_*\",color=\"r\")\n",
    "\n",
    "title(L\"problem 1(b): gradients $\\nabla f_0, \\nabla f_1, \\nabla f_2$ at $x_*$\")\n",
    "\n",
    "\n",
    "# gradients\n",
    "arrow(0,0,0,5,color=\"b\", width=0.04, head_width=0.2, head_length=0.8)\n",
    "text(0.1, 5, L\"\\nabla f_1 = \\nabla (x_2 - x_1^3)\", color=\"b\", fontsize=14)\n",
    "\n",
    "arrow(0,0,0,-5,color=\"c\", width=0.04, head_width=0.2, head_length=0.8)\n",
    "text(0.1, -5, L\"\\nabla f_2 = \\nabla (-x_2)\", color=\"c\", fontsize=14)\n",
    "\n",
    "arrow(0,0,1,0,color=\"r\", width=0.2, head_width=0.8, head_length=0.2, zorder=10)\n",
    "text(1, -2, L\"\\nabla f_0 = \\nabla (x_1)\", color=\"r\", fontsize=14)\n",
    "\n",
    "# center the axes:\n",
    "ax = gca()\n",
    "ax.spines[\"left\"].set_position(\"zero\")\n",
    "ax.spines[\"right\"].set_color(\"none\")\n",
    "ax.spines[\"bottom\"].set_position(\"zero\")\n",
    "ax.spines[\"top\"].set_color(\"none\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a49c4b1",
   "metadata": {},
   "source": [
    "## Problem 2 (6+6+6 points)\n",
    "\n",
    "Consider the convex problem:\n",
    "$$\n",
    "\\min_{x\\in \\mathbb{R}^n} \\Vert b - Ax \\Vert_2^2 \\\\\n",
    "\\mbox{  subject to } \\Vert x \\Vert_2^2 \\le r^2\n",
    "$$\n",
    "for some $r > 0$, $m \\times n$ matrix $A$ (of rank $n$), and $b \\in \\mathbb{R}^m$ — that is, least-squares optimization with the solution constrained to lie inside a sphere of radius $r$.\n",
    "\n",
    "**(a)** What is the Lagrange dual function $g(\\lambda)$?   (You can give a closed-form expression.  Hint: review Tikhonov-regularized least-squares.)   Define a corresponding Julia function `g(λ; r=1.0)` for the sample parameters given below (this syntax defines an optional keyword argument `r` that defaults to $r=1$).  Make a plot of $g(\\lambda)$ for $r=1$ and $r=0.5$ for $\\lambda \\ge 0$ to verify that it looks concave with a single maximum.\n",
    "\n",
    "**(b)** If the unconstrained least-square solution $\\hat{x} = (A^T A)^{-1} A^T b$ satisfies $\\Vert \\hat{x} \\Vert_2 < r$, then what must be true of the derivative $g'(0)$?  What if $\\Vert \\hat{x} \\Vert_2 > r$?\n",
    "\n",
    "Check in Julia that $g'(0)$ matches your expectations by computing the derivative using automatic differentiation:\n",
    "```jl\n",
    "using ForwardDiff\n",
    "dgdλ(λ; r=1.0) = ForwardDiff.derivative(λ -> g(λ; r=r), λ)\n",
    "```\n",
    "and evaluating it at `dgdλ(0; r=???)` for two values of `r`: one $r > \\Vert \\hat{x} \\Vert_2$ (so that the constraint is inactive) and one $r < \\Vert \\hat{x} \\Vert_2$ (so that the constraint is active).\n",
    "\n",
    "(In principle, you could take this derivative by hand using matrix calculus, but it's pretty error-prone.)\n",
    "\n",
    "**(c)** You can take the *second* derivative of $g(\\lambda)$ via AD by:\n",
    "```jl\n",
    "d²gdλ²(λ; r=0.5) = ForwardDiff.derivative(λ -> dgdλ(λ; r=0.5), λ)\n",
    "```\n",
    "Use this to implement a Newton iteration to maximize $g(\\lambda)$ (for $\\lambda \\ge 0$) by finding a root of $g'(\\lambda)$, starting with an initial guess of $\\lambda=0$, for $r = 0.5$.  (It should converge in only a few iterations.  The solution should have $\\lambda > 0$ in this case because ...?)   To at least 8 significant digits, give the resulting dual optimum $\\lambda_*$ and the primal optimum $x_*$ (strong duality holds in this convex problem!), and check that $x_*$ is feasible."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c6ac1c",
   "metadata": {},
   "source": [
    "### Solution:\n",
    "\n",
    "**(a)** The dual function $g(\\lambda)$ is the found by minimizing the Lagrangian over $x$, i.e.\n",
    "$$\n",
    "g(\\lambda) = \\min_x \\underbrace{\\left[ \\Vert b - Ax \\Vert_2^2 + \\lambda (\\Vert x \\Vert_2^2 - r^2) \\right]}_{L(x,\\lambda)}\n",
    "$$\n",
    "where $\\Vert x \\Vert_2^2 - r^2 \\le 0$ is our inequality constraint.  But this is exactly in the form of Tikhonov-regularized least-squares (plus a constant $-\\lambda r^2$), so we know from class that the solution is\n",
    "$$\n",
    "\\hat{x}_\\lambda = (A^T A + \\lambda I)^{-1} (A^T b)\n",
    "$$\n",
    "and hence $g(\\lambda) = \\boxed{L(\\hat{x}_\\lambda, \\lambda)}$ in terms of the formulas above.   We can implement this as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1a89466c",
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "\n",
    "m = 5\n",
    "n = 4\n",
    "A = [ -9   2  -2   3\n",
    "      -5  -3   9   3\n",
    "      -1  -6   9  -2\n",
    "      -3  -4   5   4\n",
    "      -8   9  -6   4 ]\n",
    "b = [1,2,3,4,5];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e5a615e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "g (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function g(λ; r=1)\n",
    "    x̂_λ = (A'A + λ*I) \\ (A'b)\n",
    "    return norm(b - A*x̂_λ)^2 + λ*(norm(x̂_λ)^2 - r^2)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4001e6f4",
   "metadata": {},
   "source": [
    "Plotting it, we obtain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "521ec11c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28192 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 640x480 with 1 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(-5.0, 20.0)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using PyPlot\n",
    "λ = range(0,100, length=1000)\n",
    "plot(λ, g.(λ, r=0.5), \"r-\")\n",
    "plot(λ, g.(λ, r=1), \"b--\")\n",
    "title(L\"Problem 2(a): $g(\\lambda)$\")\n",
    "xlabel(L\"\\lambda\")\n",
    "ylabel(L\"g(\\lambda)\")\n",
    "legend([L\"r=0.5\", L\"r=1\"])\n",
    "ylim(-5, 20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "864bda6b",
   "metadata": {},
   "source": [
    "**(b)** If the unconstrained least-square solution satisfies the constraint, then the constraint is **inactive** and we expect that the corresponding Lagrange multiplier λ should be **zero** (corresponding to the \"complementary slackness\" KKT condition).  In the dual problem, this corresponds to $$g(\\lambda)$$ for $$\\lambda \\ge 0$$ having its maximum at $\\lambda = 0$, corresponding to $\\boxed{g'(0) < 0}$ for $\\Vert \\hat{x} \\Vert_2 < r$.\n",
    "\n",
    "Conversely, if $\\Vert \\hat{x} \\Vert_2 > r$, then the constraint is **active** and we expect a dual maximum for $\\lambda > 0$, which requires $\\boxed{g'(0) > 0}$ in that case.\n",
    "\n",
    "To get an example of this, let's check the norm of the unconstrained least-square solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b7facd32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6508599263498679"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norm(A \\ b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32751a1e",
   "metadata": {},
   "source": [
    "i.e. $\\Vert \\hat{x} \\Vert_2 = \\Vert (A^T A)^{-1} A^T b \\Vert_2 \\approx 0.651$, so the constraint should be inactive for any radius larger than that.\n",
    "\n",
    "For example, the constraint should be *inactive* for $r=1$ and *active* for $r = 0.5$, and we can see by inspection of the plot from part (a) that $g'(0)$ is $< 0$ and $> 0$ in these two cases, respectively.  Quantitatively, we can use AD to get the derivative as suggested:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e3a11067",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dgdλ(0; r = 1.0) = -0.5763813562718445\n",
      "dgdλ(0; r = 0.5) = 0.17361864372815555\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.17361864372815555"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using ForwardDiff\n",
    "dgdλ(λ; r=1.0) = ForwardDiff.derivative(λ -> g(λ; r=r), λ)\n",
    "\n",
    "@show dgdλ(0; r=1.0)\n",
    "@show dgdλ(0; r=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bf7db46",
   "metadata": {},
   "source": [
    "Hurray, the signs match our expectations!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41241ac1",
   "metadata": {},
   "source": [
    "**(c)** As suggested, we can get a second derivative by:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ca12a000",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "d²gdλ² (generic function with 1 method)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d²gdλ²(λ; r=0.5) = ForwardDiff.derivative(λ -> dgdλ(λ; r=0.5), λ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484a2d51",
   "metadata": {},
   "source": [
    "We can use this to implement a Newton iteration for the dual optimum.  A Newton iteration to find a maximum $g'(\\lambda) = 0$ is just\n",
    "$$\n",
    "\\lambda \\longleftarrow \\lambda - \\frac{g'(\\lambda)}{g''(\\lambda)}\n",
    "$$\n",
    "We can start at $\\lambda = 0$.  Of course, we first check whether $g'(0) > 0$ to be sure the constraint is active; otherwise, the dual optimum is simply $\\lambda = 0$.\n",
    "\n",
    "Given the dual optimum $\\lambda_*$, by strong duality we then have the corresponding primal optimum $x_* = \\hat{x}_{\\lambda_*}$ from part (a)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0d7df6ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dualopt (generic function with 1 method)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# find the dual optimum (maximizing g(λ))\n",
    "# and return the dual and primal optima λ,x\n",
    "function dualopt(r; rtol=10*eps(Float64), maxiter=100)\n",
    "    λ = zero(rtol)\n",
    "    g′ = dgdλ(λ; r=r)\n",
    "    if g′ > 0 # active constraint\n",
    "        for iter=1:maxiter # Newton iterations\n",
    "            g″ = d²gdλ²(λ; r=r)\n",
    "            @show Δλ = g′ / g″\n",
    "            λ -= Δλ\n",
    "            abs(Δλ) ≤ rtol * abs(λ) && break # converged\n",
    "            g′ = dgdλ(λ; r=r)\n",
    "        end\n",
    "    end\n",
    "    x̂_λ = (A'A + λ*I) \\ (A'b) # primal optimum x\n",
    "    return λ, x̂_λ\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fff24d2b",
   "metadata": {},
   "source": [
    "For $r = 0.5$, this gives (to nearly machine precision in only 7 Newton steps!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4c37fc10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Δλ = g′ / g″ = -2.976939515696773\n",
      "Δλ = g′ / g″ = -2.197446051423321\n",
      "Δλ = g′ / g″ = -0.7226090155288597\n",
      "Δλ = g′ / g″ = -0.053830681767815296\n",
      "Δλ = g′ / g″ = -0.000263613901856008\n",
      "Δλ = g′ / g″ = -6.262750976944395e-9\n",
      "Δλ = g′ / g″ = -2.9522501729702895e-15\n",
      "(λ_opt, x_opt) = dualopt(0.5) = (5.951088884581378, [-0.2084478531413324, 0.23289743085848733, 0.29901616790443714, 0.25079396035796236])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(5.951088884581378, [-0.2084478531413324, 0.23289743085848733, 0.29901616790443714, 0.25079396035796236])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@show λ_opt, x_opt = dualopt(0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd47afd",
   "metadata": {},
   "source": [
    "i.e. $\\boxed{\\lambda_* \\approx 5.951088884581378}$ and\n",
    "$\\boxed{x_* \\approx [-0.2084478531413324, 0.23289743085848733, 0.29901616790443714, 0.25079396035796236]}$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4f2f1ef",
   "metadata": {},
   "source": [
    "We should also check that $x_*$ is feasible, and in fact that the constraint is active ($\\Vert x_* \\Vert_2 = r$) since $\\lambda_* > 0$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2c3d0137",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.5000000000000001"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norm(x_opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "530e3856",
   "metadata": {},
   "source": [
    "Yes, it is feasible and the constraint is active (up to roundoff errors)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7878762b",
   "metadata": {},
   "source": [
    "## Problem 3 (5+5+6 points)\n",
    "\n",
    "In this problem, you will use ADMM to solve the (primal) optimization problem from problem 2 above, for the parameters from problem 2c, using the equivalent formulation:\n",
    "$$\n",
    "\\min_{x \\in \\mathbb{R}^n} \\left( \\Vert b - Ax \\Vert_2^2 + \\begin{cases} 0 & \\Vert x \\Vert_2 \\le r \\\\ \\infty & \\mbox{otherwise} \\end{cases} \\right)\n",
    "$$\n",
    "where the second term is the \"indicator\" function of the feasible set (the radius-$r$ ball) as in lecture 24 and section III.4 of the text.\n",
    "\n",
    "A basic iteration of ADMM consists of 3 steps, as described in class and in the textbook:\n",
    "\n",
    "1. $x^{(k+1)} = \\mbox{arg }\\min_x \\Vert b - Ax \\Vert_2^2  + \\frac{\\rho}{2} \\Vert x - z^{(k)} + s^{(k)} \\Vert_2^2$ (for some penalty parameter $\\rho > 0$)\n",
    "2. $z^{(k+1)}$ is the projection of $x^{(k+1)} + s^{(k)}$ onto the *closest* point in the feasible set.\n",
    "3. $s^{(k+1)} = s^{(k)} + x^{(k+1)} - z^{(k+1)}$\n",
    "\n",
    "**(a)** Give a closed-form solution for step 1.  (Hint: a problem from pset 3 should be helpful.)\n",
    "\n",
    "**(b)** Give a closed-form solution for step 2.\n",
    "\n",
    "**(c)** Implement this iteration in Julia to solve this problem with the parameters from 2c above, starting from $x = z = s = \\vec{0}$.   Make a (semi-log) plot of the error $\\Vert x^{(k)} - x_* \\Vert_2$ versus $k$, where $x_*$ is your solution from 2c, for $\\rho = 1$ and $\\rho = 10$.  (The error should converge to zero!)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee954031",
   "metadata": {},
   "source": [
    "### Solution:\n",
    "\n",
    "**(a)** This is essentiallly the same as problem 3 of pset 3, with the \"prior\" $x_p = z^{(k)} - s^{(k)}$ and a penalty strength $\\delta^2 = \\rho/2$.  Quoting the solutions from pset 3, we then have:\n",
    "$$\n",
    "\\boxed{x^{(k+1)} = \\left(A^T A + \\frac{\\rho}{2} I\\right)^{-1} \\left[A^T b + \\frac{\\rho}{2} (z^{(k)} - s^{(k)}) \\right]}\n",
    "$$\n",
    "\n",
    "**(b)** To find the closest point to a vector $v$ on a sphere of radius $r$, if $\\Vert v \\Vert_2 > r$ we simply shrink $v$ onto the sphere: $v \\to v r / \\Vert v \\Vert_2$.  This gives:\n",
    "\n",
    "$$\n",
    "\\boxed{z^{(k+1)} =  \\underbrace{\\left( x^{(k+1)} + s^{(k)} \\right)}_{v^{(k)}} \\frac{\\min \\{r, \\Vert v^{(k)} \\Vert_2\\}}{\\Vert v^{(k)} \\Vert_2} }\n",
    "$$\n",
    "where the second term shrinks the length if it is $> r$.  (You can write this in other ways, of course, e.g. as an \"if–then\" statement.)\n",
    "\n",
    "**(c)** We'll implement a function `x,z,s = ADMM(x,z,s)` to take an ADMM step using the formulas above, and then iterate it for 1000 iterations and plot the convergence towards the optimum `x_opt` from problem 2(c) above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e0f4080f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ADMM (generic function with 1 method)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# perform 1 ADMM step, returning the new x,z,s vectors\n",
    "function ADMM(x,z,s; r=0.5, ρ=1)\n",
    "    x = (A'A + (ρ/2)*I) \\ (A'b + (ρ/2)*(z - s))\n",
    "    v = x + s\n",
    "    vlen = norm(v)\n",
    "    z = vlen ≤ r ? v : v * r/vlen\n",
    "    s = s + x - z\n",
    "    return x,z,s\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "483e87b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = z = s = zeros(n)\n",
    "err1 = [norm(x - x_opt)]\n",
    "for i = 2:1000\n",
    "    x,z,s = ADMM(x,z,s; ρ=1)\n",
    "    push!(err1, norm(x - x_opt))\n",
    "end\n",
    "\n",
    "x = z = s = zeros(n)\n",
    "err10 = [norm(x - x_opt)]\n",
    "for i = 2:1000\n",
    "    x,z,s = ADMM(x,z,s; ρ=10)\n",
    "    push!(err10, norm(x - x_opt))\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f2ce39f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 32836 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 640x480 with 1 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "PyObject <matplotlib.legend.Legend object at 0x295a97af0>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "semilogy(err1, \"b.-\")\n",
    "semilogy(err10, \"r.-\")\n",
    "xlabel(\"iteration\")\n",
    "ylabel(L\"error $\\Vert x - x_*\\Vert_2$\")\n",
    "title(\"Problem 3c: ADMM convergence\")\n",
    "legend([L\"\\rho = 1\", L\"\\rho = 10\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5132c6e8",
   "metadata": {},
   "source": [
    "As desired, the error is clearly converging towards zero — ADMM is converging towards the same solution as 2(c) above.\n",
    "\n",
    "Of course, it's not converging nearly as fast as a Newton iteration on the dual problem.  In this problem, $\\rho = 10$ is converging significantly faster than $\\rho = 1$.  \n",
    "\n",
    "**Optional commentary:**\n",
    "\n",
    "However, it's not simply a case of \"bigger is better\".  For any given problem, there is an optimum $\\rho$ for ADMM convergence.  Here, if we increase to $\\rho = 100$ the convergence is better, but if we go to $\\rho = 1000$ it gets *worse*.  (You were not required to investigate this, however.)  This kind of hyper-parameter tuning is painful to do by hand; a more sophisticated approach might be to try to update $\\rho$ dynamically depending on the convergence of the objective function and the constraint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "040588ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43672 characters]",
      "text/plain": [
       "Figure(PyObject <Figure size 640x480 with 1 Axes>)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "PyObject <matplotlib.legend.Legend object at 0x29798e580>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = z = s = zeros(n)\n",
    "err100 = [norm(x - x_opt)]\n",
    "for i = 2:1000\n",
    "    x,z,s = ADMM(x,z,s; ρ=100)\n",
    "    push!(err100, norm(x - x_opt))\n",
    "end\n",
    "\n",
    "x = z = s = zeros(n)\n",
    "err1000 = [norm(x - x_opt)]\n",
    "for i = 2:1000\n",
    "    x,z,s = ADMM(x,z,s; ρ=1000)\n",
    "    push!(err1000, norm(x - x_opt))\n",
    "end\n",
    "\n",
    "semilogy(err10, \"r.-\")\n",
    "semilogy(err100, \"k.-\")\n",
    "semilogy(err1000, \"g.-\")\n",
    "xlabel(\"iteration\")\n",
    "ylabel(L\"error $\\Vert x - x_*\\Vert_2$\")\n",
    "title(\"Optional Problem 3c: ADMM convergence for larger ρ\")\n",
    "legend([L\"\\rho = 10\", L\"\\rho = 100\", L\"\\rho = 1000\"])"
   ]
  }
 ],
 "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
}
