{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true,
    "tags": [
     "active-ipynb",
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple Optimization using Simultaneous Derivatives\n",
    "",
    "\n",
    "Consider a set of points in 2-d space that are to be arranged along a circle such that the radius\n",
    "of the circle is maximized, subject to constraints.  We start out with the points randomly\n",
    "distributed within a unit circle centered about the origin.  The locations of our points are\n",
    "determined by the values of the *x* and *y* arrays defined in our problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "SIZE = 10\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('arctan_yox', om.ExecComp('g=arctan(y/x)', has_diag_partials=True,\n",
    "                                                g=np.ones(SIZE), x=np.ones(SIZE), y=np.ones(SIZE)),\n",
    "                      promotes_inputs=['x', 'y'])\n",
    "\n",
    "p.model.add_subsystem('circle', om.ExecComp('area=pi*r**2'), promotes_inputs=['r'])\n",
    "\n",
    "p.model.add_subsystem('r_con', om.ExecComp('g=x**2 + y**2 - r', has_diag_partials=True,\n",
    "                                           g=np.ones(SIZE), x=np.ones(SIZE), y=np.ones(SIZE)),\n",
    "                      promotes_inputs=['r', 'x', 'y'])\n",
    "\n",
    "thetas = np.linspace(0, np.pi/4, SIZE)\n",
    "p.model.add_subsystem('theta_con', om.ExecComp('g = x - theta', has_diag_partials=True,\n",
    "                                               g=np.ones(SIZE), x=np.ones(SIZE),\n",
    "                                               theta=thetas))\n",
    "p.model.add_subsystem('delta_theta_con', om.ExecComp('g = even - odd', has_diag_partials=True,\n",
    "                                                     g=np.ones(SIZE//2), even=np.ones(SIZE//2),\n",
    "                                                     odd=np.ones(SIZE//2)))\n",
    "\n",
    "p.model.add_subsystem('l_conx', om.ExecComp('g=x-1', has_diag_partials=True, g=np.ones(SIZE), x=np.ones(SIZE)),\n",
    "                      promotes_inputs=['x'])\n",
    "\n",
    "IND = np.arange(SIZE, dtype=int)\n",
    "ODD_IND = IND[1::2]  # all odd indices\n",
    "EVEN_IND = IND[0::2]  # all even indices\n",
    "\n",
    "p.model.connect('arctan_yox.g', 'theta_con.x')\n",
    "p.model.connect('arctan_yox.g', 'delta_theta_con.even', src_indices=EVEN_IND)\n",
    "p.model.connect('arctan_yox.g', 'delta_theta_con.odd', src_indices=ODD_IND)\n",
    "\n",
    "p.driver = om.ScipyOptimizeDriver()\n",
    "p.driver.options['optimizer'] = 'SLSQP'\n",
    "p.driver.options['disp'] = False\n",
    "\n",
    "# set up dynamic total coloring here\n",
    "p.driver.declare_coloring()\n",
    "\n",
    "p.model.add_design_var('x')\n",
    "p.model.add_design_var('y')\n",
    "p.model.add_design_var('r', lower=.5, upper=10)\n",
    "\n",
    "# nonlinear constraints\n",
    "p.model.add_constraint('r_con.g', equals=0)\n",
    "\n",
    "p.model.add_constraint('theta_con.g', lower=-1e-5, upper=1e-5, indices=EVEN_IND)\n",
    "p.model.add_constraint('delta_theta_con.g', lower=-1e-5, upper=1e-5)\n",
    "\n",
    "# this constrains x[0] to be 1 (see definition of l_conx)\n",
    "p.model.add_constraint('l_conx.g', equals=0, linear=False, indices=[0,])\n",
    "\n",
    "# linear constraint\n",
    "p.model.add_constraint('y', equals=0, indices=[0,], linear=True)\n",
    "\n",
    "p.model.add_objective('circle.area', ref=-1)\n",
    "\n",
    "p.setup(mode='fwd')\n",
    "\n",
    "# the following were randomly generated using np.random.random(10)*2-1 to randomly\n",
    "# disperse them within a unit circle centered at the origin.\n",
    "p.set_val('x', np.array([ 0.55994437, -0.95923447,  0.21798656, -0.02158783,  0.62183717,\n",
    "                          0.04007379,  0.46044942, -0.10129622,  0.27720413, -0.37107886]))\n",
    "p.set_val('y', np.array([ 0.52577864,  0.30894559,  0.8420792 ,  0.35039912, -0.67290778,\n",
    "                         -0.86236787, -0.97500023,  0.47739414,  0.51174103,  0.10052582]))\n",
    "p.set_val('r', .7)\n",
    "\n",
    "p.run_driver()\n",
    "\n",
    "print(p['circle.area'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from numpy.testing import assert_almost_equal\n",
    "assert_almost_equal(p['circle.area'], np.pi, decimal=7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Total derivatives with respect to *x* and *y* will be solved for simultaneously based on the\n",
    "color of the points shown below.  Note that we have two colors and our *x* and *y* design\n",
    "variables are of size 10. We have a third design variable, *r*, that is size 1.\n",
    "This means that if we don't solve for derivatives simultaneously, we must\n",
    "perform 21 linear solves (10 + 10 + 1) to solve for total derivatives with respect to\n",
    "all of our design variables.  But, since all of our design variables have only\n",
    "5 colors, we can solve for all of our total derivatives using only 5 linear solves.\n",
    "This means that using simultaneous derivatives saves us 16 linear solves each time we compute our\n",
    "total derivatives.\n",
    "\n",
    "Here's our problem at the start of the optimization:\n",
    "\n",
    "![pre optimization](images/circle_example1.png)\n",
    "\n",
    "After the optimization, all of our points lie along the unit circle.  The final radius is 1.0\n",
    "(which corresponds to an area of PI) because we constrained our *x[0]* value to be 1.0.\n",
    "\n",
    "![pre optimization](images/circle_example2.png)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
