{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "deletable": false,
    "editable": false,
    "id": "bx-ZeOJAQGQS",
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "d7142079a02673557fd4d7c62d09fb35",
     "grade": false,
     "grade_id": "cell-59e5add4db8ca4cf",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "# Part 1: Optimising functions\n",
    "    \n",
    "In this lab we will play with some of the optimisation methods we learned in the lecture by exploring how they work on some analytic functions (both convex and non-convex)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "14fde3e4d56930cde83320df869ce52a",
     "grade": false,
     "grade_id": "cell-2c171e030a9385fc",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "deletable": false,
    "editable": false,
    "id": "cG3yqv0BQQVx",
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "d6d5a1e2037a78f4a83c58ba02f204b6",
     "grade": false,
     "grade_id": "cell-3c9fb8c6c798fa3e",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "## A Simple Function\n",
    "\n",
    "For this first task, we are going to try to optimise the following using Stochastic Gradient Descent:\n",
    "\n",
    "\\begin{equation}\n",
    "min_{\\textbf{x}} (\\textbf{x}[0] - 5)^2 + \\textbf{x}[1]^2 + (\\textbf{x}[2] - 1)^2\\; ,\n",
    "\\end{equation}\n",
    "\n",
    "Use the following block the write down the analytic minima of the above function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "086b733ee28421cd24eb19ade35c6ed4",
     "grade": true,
     "grade_id": "cell-b78040990d924cd0",
     "locked": false,
     "points": 1,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# YOUR CODE HERE\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "62407a53e2c8038e31413b6e939f397c",
     "grade": false,
     "grade_id": "cell-31897742431a78ea",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "### Implement the function\n",
    "\n",
    "First, complete the following code block to implement the above function using PyTorch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "id": "F4BwVbzTQNv6",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "82360d8bcdfd9ec6dccea7295fd4440d",
     "grade": false,
     "grade_id": "cell-205644c1f2119166",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def function(x):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "deletable": false,
    "editable": false,
    "id": "hfTJ2GI_Qcej",
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "7480042c8d67ef7613ee4bfea5d4971c",
     "grade": false,
     "grade_id": "cell-3216df41fb235ae6",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "### Optimising\n",
    "\n",
    "We need two more things before we can start optimising.\n",
    "We need our initial guess - which we've set to [2.0, 1.0, 10.0] and we need to how many epochs to take."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "vvRXozg6QgNV",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "c117216da51f29eddbad68c8be917e46",
     "grade": false,
     "grade_id": "cell-5d4d7fa7ed77956f",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "p = torch.tensor([2.0, 1.0, 10.0], requires_grad=True)\n",
    "epochs = 5000"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "deletable": false,
    "editable": false,
    "id": "nGNScu9QQgbZ",
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "1de659ee2ace2ae86f8601c28e1033e7",
     "grade": false,
     "grade_id": "cell-a71c41511464a13e",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "We define the optimisation loop in the standard way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "editable": false,
    "id": "vkmkfXkaQhou",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "fb907848ef17ec996ee54c0f8aeb932c",
     "grade": false,
     "grade_id": "cell-3ae9b149f3a8a224",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "opt = optim.SGD([p], lr=0.001)\n",
    "\n",
    "for i in range(epochs):\n",
    "    opt.zero_grad()\n",
    "    output = function(p)\n",
    "    output.backward()\n",
    "    opt.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "deletable": false,
    "editable": false,
    "id": "hFXAQAXqQi3H",
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "55c0aa9e2acd143f369fd00849faa6c5",
     "grade": false,
     "grade_id": "cell-f63a8e54eca095d1",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Use the following block to print out the final value of `p`. Does it match the value you expected?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "deletable": false,
    "id": "6DEGWh_rQkJV",
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b4d7857d166ffa2f5634a181ce807851",
     "grade": false,
     "grade_id": "cell-7124ebe0293e33c9",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# YOUR CODE HERE\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualising Himmelblau's Function\n",
    "\n",
    "We'll now have a go at a more complex example, which we also visualise, with multiple optima; [Himmelblau's function](https://en.wikipedia.org/wiki/Himmelblau%27s_function). This is defined as:\n",
    "\n",
    "\\begin{equation}\n",
    "f(x, y) = (x^2 + y - 11)^2 + (x + y^2 - 7)^2\\; ,\n",
    "\\end{equation}\n",
    "and has minima at\n",
    "\\begin{equation}\n",
    "f(3, 2) = f(-2.805118, 3.131312) = f(-3.779310, -3.283186) = f(3.584428, -1.848126) = 0\\; .\n",
    "\\end{equation}\n",
    "\n",
    "Use the following block to first define the function (the inputs $x, y$ are packed into a vector as for the previous quadratic function above):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "1cb1ab20f7bc9cdcd7a221a5134578af",
     "grade": false,
     "grade_id": "cell-05f43ebf2fd7e68c",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "def himm(x):\n",
    "    # YOUR CODE HERE\n",
    "    raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "a87861c51de11fc596ef3de86b950c70",
     "grade": false,
     "grade_id": "cell-fc81ef599d739b56",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "The following will plot its surface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b6f21db5449a4ad97911d24936609a48",
     "grade": false,
     "grade_id": "cell-7d401c3457a5f7e9",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from matplotlib.colors import LogNorm\n",
    "\n",
    "xmin, xmax, xstep = -5, 5, .2\n",
    "ymin, ymax, ystep = -5, 5, .2\n",
    "x, y = np.meshgrid(np.arange(xmin, xmax + xstep, xstep), np.arange(ymin, ymax + ystep, ystep))\n",
    "z = himm(torch.tensor([x, y])).numpy()\n",
    "\n",
    "fig = plt.figure(figsize=(8, 5))\n",
    "ax = plt.axes(projection='3d', elev=50, azim=-50)\n",
    "ax.plot_surface(x, y, z, norm=LogNorm(), rstride=1, cstride=1, \n",
    "                edgecolor='none', alpha=.8, cmap=plt.cm.jet)\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$y$')\n",
    "ax.set_zlabel('$z$')\n",
    "\n",
    "ax.set_xlim((xmin, xmax))\n",
    "ax.set_ylim((ymin, ymax))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "1b6e3379a1d9fb21879d0845ca7ba072",
     "grade": false,
     "grade_id": "cell-f938244581085c83",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Check that the above plot looks correct  by comparing to the picture on the [Wikipedia page](https://en.wikipedia.org/wiki/Himmelblau%27s_function)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "198bab34933b991f896a519e2f4162fd",
     "grade": false,
     "grade_id": "cell-574019e663811dcc",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "### Optimising\n",
    "\n",
    "Let's see how it looks for a few different optimisers from a range of starting points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "f52952455140dbe3b789e48da231adb9",
     "grade": false,
     "grade_id": "cell-2ab57db6f9c3ff31",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "xmin, xmax, xstep = -5, 5, .2\n",
    "ymin, ymax, ystep = -5, 5, .2\n",
    "x, y = np.meshgrid(np.arange(xmin, xmax + xstep, xstep), np.arange(ymin, ymax + ystep, ystep))\n",
    "z = himm(torch.tensor([x, y])).numpy()\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 8))\n",
    "ax.contourf(x, y, z, levels=np.logspace(0, 5, 35), norm=LogNorm(), cmap=plt.cm.gray)\n",
    "\n",
    "p = torch.tensor([[0.0],[0.0]], requires_grad=True)\n",
    "opt = optim.SGD([p], lr=0.01)\n",
    "\n",
    "path = np.empty((2,0))\n",
    "path = np.append(path, p.data.numpy(), axis=1)\n",
    "\n",
    "for i in range(50):\n",
    "    opt.zero_grad()\n",
    "    output = himm(p)\n",
    "    output.backward()\n",
    "    opt.step()\n",
    "    path = np.append(path, p.data.numpy(), axis=1)\n",
    "\n",
    "ax.plot(path[0], path[1], color='red', label='SGD', linewidth=2)\n",
    "\n",
    "ax.legend()\n",
    "ax.set_xlabel('$x$')\n",
    "ax.set_ylabel('$y$')\n",
    "\n",
    "ax.set_xlim((xmin, xmax))\n",
    "ax.set_ylim((ymin, ymax))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "a7281ca078d692fbe12ddfe12a0bd51c",
     "grade": false,
     "grade_id": "cell-f92c8fbe7bc00ba2",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Use the following block to run SGD with momentum (lr=0.01, momentum=0.9) from the same initial point, saving the position at each timestep into a variable called `path_mom`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "5a2c8c4d77435e085753249bb03dd44a",
     "grade": false,
     "grade_id": "cell-8beae2091cb7be41",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# YOUR CODE HERE\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "f581993736855e0316c24d86dafede92",
     "grade": false,
     "grade_id": "cell-1a3d5f0a22670713",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "The following will plot the path taken when momentum was used, as well as the original plain SGD path:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "ef449c500d6e15915d34f0ae82f7474f",
     "grade": false,
     "grade_id": "cell-b5ee1281d6ac6ec8",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "ax.plot(path_mom[0], path_mom[1], color='yellow', label='SGDM', linewidth=2)\n",
    "ax.legend()\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "8ce2cd5108842b8589b0a7c0e789a54e",
     "grade": false,
     "grade_id": "cell-654d173df302005c",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "Now explore what happens when you start from different points. What effect do you get with different optimisers? "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "98b599e2b344f4d5643487f8bb78ed5d",
     "grade": false,
     "grade_id": "cell-c95af05c3c20a927",
     "locked": false,
     "schema_version": 3,
     "solution": true,
     "task": false
    }
   },
   "outputs": [],
   "source": [
    "# YOUR CODE HERE\n",
    "raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "3_1_FuntionOptimisation.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}