{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HGDL Constrained Optimization of Rosenbrock's Function\n",
    "In this script, we show how HGDL is used for constrained optimization. Unconstrained optimization is simpler and automatically included. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install hgdl==2.1.9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## First some function to make nice plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "def plot(x,y,z,data = None, constr = None):\n",
    "    fig = go.Figure()\n",
    "    fig.add_trace(go.Surface(x = x, y = y,z=z))\n",
    "    if data is not None: \n",
    "        fig.add_trace(go.Scatter3d(x=data[:,0], y=data[:,1], z=data[:,2] + 50,\n",
    "                                   mode='markers'))\n",
    "    if constr is not None: \n",
    "        fig.add_trace(go.Scatter3d(x=constr[:,0], y=constr[:,1], z=constr[:,2],\n",
    "                                   mode='markers')) \n",
    "\n",
    "    fig.update_layout(title='Surface Plot', autosize=True,\n",
    "                  width=800, height=800, font=dict(\n",
    "                  family=\"Courier New, monospace\",\n",
    "                  size=18),\n",
    "                  margin=dict(l=65, r=50, b=65, t=90))\n",
    "\n",
    "    fig.show()\n",
    "\n",
    "def make_plot(bounds, function, data = None, constraint = None):\n",
    "    x1,x2 = np.linspace(bounds[0,0],bounds[0,1],100),np.linspace(bounds[1,0],bounds[1,1],100)\n",
    "    x_pred = np.transpose([np.tile(x1, len(x2)), np.repeat(x2, len(x1))])\n",
    "    x1,x2 = np.meshgrid(x1,x2)\n",
    "    z = np.zeros((10000))\n",
    "    func = np.zeros((10000))\n",
    "    cons = np.zeros((10000))\n",
    "    for i in range(10000): \n",
    "        z[i] = rosen(x_pred[i])\n",
    "        if constraint: cons[i] = constraint(x_pred[i])\n",
    "\n",
    "    \n",
    "    plot(x1, x2, z.reshape(100,100).T, data = data)\n",
    "\n",
    "    if constraint: \n",
    "        fig = go.Figure()\n",
    "        fig.add_trace(go.Surface(x = x1, y = x2,z = cons.reshape(100,100).T))\n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining the Constraints and some Bounds\n",
    "Keep in mind that not all local optimizers allow any combination of bounds and constraints\n",
    "Visit https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html\n",
    "for more information on that"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bounds = np.array([[-4,4],[-4,4]])\n",
    "def g1(x):\n",
    "    return (np.linalg.norm(x)**2/10.0) - 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now we import HGDL and run a constrained optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from hgdl.hgdl import HGDL as hgdl\n",
    "from hgdl.support_functions import *\n",
    "import time\n",
    "from scipy.optimize import rosen, rosen_der, rosen_hess\n",
    "\n",
    "\n",
    "#constraint definitions form scipy\n",
    "from scipy.optimize import NonlinearConstraint\n",
    "nlc = NonlinearConstraint(g1, 0, 10)\n",
    "\n",
    "\n",
    "a = hgdl(rosen, rosen_der, bounds,\n",
    "            hess = rosen_hess, ##if this is None, the Hessian will be approximated if the local optimizer needs it\n",
    "            #global_optimizer = \"random\", #there are a few options to choose from for the global optimizer\n",
    "            global_optimizer = \"genetic\",\n",
    "            local_optimizer = \"L-BFGS-B\", #dNewton is an example and will be changed automatically to \"SLSQP\" because constraints are used\n",
    "            number_of_optima = 30000, #the number fo optima that will be stored and used for deflation\n",
    "            args = (), num_epochs = 1000, #the number fo total epochs. Since this is an asynchronous algorithms, this number can be very high \n",
    "            constraints = (nlc,) #the constraints\n",
    "            #constraints = () #if no constraints are used\n",
    "            )\n",
    "    \n",
    "a.optimize(x0=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = a.get_latest()\n",
    "for entry in res: print(entry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = a.kill_client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Making a Plot\n",
    "You should see the constraints and the found optima. If everything worked, the found points are in between the two constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "data = [np.append(entry[\"x\"],entry[\"f(x)\"]) for entry in res]\n",
    "make_plot(bounds, rosen, data = np.array(data), constraint = g1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}