{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "upset-transaction",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "allied-means",
   "metadata": {},
   "source": [
    "# Moving Some Systems Outside of the Optimization Loop\n",
    "\n",
    "Sometimes a model contains systems that don't need to be executed during each iteration of an\n",
    "optimization loop.  If those systems are expensive to execute, leaving them out of the optimization\n",
    "loop can significantly improve performance.\n",
    "\n",
    "If the `Problem` option `group_by_pre_opt_post` is True, OpenMDAO will separate all of the \n",
    "Components in the model into `pre`, `optimization`, or `post` sets based on the dependency graph. \n",
    "Only those Components found in\n",
    "the `optimization` set will be executed during each iteration of the optimization, while the ones\n",
    "in the `pre` and `post` sets will only execute before or after the optimization loop.\n",
    "\n",
    "It's also possible, by setting the `always_opt` option on a component, to force that component\n",
    "to be included in the optimization loop regardless of its dependencies.  Note that depending\n",
    "on what its dependecies actually are, moving it to the optimization loop may also force other\n",
    "components to move there even if they would normally be in `pre` or `post`.\n",
    "\n",
    "## Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73798c74",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.exec_comp_for_test import ExecComp4Test\n",
    "\n",
    "size = 3\n",
    "\n",
    "prob = om.Problem(group_by_pre_opt_post=True)\n",
    "prob.driver = om.ScipyOptimizeDriver(optimizer='SLSQP', disp=False)\n",
    "\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('pre1', ExecComp4Test('y=2.*x', x=np.ones(size), y=np.zeros(size)))\n",
    "model.add_subsystem('pre2', ExecComp4Test('y=3.*x', x=np.ones(size), y=np.zeros(size)))\n",
    "\n",
    "model.add_subsystem('iter1', ExecComp4Test('y=x1 + x2*4. + x3',\n",
    "                                            x1=np.ones(size), x2=np.ones(size),\n",
    "                                            x3=np.ones(size), y=np.zeros(size)))\n",
    "model.add_subsystem('iter2', ExecComp4Test('y=.5*x', x=np.ones(size), y=np.zeros(size)))\n",
    "model.add_subsystem('iter4', ExecComp4Test('y=7.*x', x=np.ones(size), y=np.zeros(size)))\n",
    "model.add_subsystem('iter3', ExecComp4Test('y=6.*x', x=np.ones(size), y=np.zeros(size)))\n",
    "\n",
    "model.add_subsystem('post1', ExecComp4Test('y=8.*x', x=np.ones(size), y=np.zeros(size)))\n",
    "model.add_subsystem('post2', ExecComp4Test('y=x1*9. + x2*5', x1=np.ones(size),\n",
    "                                            x2=np.ones(size), y=np.zeros(size)))\n",
    "\n",
    "# we don't want ExecComps to be colored because it makes the iter counting more complicated\n",
    "for comp in model.system_iter(typ=ExecComp4Test):\n",
    "    comp.options['do_coloring'] = False\n",
    "    comp.options['has_diag_partials'] = True\n",
    "\n",
    "\n",
    "model.connect('pre1.y', ['iter1.x1', 'post2.x1'])\n",
    "model.connect('pre2.y', 'iter1.x2')\n",
    "model.connect('iter1.y', ['iter2.x', 'iter4.x'])\n",
    "model.connect('iter2.y', 'post2.x2')\n",
    "model.connect('iter3.y', 'post1.x')\n",
    "model.connect('iter4.y', 'iter3.x')\n",
    "\n",
    "prob.model.add_design_var('iter1.x3', lower=0, upper=10)\n",
    "prob.model.add_constraint('iter2.y', upper=10.)\n",
    "prob.model.add_objective('iter3.y', index=0)\n",
    "\n",
    "prob.setup()\n",
    "prob.run_driver()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bea1bf63",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "allow-assert"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_check_totals\n",
    "\n",
    "assert prob.model.pre1.num_nl_solves == 1\n",
    "assert prob.model.pre2.num_nl_solves == 1\n",
    "\n",
    "assert prob.model.iter1.num_nl_solves == 3\n",
    "assert prob.model.iter2.num_nl_solves == 3\n",
    "assert prob.model.iter3.num_nl_solves == 3\n",
    "assert prob.model.iter4.num_nl_solves == 3\n",
    "\n",
    "assert prob.model.post1.num_nl_solves == 1\n",
    "assert prob.model.post2.num_nl_solves == 1\n",
    "\n",
    "data = prob.check_totals(out_stream=None)\n",
    "assert_check_totals(data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "522e97b4",
   "metadata": {},
   "source": [
    "## Check pre/post via the command line\n",
    "\n",
    "Use the `openmdao list_pre_post` command to view the sets of components found in the `pre` and `post` sets.  Any components not shown will be found in the `optimization` set and will run during each iteration of the optimization loop.\n",
    "Running the command on a script containing the example model above would show the following:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "17ac1495",
   "metadata": {},
   "source": [
    "```\n",
    "    Pre-optimization systems:\n",
    "        pre1\n",
    "        pre2\n",
    "\n",
    "    Post-optimization systems:\n",
    "        post1\n",
    "        post2\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "115f00b0",
   "metadata": {},
   "source": [
    "Note that the `pre` and `post` sets will only be shown if the `driver` is an optimizer and the `group_by_pre_opt_post` option of the `Problem` is True."
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "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.11.4"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
