{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HGDL Constrained Optimization of Schwefel'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(data = None, constraint1 = None, constraint2 = None):\n",
    "    x1,x2 = np.linspace(-500,500,100),np.linspace(-500,500,100)\n",
    "    x_pred = np.transpose([np.tile(x1, len(x2)), np.repeat(x2, len(x1))])\n",
    "    r1 = np.sqrt(160000.)\n",
    "    r2 = np.sqrt(40000.)\n",
    "    c1,c2 = r1*np.cos(np.linspace(0,2.*np.pi,100)),r1*np.sin(np.linspace(0,2.*np.pi,100))\n",
    "    d1,d2 = r2*np.cos(np.linspace(0,2.*np.pi,100)),r2*np.sin(np.linspace(0,2.*np.pi,100))\n",
    "    c3 = np.zeros((len(c2)))\n",
    "    d3 = np.zeros((len(c2)))\n",
    "    x1,x2 = np.meshgrid(x1,x2)\n",
    "    z = np.zeros((10000))\n",
    "    zc1 = np.zeros((10000))\n",
    "    zc2 = np.zeros((10000))\n",
    "    for i in range(10000): \n",
    "        z[i] = schwefel(x_pred[i], 1, 1)\n",
    "        if constraint1: zc1[i] = constraint1(x_pred[i])\n",
    "        if constraint2: zc2[i] = constraint2(x_pred[i])\n",
    "    for i in range(len(c1)):\n",
    "        c3[i] = schwefel(np.array([c1[i],c2[i]]))\n",
    "        d3[i] = schwefel(np.array([d1[i],d2[i]]))\n",
    "    plot(x1,x2,z.reshape(100,100).T, data = data, \n",
    "         constr = np.row_stack([np.column_stack([c1,c2,c3]),np.column_stack([d1,d2,d3])]))"
   ]
  },
  {
   "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": [
    "\n",
    "bounds = np.array([[-500,500],[-500,500]])\n",
    "def g1(x):\n",
    "    return (np.linalg.norm(x)**2/10.0) - 16000.0\n",
    "def g2(x):\n",
    "    return (np.linalg.norm(x)**2/10.0) - 4000.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now we import HGDL and run a constrained optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from hgdl.hgdl import HGDL as hgdl\n",
    "from hgdl.support_functions import *\n",
    "import time\n",
    "\n",
    "#example arguments that will be passed to the function, the gradient and the Hessian function\n",
    "a  = 5\n",
    "b  = 6\n",
    "\n",
    "#constraint definitions form scipy\n",
    "from scipy.optimize import NonlinearConstraint\n",
    "nlc = NonlinearConstraint(g1, -np.inf, 0)\n",
    "nlc = NonlinearConstraint(g2, 0,np.inf)\n",
    "\n",
    "\n",
    "a = hgdl(schwefel, schwefel_gradient, bounds,\n",
    "            hess = None, ##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 = \"dNewton\", #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 = (a,b), num_epochs = 1000, #the number of total epochs. Since this is an asynchronous algorithms, this number can be very high \n",
    "            constraints = (nlc,) #the constraints\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": "markdown",
   "metadata": {},
   "source": [
    "## Making a Plot\n",
    "You should see the constraints and the foudn 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(data = np.array(data), constraint1 = g1, constraint2 = g2)"
   ]
  },
  {
   "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
}