{
 "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": [
    "# Providing an Initial Guess for Implicit States in a Group\n",
    "\n",
    "In the documentation for [ImplicitComponent](../working_with_components/implicit_component),\n",
    "you saw that you can provide an initial guess for implicit states within the\n",
    "component using it's *guess_nonlinear* method.\n",
    "\n",
    "`Group` also provides a *guess_nonlinear* method in which you can supply\n",
    "the starting value for implicit state variables in any of it's subsystems.\n",
    "\n",
    "The following example demonstrates the capability of setting the initial guess\n",
    "value at the group level, using the input from one component to compute the guess\n",
    "for another.  In this case, a `Discipline` group solves a system of equations\n",
    "using a [BalanceComp](../../building_blocks/components/balance_comp). It answers the question:\n",
    "\"What is $x$ such that $x^2$ is equal to twice our input value\".\n",
    "\n",
    "![Group guess_nonlinear Example](images/guess_example.png)\n",
    "\n",
    "Given our knowledge of the relationship between the two equations, we can supply\n",
    "an initial guess for the implicit state variable, $x$, that makes it\n",
    "unnecessary for the solver to iterate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class Discipline(om.Group):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_subsystem('comp0', om.ExecComp('y=x**2'))\n",
    "        self.add_subsystem('comp1', om.ExecComp('z=2*external_input'),\n",
    "                           promotes_inputs=['external_input'])\n",
    "\n",
    "        self.add_subsystem('balance', om.BalanceComp('x', lhs_name='y', rhs_name='z'),\n",
    "                           promotes_outputs=['x'])\n",
    "\n",
    "        self.connect('comp0.y', 'balance.y')\n",
    "        self.connect('comp1.z', 'balance.z')\n",
    "\n",
    "        self.connect('x', 'comp0.x')\n",
    "\n",
    "        self.nonlinear_solver = om.NewtonSolver(iprint=2, solve_subsystems=True)\n",
    "        self.linear_solver = om.DirectSolver()\n",
    "\n",
    "    def guess_nonlinear(self, inputs, outputs, residuals):\n",
    "        # Check residuals\n",
    "        if np.abs(residuals['x']) > 1.0E-2:\n",
    "            # inputs are addressed using full path name, regardless of promotion\n",
    "            external_input = inputs['comp1.external_input']\n",
    "\n",
    "            # balance drives x**2 = 2*external_input\n",
    "            x_guess = (2*external_input)**.5\n",
    "\n",
    "            # outputs are addressed by the their promoted names\n",
    "            outputs['x'] = x_guess # perfect guess should converge in 0 iterations\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('discipline', Discipline(), promotes_inputs=['external_input'])\n",
    "\n",
    "p.setup()\n",
    "p.set_val('external_input', 1.)\n",
    "p.run_model()\n",
    "\n",
    "print(p.model.nonlinear_solver._iter_count)\n",
    "\n",
    "print(p.get_val('discipline.x'))"
   ]
  },
  {
   "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(p.model.nonlinear_solver._iter_count == 0)\n",
    "assert_near_equal(p.get_val('discipline.x'), 1.41421356, 1e-6)"
   ]
  }
 ],
 "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"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
