{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Advanced Recording Example\n",
    "\n",
    "Below we demonstrate a more advanced example of case recording including the four different objects\n",
    "that a recorder can be attached to. We will then show how to extract various data from the model and finally,\n",
    "relate that back to an XDSM for illustrative purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_SellarMDAWithUnits\", get_code(\"openmdao.test_suite.components.sellar_feature.SellarMDAWithUnits\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `SellarMDAWithUnits` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_SellarMDAWithUnits`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.sellar_feature import SellarMDAWithUnits\n",
    "\n",
    "# build the model\n",
    "prob = om.Problem(model=SellarMDAWithUnits())\n",
    "\n",
    "model = prob.model\n",
    "model.add_design_var('z', lower=np.array([-10.0, 0.0]),\n",
    "                          upper=np.array([10.0, 10.0]))\n",
    "model.add_design_var('x', lower=0.0, upper=10.0)\n",
    "model.add_objective('obj')\n",
    "model.add_constraint('con1', upper=0.0)\n",
    "model.add_constraint('con2', upper=0.0)\n",
    "\n",
    "# setup the optimization\n",
    "prob.driver = om.ScipyOptimizeDriver(optimizer='SLSQP', tol=1e-9, disp=False)\n",
    "\n",
    "# Here we show how to attach recorders to each of the four objects:\n",
    "#   problem, driver, solver, and system\n",
    "\n",
    "# Create a recorder\n",
    "recorder = om.SqliteRecorder('cases.sql')\n",
    "\n",
    "# Attach recorder to the problem\n",
    "prob.add_recorder(recorder)\n",
    "\n",
    "# Attach recorder to the driver\n",
    "prob.driver.add_recorder(recorder)\n",
    "\n",
    "# To attach a recorder to a subsystem or solver, you need to call `setup`\n",
    "# first so that the model hierarchy has been generated\n",
    "prob.setup()\n",
    "\n",
    "# Attach recorder to a subsystem\n",
    "model.obj_cmp.add_recorder(recorder)\n",
    "\n",
    "# Attach recorder to a solver\n",
    "model.cycle.nonlinear_solver.add_recorder(recorder)\n",
    "\n",
    "prob.set_solver_print(0)\n",
    "prob.run_driver()\n",
    "prob.record(\"final_state\")\n",
    "prob.cleanup()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following XDSM diagram shows the SellarMDA component equations and their inputs and outputs. Through\n",
    "the different recorders we can access the different parts of the model. We'll take you through an\n",
    "example of each object and relate it back to the diagram.\n",
    "\n",
    "![sellar_xdsm](images/sellar_xdsm.jpg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Driver\n",
    "If we want to view the convergence of the model, the best place to find that by looking at the cases \n",
    "recorded by the `driver`. By default, a recorder attached to a driver will record the design variables,\n",
    "constraints and objectives, so we will print them for the model at the end of the optimization. We'll \n",
    "use the helper methods like `get_objectives`, `get_design_vars`, `get_constraints` to return the info \n",
    "we're seeking.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "# Instantiate your CaseReader\n",
    "cr = om.CaseReader(\"cases.sql\")\n",
    "\n",
    "# List driver cases (do not recurse to system/solver cases)\n",
    "driver_cases = cr.list_cases('driver', recurse=False);\n",
    "last_case = cr.get_case(driver_cases[-1])\n",
    "\n",
    "objectives = last_case.get_objectives()\n",
    "design_vars = last_case.get_design_vars()\n",
    "constraints = last_case.get_constraints()\n",
    "\n",
    "print(objectives['obj'])\n",
    "print(design_vars['x'])\n",
    "print(design_vars['z'])\n",
    "print(constraints['con1'])\n",
    "print(constraints['con2'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(objectives['obj'], 3.18339395, 1e-8)\n",
    "assert_near_equal(design_vars['x'], 0., 1e-8)\n",
    "assert_near_equal(design_vars['z'], [1.97763888, 1.25035459e-15], 1e-8)\n",
    "assert_near_equal(constraints['con1'], -1.68550507e-10, 1e-8)\n",
    "assert_near_equal(constraints['con2'], -20.24472223, 1e-8)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Design Variables\n",
    "\n",
    "When inspecting or debugging a model, it can be helpful to visualize the path of the design\n",
    "variables to their final values. To do this, we can list the cases of the driver and plot the data\n",
    "with respect to the iteration number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import openmdao.api as om\n",
    "\n",
    "# Instantiate your CaseReader\n",
    "cr = om.CaseReader(\"cases.sql\")\n",
    "\n",
    "# Get driver cases (do not recurse to system/solver cases)\n",
    "driver_cases = cr.get_cases('driver', recurse=False)\n",
    "\n",
    "# Plot the path the design variables took to convergence\n",
    "# Note that there are two lines in the right plot because \"Z\"\n",
    "# contains two variables that are being optimized\n",
    "dv_x_values = []\n",
    "dv_z_values = []\n",
    "for case in driver_cases:\n",
    "    dv_x_values.append(case['x'])\n",
    "    dv_z_values.append(case['z'])\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2)\n",
    "fig.subplots_adjust(left=None, bottom=None, right=None, top=None, wspace=0.5, hspace=None)\n",
    "\n",
    "ax1.plot(np.arange(len(dv_x_values)), np.array(dv_x_values))\n",
    "ax1.set(xlabel='Iterations', ylabel='Design Var: X', title='Optimization History')\n",
    "ax1.grid()\n",
    "\n",
    "ax2.plot(np.arange(len(dv_z_values)), np.array(dv_z_values))\n",
    "ax2.set(xlabel='Iterations', ylabel='Design Var: Z', title='Optimization History')\n",
    "ax2.grid()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem\n",
    "\n",
    "A `Problem` recorder is best if you want to record an arbitrary case before or after a running the\n",
    "model. In our case, we have placed our point at the end of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "# Instantiate your CaseReader\n",
    "cr = om.CaseReader(\"cases.sql\")\n",
    "\n",
    "# get list of cases recorded on problem\n",
    "problem_cases = cr.list_cases('problem')\n",
    "\n",
    "# get list of variables recorded on problem\n",
    "problem_vars = cr.list_source_vars('problem')\n",
    "print(problem_vars['outputs'])\n",
    "\n",
    "# get the recorded case and check values\n",
    "case = cr.get_case('final_state')\n",
    "print(case.get_design_vars())\n",
    "print(case.get_constraints())\n",
    "print(case.get_objectives())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert(problem_cases == ['final_state'])\n",
    "assert(sorted(problem_vars['outputs']) == ['con1', 'con2', 'obj', 'x', 'y1', 'y2', 'z'])\n",
    "assert_near_equal(case.get_design_vars(), {'x': 0., 'z': [1.9776, 1.25e-15]}, tolerance=1e-4)\n",
    "assert_near_equal(case.get_constraints(), {'con1': 0., 'con2': -20.2447}, tolerance=1e-4)\n",
    "assert_near_equal(case.get_objectives(), {'obj': 3.18339395}, tolerance=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## System\n",
    "Suppose we want to know the values of the inputs to the objective function. \n",
    "For this we will use the `system` recorder. Using the `list_cases` method, we'll get the list of \n",
    "cases that were recorded by the objective component `root.obj_comp`. Then we use `get_case` to \n",
    "get each case and determine the input names using the keys of the case's inputs dictionary. \n",
    "Since we want to find the values going into the objective function for each execution, \n",
    "we iterate through those 14 cases and get the values in each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Instantiate your CaseReader\n",
    "cr = om.CaseReader(\"cases.sql\")\n",
    "\n",
    "# Get a list of cases for the objective component\n",
    "system_cases = cr.list_cases('root.obj_cmp', recurse=False)\n",
    "\n",
    "# Number of cases recorded for 'obj_cmp'\n",
    "print(f\"Number of cases: {len(system_cases)}\")\n",
    "\n",
    "for case_num, case_id in enumerate(system_cases):\n",
    "    case = cr.get_case(case_id)\n",
    "    \n",
    "    # Get the names of all the inputs to the objective component\n",
    "    inputs = case.inputs.keys()\n",
    "    \n",
    "    # Get the rounded value of each input (as a python list)\n",
    "    values = [(name, case[name].round(10).tolist()) for name in inputs]\n",
    "    \n",
    "    # Print the input values for this case\n",
    "    print(f\"{case_num:>2}: {values}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert(sorted(inputs) == ['x', 'y1', 'y2', 'z'])\n",
    "\n",
    "assert(len(system_cases) == 14)\n",
    "\n",
    "assert_near_equal([case['y1'].item() for case in cr.get_cases('root.obj_cmp')],\n",
    "                  [25.6, 25.6, 8.33, 4.17, 3.30, 3.18, 3.16,\n",
    "                   3.16, 3.16, 3.16, 3.16, 3.16, 3.16, 3.16],\n",
    "                  tolerance=1e-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solver\n",
    "\n",
    "Similar to the `system` recorder, we can query the cases recorded by the `solver`. You can also\n",
    "access the values of inputs to the equation with the solver but in this case we'll focus on the\n",
    "coupling variables `y1` and `y2`.\n",
    "\n",
    "We'll use `list_cases` to get the cases recorded for `'root.cycle.nonlinear_solver'`, find out how many\n",
    "cases there are and plot a subset of the solver iterations (covering multiple driver iterations)\n",
    "to see the final convergence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import openmdao.api as om\n",
    "\n",
    "# Instantiate your CaseReader\n",
    "cr = om.CaseReader(\"cases.sql\")\n",
    "\n",
    "# Get list of cases, without displaying them\n",
    "solver_cases = cr.list_cases('root.cycle.nonlinear_solver', out_stream=None)\n",
    "\n",
    "# we will plot the last half of the iterations\n",
    "num_cases = len(solver_cases)\n",
    "sub_cases = num_cases//2\n",
    "iterations = np.arange(num_cases-sub_cases, num_cases)\n",
    "\n",
    "# Plot the final convergence of the two coupling variables\n",
    "y1_history = []\n",
    "y2_history = []\n",
    "\n",
    "for case_id in solver_cases[-sub_cases:]:\n",
    "    case = cr.get_case(case_id)\n",
    "    y1_history.append(case['y1'])\n",
    "    y2_history.append(case['y2'])\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(2, 1)\n",
    "\n",
    "ax1.plot(iterations, np.array(y1_history))\n",
    "ax1.set(ylabel='Coupling Output: y1', title='Solver History')\n",
    "ax1.grid()\n",
    "\n",
    "ax2.plot(iterations, np.array(y2_history))\n",
    "ax2.set(ylabel='Coupling Parameter: y2', xlabel='Iterations')\n",
    "ax2.grid()\n",
    "\n",
    "plt.show()\n",
    "\n",
    "# Get the final values\n",
    "print(f\"Number of cases: {len(solver_cases)}\")\n",
    "print(\"Final values:\")\n",
    "case = cr.get_case(solver_cases[-1])\n",
    "print('  y1:', case['y1'])\n",
    "print('  y2:', case['y2'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert(len(solver_cases) == 76)\n",
    "assert_near_equal(case['y1'], 3.16, 1e-8)\n",
    "assert_near_equal(case['y2'], 3.75527777, 1e-8)"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
