{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization\n",
    "\n",
    "The notebook will walk through a simple example of performing optimization of a CSDL model using [modopt](https://github.com/LSDOlab/modopt). \n",
    "\n",
    "Begin by initializing a csdl `Recorder` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import csdl_alpha as csdl\n",
    "import numpy as np\n",
    "\n",
    "recorder = csdl.Recorder(inline=True)\n",
    "recorder.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, we will use the [toy sellar optimization problem](https://openmdao.org/newdocs/versions/latest/basic_user_guide/multidisciplinary_optimization/sellar.html) as exactly defined in OpenMDAO's documentation. The problem is a simple system with coupled disciplines. See their documentation page for more details.\n",
    "\n",
    " We define the problem here using namespaces, nonlinear solvers, as well as special methods to specify optimization variables: `set_as_design_variable`, `set_as_constraint` and `set_as_objective`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n"
     ]
    }
   ],
   "source": [
    "z1 = csdl.Variable(name = 'z1', value = 5)\n",
    "z2 = csdl.Variable(name = 'z2', value = 2)\n",
    "x = csdl.Variable(name = 'x', value = 1.0)\n",
    "y2 = csdl.ImplicitVariable(name = 'y2', value = 1.0)\n",
    "\n",
    "z1.set_as_design_variable(lower=-10.0, upper=10.0) # design variable\n",
    "z2.set_as_design_variable(lower= 0.0, upper=10.0) # design variable\n",
    "x.set_as_design_variable(lower= 0.0, upper=10.0) # design variable\n",
    "\n",
    "# Define each \"component\" from the example\n",
    "with csdl.namespace('Discipline 1'):\n",
    "    y1 = z1**2 + z2 + x - 0.2*y2\n",
    "    y1.add_name('y1')\n",
    "\n",
    "with csdl.namespace('Discipline 2'):\n",
    "    residual = csdl.sqrt(y1) + z1 + z2 - y2\n",
    "    residual.add_name('residual')\n",
    "\n",
    "with csdl.namespace('Objective'):\n",
    "    f = x**2 + z2 + y1 + csdl.exp(-y2)\n",
    "    f.add_name('f')\n",
    "    f.set_as_objective() # objective\n",
    "\n",
    "with csdl.namespace('Constraint 1'):\n",
    "    g1 = 3.16 - y1\n",
    "    g1.add_name('g1')\n",
    "    g1.set_as_constraint(upper=0.0) # constraint\n",
    "\n",
    "with csdl.namespace('Constraint 2'):\n",
    "    g2 = y2 - 24.0\n",
    "    g2.add_name('g2')\n",
    "    g2.set_as_constraint(upper=0.0) # constraint\n",
    "\n",
    "# Specifiy coupling\n",
    "with csdl.namespace('Couple'):\n",
    "    solver = csdl.nonlinear_solvers.Newton()\n",
    "    solver.add_state(y2, residual, tolerance=1e-8)\n",
    "    solver.run()\n",
    "\n",
    "recorder.stop()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we set three design variables, two constraints and an objective. Because we instantiated the `Recorder` with `inline = True`, we can see the values of the initial evaluation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intital run values:\n",
      "y1:   [25.58830237]\n",
      "y2:   [12.05848815]\n",
      "f:    [28.58830816]\n",
      "g1:   [-22.42830237]\n",
      "g2:   [-11.94151185]\n",
      "z1:   [5.]\n",
      "z2:   [2.]\n"
     ]
    }
   ],
   "source": [
    "print('intital run values:')\n",
    "print('y1:  ',  y1.value)\n",
    "print('y2:  ',  y2.value)\n",
    "print('f:   ',  f.value)\n",
    "print('g1:  ',  g1.value)\n",
    "print('g2:  ',  g2.value)\n",
    "print('z1:  ',  z1.value)\n",
    "print('z2:  ',  z2.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CSDL makes it easy to evaluate the model by providing the `Simulator` class. As the name suggests, the `Simulator` class is an interface for *evaluating* your model and any derivatives *after* the model is defined.\n",
    "\n",
    "We currently provide two experimental simulators: `PySimulator` and `JaxSimulator` which computes your model using numpy and [jax](https://github.com/google/jax) respectively. For this example, we will use the `PySimulator` for easier use. See the 'JAX' tutorial to learn more about running your model efficiently with `JaxSimulator`. Creating a `PySimulator` is as simple as instantiating it with your `Recorder` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = csdl.experimental.PySimulator(recorder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To interface our simulator with `modopt`, we use `modopt`'s `CSDLAlphaProblem`. We can then solve the optimization using `modopt`'s API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "Deleting self.pFpx ...\n",
      "Deleting self.pCpx, pCpx_dict ...\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing objective >>>>>>>>>>\n",
      "---------Computed objective---------\n",
      "Computing gradient >>>>>>>>>>\n",
      "---------Computed gradient---------\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing Jacobian >>>>>>>>>>\n",
      "---------Computed Jacobian---------\n",
      "Computing objective >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "---------Computed objective---------\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing gradient >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "---------Computed gradient---------\n",
      "Computing Jacobian >>>>>>>>>>\n",
      "---------Computed Jacobian---------\n",
      "Computing objective >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed objective---------\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing gradient >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed gradient---------\n",
      "Computing Jacobian >>>>>>>>>>\n",
      "---------Computed Jacobian---------\n",
      "Computing objective >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed objective---------\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing gradient >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed gradient---------\n",
      "Computing Jacobian >>>>>>>>>>\n",
      "---------Computed Jacobian---------\n",
      "Computing objective >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed objective---------\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing gradient >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed gradient---------\n",
      "Computing Jacobian >>>>>>>>>>\n",
      "---------Computed Jacobian---------\n",
      "Computing objective >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed objective---------\n",
      "Computing constraints >>>>>>>>>>\n",
      "---------Computed constraints---------\n",
      "Computing gradient >>>>>>>>>>\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "---------Computed gradient---------\n",
      "Computing Jacobian >>>>>>>>>>\n",
      "---------Computed Jacobian---------\n"
     ]
    }
   ],
   "source": [
    "from modopt import CSDLAlphaProblem\n",
    "from modopt import SLSQP\n",
    "\n",
    "# Instantiate your problem using the csdl Simulator object and name your problem\n",
    "prob = CSDLAlphaProblem(problem_name='sellar',simulator=sim)\n",
    "\n",
    "optimizer = SLSQP(prob, ftol=1e-9, maxiter=20)\n",
    "\n",
    "# Check first derivatives at the initial guess, if needed\n",
    "# optimizer.check_first_derivatives(prob.x0)\n",
    "\n",
    "# Solve your optimization problem\n",
    "optimizer.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, lets print the optimized values and verify our results match OpenMDAO's [results](https://openmdao.org/newdocs/versions/latest/basic_user_guide/reading_recording/basic_recording_example.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " \t ==============\n",
      "\t Scipy summary:\n",
      "\t ==============\n",
      "\t Problem                    : sellar\n",
      "\t Solver                     : scipy_slsqp\n",
      "\t Success                    : True\n",
      "\t Message                    : Optimization terminated successfully\n",
      "\t Objective                  : 3.183393951728967\n",
      "\t Gradient norm              : 4.022879450485861\n",
      "\t Total time                 : 0.01696324348449707\n",
      "\t Major iterations           : 6\n",
      "\t Total function evals       : 6\n",
      "\t Total gradient evals       : 6\n",
      "\t ==================================================\n",
      "optimized values:\n",
      "y1:   [3.16]\n",
      "y2:   [3.75527777]\n",
      "f:    [3.18339395]\n",
      "g1:   [-8.95310492e-11]\n",
      "g2:   [-20.24472223]\n",
      "z1:   [1.97763888]\n",
      "z2:   [0.]\n"
     ]
    }
   ],
   "source": [
    "optimizer.print_results()\n",
    "\n",
    "print('optimized values:')\n",
    "print('y1:  ',  y1.value)\n",
    "print('y2:  ',  y2.value)\n",
    "print('f:   ',  f.value)\n",
    "print('g1:  ',  g1.value)\n",
    "print('g2:  ',  g2.value)\n",
    "print('z1:  ',  z1.value)\n",
    "print('z2:  ',  z2.value)\n",
    "\n",
    "assert np.isclose(y1.value, 3.15999999)\n",
    "assert np.isclose(y2.value, 3.75527776)\n",
    "assert np.isclose(f.value, 3.18339394)\n",
    "assert np.isclose(g1.value, -8.95310492e-11)\n",
    "assert np.isclose(g2.value, -20.24472224)\n",
    "assert np.isclose(z1.value, 1.97763888)\n",
    "assert np.isclose(z2.value, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We often want to run the model stand-alone without any optimization. `Simulator` provides this functionality by the `run` and `compute_totals` methods. Use the getters and setters to see and change the values of variables. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "current f:       [3.18339395]\n",
      "nonlinear solver: newton_nlsolver converged in 3 iterations.\n",
      "current dfdz1:   [[3.50848986]] \n",
      "\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "\n",
      "new f:            [7.71844379]\n",
      "current dfdz1:    [[5.54894571]]\n"
     ]
    }
   ],
   "source": [
    "print('current f:      ', sim[f])\n",
    "print('current dfdz1:  ', sim.compute_totals(f,z1)[f,z1], '\\n')\n",
    "\n",
    "sim[z1] += 1.0\n",
    "sim.run()\n",
    "dfdz1 = sim.compute_totals(f,z1)[f,z1]\n",
    "\n",
    "print('\\nnew f:           ', sim[f])\n",
    "print('current dfdz1:   ', dfdz1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use `check_totals` to verify analytical derivatives using finite difference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "nonlinear solver: newton_nlsolver converged in 2 iterations.\n",
      "\n",
      "Derivative Verification Results\n",
      "-------------------------------\n",
      "ofs (3)   wrts (3)   norm                  fd norm               error                    rel error                tags              \n",
      "------------------------------------------------------------------------------------------------------------------------------\n",
      "f         z1         5.548945712851889     5.548946721667392     1.0088155022458523e-06   1.8180306152637114e-07   obj(1,),dv(1,),   \n",
      "f         z2         1.7685933165868333    1.7685933180189295    1.4320962193892228e-09   8.097374364126682e-10    obj(1,),dv(1,),   \n",
      "f         x          0.9646991795541301    0.9647001810719757    1.0015178456557194e-06   1.0381648778616708e-06   obj(1,),dv(1,),   \n",
      "g1        z1         5.555277766964987     5.5552787694423955    1.0024774086403454e-06   1.8045492409032929e-07   con(1,),dv(1,),   \n",
      "g1        z2         0.7721994304206554    0.7721994297327228    6.87932599774399e-10     8.908742654893042e-10    con(1,),dv(1,),   \n",
      "g1        x          0.9652492880258193    0.9652492893863496    1.360530244021163e-09    1.4095117799942754e-09   con(1,),dv(1,),   \n",
      "g2        z1         2.0000000000580482    1.9999999913977717    8.66027649593093e-09     4.330138266589884e-09    con(1,),dv(1,),   \n",
      "g2        z2         1.1390028478967227    1.1390028475943836    3.023390426903916e-10    2.654418672692021e-10    con(1,),dv(1,),   \n",
      "g2        x          0.17375355987090357   0.17375354843807145   1.1432832125457537e-08   6.57991288709271e-08     con(1,),dv(1,),   \n"
     ]
    }
   ],
   "source": [
    "dfdz1 = sim.check_totals()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "csdl_a",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}