{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b82f73f7-c634-4424-8e78-223edb55f555",
   "metadata": {},
   "source": [
    "# Minimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12eb3a6-47b2-463e-8bdd-a13788e682fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "using Optim\n",
    "using PyPlot\n",
    "using Printf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15c42ca7-b378-484e-a01f-3a81c9995986",
   "metadata": {},
   "source": [
    "## The Rosenbrock function\n",
    "\n",
    "The function is used as a performance test problem for optimization algorithms. It is also known as Rosenbrock's valley or Rosenbrock's banana function. \n",
    "\n",
    "The function is defined by\n",
    "\n",
    "$$f(x,y) = (a-x)^2 + b \\,(y-x^{2})^2 . $$\n",
    "\n",
    "It has a global minimum at $(x,y) = (a, a^2)$, where \n",
    "$f(x,y) = 0$. Usually, the parameters are set such that \n",
    "$a = 1$ and $b=100$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aeb13a3-5629-4e44-aeaa-0d4f5f2a4a4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "f(x) = (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a8bfcd-2303-4efe-af9f-964a7d6f60e2",
   "metadata": {},
   "source": [
    "The global minimum of the function is inside a long, narrow, parabolic-shaped flat valley. \n",
    "\n",
    "### Contour levels of Rosenbrock function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aae968f-0f10-4c8d-bc5b-fa87f4b62cfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\"\"\"\n",
    "    xx, yy, zz = rosenbrock_values(;xmin=-2.0, xmax=2.0, ymin=-1.0, ymax=3.0, np=500)\n",
    "\n",
    "Calculate the values of rosenbrock function on a rectangular grid.\n",
    "\"\"\"\n",
    "function rosenbrock_values(;xmin=-2.0, xmax=2.0, ymin=-1.0, ymax=3.0, np=500)\n",
    "    xx = range(xmin, xmax, np)\n",
    "    yy = range(ymin, ymax, np)\n",
    "\n",
    "    combine(x, y) = [x, y]  # construct a vector out of components\n",
    "    xy = combine.(xx, yy')\n",
    "\n",
    "    zz = f.(xy)\n",
    "    \n",
    "    return xx, yy, zz\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "679ef250-0c73-4f2d-8448-236702d41fc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "xx, yy, zz = rosenbrock_values();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "831b311f-7d78-4363-a2cd-32aaaad9a17c",
   "metadata": {},
   "source": [
    "Manually set the contour levels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1cd87c5-3470-4384-8f47-c6075ab343ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "nlevels = 15\n",
    "lev_exp = range(floor(log10(minimum(zz))-1), ceil(log10(maximum(zz)+1)), nlevels)\n",
    "levs = 10 .^ lev_exp;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3dfcc53-7ca2-49b4-b253-b5594067ab8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "figure(figsize=(9,6))\n",
    "contourf(xx, yy, zz', norm=\"log\", levs, cmap=\"coolwarm\")\n",
    "colorbar(ticks=[1e-7, 1e-5, 1e-3, 1e-1, 1e1, 1e3])\n",
    "contour(xx, yy, zz', levs, linewidths=0.5)\n",
    "axis(\"square\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec4a61c1-c71d-4dc8-97de-14bf978489be",
   "metadata": {},
   "source": [
    "## Minimization options:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8aaa0b1d-b5d3-4d10-9e7d-e97470868411",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "opts = Optim.Options(store_trace=true, extended_trace=true)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb6ee161-a1a0-43fc-ac9c-8cfa470eb7d6",
   "metadata": {},
   "source": [
    "## Helper animation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69eaa601-ada5-45f1-b9c2-06ca9ae80b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\"\"\"\n",
    "    animation(xpath, xx, yy, zz, levs; sleeptime=0.2)\n",
    "\n",
    "Animate the minimization path `xpath`; xx, yy, zz, and levs are for contour plotting\n",
    "\"\"\"\n",
    "function animation(xpath, xx, yy, zz, levs; sleeptime=0.2)\n",
    "    nit = length(xpath)\n",
    "    fig = figure(figsize=(9,6))\n",
    "    for i = 1:nit\n",
    "        ax = gca()\n",
    "        ax.set_aspect(\"equal\")\n",
    "        contourf(xx, yy, zz', norm=\"log\", levs, cmap=\"coolwarm\")\n",
    "        contour(xx, yy, zz', levs, linewidths=0.5)\n",
    "        scatter(xpath[i][1], xpath[i][2], color=\"blue\")\n",
    "\n",
    "        titl = @sprintf(\"Step %3d/%3d\", i, nit)\n",
    "        title(titl, family=\"monospace\")\n",
    "        xlabel(\"x\")\n",
    "        ylabel(\"y\")\n",
    "        grid(true)\n",
    "    \n",
    "        display(fig)\n",
    "        IJulia.clear_output(true)\n",
    "        clf()\n",
    "\n",
    "        sleep(sleeptime)\n",
    "    end\n",
    " \n",
    "    ax = gca()\n",
    "    ax.set_aspect(\"equal\")\n",
    "    titl = @sprintf(\"Step %3d/%3d\", nit, nit)\n",
    "    title(titl, family=\"monospace\")\n",
    "    xlabel(\"x\")\n",
    "    ylabel(\"y\")\n",
    "    grid(true)\n",
    "    contourf(xx, yy, zz', norm=\"log\", levs, cmap=\"coolwarm\")\n",
    "    # contour(xx, yy, cbrt.(zz), nlevels, cmap=\"cool\")   # \"flatten\" the function using cbrt\n",
    "    contour(xx, yy, zz', levs, linewidths=0.5)\n",
    "    plot(first.(xpath), last.(xpath), marker=\".\", color=\"black\")\n",
    "    scatter(xpath[end][1], xpath[end][2], color=\"blue\")\n",
    "    \n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad99a15d-47d5-46b8-9eb0-864ab0bbd7fc",
   "metadata": {},
   "source": [
    "## Unconstrained Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f12ead83-eba6-47f0-9c2f-36675bb3dd9c",
   "metadata": {},
   "source": [
    "Initial approximation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df839093-475c-4519-985b-d45498e46fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "initial_x = [-1.9, 2.0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9bc9a6f-e75f-4547-8c81-ca2e12425088",
   "metadata": {},
   "source": [
    "The methods below require a gradient. One can use automatic differentiation to \n",
    "calculate the gradient, by using the autodiff keyword and setting it to :forward:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "284745d7-abce-4898-8372-45a494810b57",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "res = optimize(f, initial_x, LBFGS(), opts; autodiff = :forward)\n",
    "# res = optimize(f, initial_x, GradientDescent(), opts; autodiff = :forward)\n",
    "# res = optimize(f, initial_x, ConjugateGradient(), opts; autodiff = :forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd72afe-f40e-43e7-b7e6-32428fb2da3a",
   "metadata": {},
   "source": [
    "Position of the minimum:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e1a287e-f9e3-4102-8f0b-47b457c55b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "Optim.minimizer(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33653620-5870-4cc5-86eb-e52cd61fbe0c",
   "metadata": {},
   "source": [
    "The animation of minimization path:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "647d9970-a284-47cf-bc0b-74b8a111dedf",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "xpath = Optim.x_trace(res);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dfd7cef-2d34-4360-94b9-b1beeabce468",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "animation(xpath, xx, yy, zz, levs, sleeptime=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "782a8905-d3d5-496a-aaab-04e66145be14",
   "metadata": {},
   "source": [
    "Decrease of the objective value during the minimization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "241faf3b-3ef8-4cae-a4fc-9c9e5ee7cefd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "semilogy(Optim.f_trace(res))\n",
    "grid(true)\n",
    "xlabel(\"minimization step\")\n",
    "ylabel(\"objective value\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36019c2b-251f-4f27-9e01-3d498dee57c6",
   "metadata": {},
   "source": [
    "## Box Constrained Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9edf84b3-cbcc-4afb-a0d9-d23030030aa6",
   "metadata": {},
   "source": [
    "Constrants:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84d7f908-a9c7-4601-9893-948abe4c64ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "lower = [-2.0, 1.5]\n",
    "upper = [Inf, Inf];"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "596d34e7-1de6-424e-bb73-f3819ed919b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "initial_x = [0.2, 2.55];"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f7d9e2d-1284-4f9c-b53c-6337fcb54ec3",
   "metadata": {},
   "source": [
    "This performs optimization with a barrier penalty, successively scaling down the barrier coefficient and using the chosen inner_optimizer for convergence at each step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f18ec963-618f-4337-b17f-120350f34fad",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "res = optimize(f, lower, upper, initial_x, Fminbox(LBFGS()), opts; autodiff = :forward)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4731e61-a034-4b3c-b707-76c92f5c8bd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "xpath = Optim.x_trace(res);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "148c4df6-ebed-453a-9425-38597f517952",
   "metadata": {},
   "source": [
    "The animation of minimization path:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1621ad39-a026-4324-b281-67ef07011bf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "animation(xpath, xx, yy, zz, levs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b907bcdc-5aef-47fa-8609-0cfdc960ad6b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.7",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}