{
 "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  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sellar - A Two-Discipline Problem with a Nonlinear Solver"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the monodisciplinary tutorials, we built and optimized models comprised of only a single component.\n",
    "Now, we'll work through a slightly more complex problem that involves two disciplines, and hence two main components.\n",
    "You'll learn how to group components together into a larger model and how to use\n",
    "a [NonlinearBlockGaussSeidel](../../features/building_blocks/solvers/nonlinear_block_gs) nonlinear solver to converge a group with coupled components."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Sellar problem is a two-discipline toy problem with each discipline described by a single equation.\n",
    "There isn't any physical significance to the equations, it's just a compact example to use as a means of understanding\n",
    "simple coupled models.\n",
    "The output of each component feeds into the input of the other, which creates a coupled model that needs to\n",
    "be converged in order for the outputs to be valid.\n",
    "You can see the coupling between the two disciplines show up through the \\\\(y_1\\\\) and \\\\(y_2\\\\) variables in the following diagram that describes the problem structure: ![sellar example](images/sellar_xdsm.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Disciplinary Components\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following component definitions, there is a call to [declare_partials](../../features/core_features/working_with_derivatives/approximating_partial_derivatives) in the `setup_partials` method that looks like this:\n",
    "\n",
    "```\n",
    "self.declare_partials('*', '*', method='fd')\n",
    "```\n",
    "\n",
    "This command tells OpenMDAO to approximate all the partial derivatives of that component using finite difference.\n",
    "The default settings will use forward difference with an absolute step size of 1e-6, but you can change the [FD settings](../../features/core_features/working_with_derivatives/approximating_partial_derivatives) to work well for your component.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class SellarDis1(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Component containing Discipline 1 -- no derivatives version.\n",
    "    \"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Global Design Variable\n",
    "        self.add_input('z', val=np.zeros(2))\n",
    "\n",
    "        # Local Design Variable\n",
    "        self.add_input('x', val=0.)\n",
    "\n",
    "        # Coupling parameter\n",
    "        self.add_input('y2', val=1.0)\n",
    "\n",
    "        # Coupling output\n",
    "        self.add_output('y1', val=1.0)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        # Finite difference all partials.\n",
    "        self.declare_partials('*', '*', method='fd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        \"\"\"\n",
    "        Evaluates the equation\n",
    "        y1 = z1**2 + z2 + x1 - 0.2*y2\n",
    "        \"\"\"\n",
    "        z1 = inputs['z'][0]\n",
    "        z2 = inputs['z'][1]\n",
    "        x1 = inputs['x']\n",
    "        y2 = inputs['y2']\n",
    "\n",
    "        outputs['y1'] = z1**2 + z2 + x1 - 0.2*y2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SellarDis2(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Component containing Discipline 2 -- no derivatives version.\n",
    "    \"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "        # Global Design Variable\n",
    "        self.add_input('z', val=np.zeros(2))\n",
    "\n",
    "        # Coupling parameter\n",
    "        self.add_input('y1', val=1.0)\n",
    "\n",
    "        # Coupling output\n",
    "        self.add_output('y2', val=1.0)\n",
    "\n",
    "    def setup_partials(self):\n",
    "        # Finite difference all partials.\n",
    "        self.declare_partials('*', '*', method='fd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        \"\"\"\n",
    "        Evaluates the equation\n",
    "        y2 = y1**(.5) + z1 + z2\n",
    "        \"\"\"\n",
    "\n",
    "        z1 = inputs['z'][0]\n",
    "        z2 = inputs['z'][1]\n",
    "        y1 = inputs['y1']\n",
    "\n",
    "        # Note: this may cause some issues. However, y1 is constrained to be\n",
    "        # above 3.16, so lets just let it converge, and the optimizer will\n",
    "        # throw it out\n",
    "        if y1.real < 0.0:\n",
    "            y1 *= -1\n",
    "\n",
    "        outputs['y2'] = y1**.5 + z1 + z2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grouping and Connecting Components\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now want to build in OpenMDAO the model that is represented by the XDSM diagram above. We’ve defined the two disciplinary components, but there are still the three outputs of the model that need to be computed. Additionally, since we have the computations split up into multiple components, we need to group them all together and tell OpenMDAO how to pass data between them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SellarMDA(om.Group):\n",
    "    \"\"\"\n",
    "    Group containing the Sellar MDA.\n",
    "    \"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "        cycle = self.add_subsystem('cycle', om.Group(), promotes=['*'])\n",
    "        cycle.add_subsystem('d1', SellarDis1(), promotes_inputs=['x', 'z', 'y2'],\n",
    "                            promotes_outputs=['y1'])\n",
    "        cycle.add_subsystem('d2', SellarDis2(), promotes_inputs=['z', 'y1'],\n",
    "                            promotes_outputs=['y2'])\n",
    "\n",
    "        cycle.set_input_defaults('x', 1.0)\n",
    "        cycle.set_input_defaults('z', np.array([5.0, 2.0]))\n",
    "\n",
    "        # Nonlinear Block Gauss Seidel is a gradient free solver\n",
    "        cycle.nonlinear_solver = om.NonlinearBlockGS()\n",
    "\n",
    "        self.add_subsystem('obj_cmp', om.ExecComp('obj = x**2 + z[1] + y1 + exp(-y2)',\n",
    "                                                  z=np.array([0.0, 0.0]), x=0.0),\n",
    "                           promotes=['x', 'z', 'y1', 'y2', 'obj'])\n",
    "\n",
    "        self.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "        self.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're working with a new class: [Group](../../features/core_features/working_with_groups/main).\n",
    "This is the container that lets you build up complex model hierarchies.\n",
    "Groups can contain other groups, components, or combinations of groups and components.\n",
    "\n",
    "You can directly create instances of `Group` to work with, or you can subclass from it to define your own custom\n",
    "groups. We're doing both of these things above. First, we defined our own custom `Group` subclass called `SellarMDA`.\n",
    "In our run script, we created an instance of `SellarMDA` to actually run it.\n",
    "Then, inside the `setup` method of `SellarMDA` we're also working directly with a `Group` instance by adding the subsystem `cycle`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = om.Problem()\n",
    "cycle = prob.model.add_subsystem('cycle', om.Group(), promotes=['*'])\n",
    "cycle.add_subsystem('d1', SellarDis1(), promotes_inputs=['x', 'z', 'y2'], promotes_outputs=['y1'])\n",
    "cycle.add_subsystem('d2', SellarDis2(), promotes_inputs=['z', 'y1'], promotes_outputs=['y2'])\n",
    "\n",
    "# Nonlinear Block Gauss-Seidel is a gradient-free solver\n",
    "cycle.nonlinear_solver = om.NonlinearBlockGS()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our `SellarMDA` Group, when instantiated, will have a three-level hierarchy, with itself as the topmost level, as\n",
    "illustrated in the following figure: \n",
    "![sellar tree](images/sellar_tree.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Why do we create the *cycle* subgroup?\n",
    "There is a circular data dependency between *d1* and *d2* that needs to be converged with a nonlinear solver in order to get a valid answer.\n",
    "It's a bit more efficient to put these two components into their own subgroup, so that we can iteratively converge them by themselves,\n",
    "before moving on to the rest of the calculations in the model.\n",
    "Models with cycles in them are often referred to as \"Multidisciplinary Analyses\" or **MDA** for short.\n",
    "You can pick which kind of solver you would like to use to converge the MDA. The most common choices are:\n",
    "\n",
    "  - [NonlinearBlockGaussSeidel](../../features/building_blocks/solvers/nonlinear_block_gs)\n",
    "  - [NewtonSolver](../../features/building_blocks/solvers/newton)\n",
    "    \n",
    "The `NonlinearBlockGaussSeidel` solver, also sometimes called a \"fixed-point iteration solver,\" is a gradient-free method\n",
    "that works well in many situations.\n",
    "More tightly-coupled problems, or problems with instances of [ImplicitComponent](../../features/core_features/working_with_components/implicit_component) that don't implement their own `solve_nonlinear` method, will require the `Newton` solver.\n",
    "\n",
    "```{Note}\n",
    "OpenMDAO comes with other nonlinear solvers you can use if they suit your problem.\n",
    "Here is the complete list of OpenMDAO's [nonlinear solvers](../../features/building_blocks/solvers/nonlinear_solvers).\n",
    "```\n",
    "    \n",
    "The subgroup, named `cycle`, is useful here, because it contains the multidisciplinary coupling of the Sellar problem.\n",
    "This allows us to assign the nonlinear solver to `cycle` to just converge those two components, before moving on to the final\n",
    "calculations for the `obj_cmp`, `con_cmp1`, and `con_cmp2` to compute the actual outputs of the problem.\n",
    "\n",
    "## Promoting variables with the same name connects them\n",
    "\n",
    "The data connections in this model are made via promotion.\n",
    "OpenMDAO will look at each level of the hierarchy and connect all output-input pairs that have the same names. When an input is promoted on multiple components, you can use \"set_input_defaults\" to define the common initial value.\n",
    "\n",
    "\n",
    "## ExecComp is a helper component for quickly defining components for simple equations\n",
    "\n",
    "A lot of times in your models, you need to define a new variable as a simple function of other variables.\n",
    "OpenMDAO provides a helper component to make this easier, called [ExecComp](../../features/building_blocks/components/exec_comp).\n",
    "It's fairly flexible, allowing you to work with scalars or arrays, to work with units, and to call basic math\n",
    "functions (e.g. *sin* or *exp*).  We have used `ExecComp` in this model to calculate our\n",
    "objectives and constraints."
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
