{
 "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": [
    "# NewtonSolver\n",
    "\n",
    "NewtonSolver implements Newton's method to solve the system that contains it. This\n",
    "is the most general solver in OpenMDAO, in that it can solve any topology including cyclic\n",
    "connections and implicit states in the system or subsystems. Newton's method requires derivatives,\n",
    "so a linear solver can also be specified. By default, NewtonSolver uses the linear solver\n",
    "that is slotted in the containing system.\n",
    "\n",
    "OpenMDAO's implementation is very close to pure Newton's method, but adds line searches,\n",
    "bounds on state variables, and techniques to unstick the solver if it stalls. For an intuitive\n",
    "understanding of how Newton's method works and how it extends to multiple dimensions, take a look\n",
    "at the following video:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "from IPython.display import IFrame\n",
    "IFrame(\n",
    "    \"https://www.youtube-nocookie.com/embed/fuHl4bgCQTI\",\n",
    "    width=1,\n",
    "    height=1,\n",
    "    extras=[\n",
    "        \"title=\\\"YouTube video player\\\"\",\n",
    "        \"frameborder=\\\"0\\\"\",\n",
    "        \"allow=\\\"accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture\\\"\",\n",
    "        \"style=\\\"width: 85%; height: 100%; aspect-ratio: 16 / 9; margin-left: auto; margin-right: auto; display: block;\\\"\"\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "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_src29\", get_code(\"openmdao.test_suite.components.sellar.SellarDis1withDerivatives\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `SellarDis1withDerivatives` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src29`\n",
    ":::"
   ]
  },
  {
   "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_src30\", get_code(\"openmdao.test_suite.components.sellar.SellarDis2withDerivatives\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `SellarDis2withDerivatives` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src30`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.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=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "model.linear_solver = om.DirectSolver()\n",
    "\n",
    "model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 1.)\n",
    "prob.set_val('z', np.array([5.0, 2.0]))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('y1'))\n",
    "print(prob.get_val('y2'))"
   ]
  },
  {
   "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(prob.get_val('y1'), 25.58830273, .00001)\n",
    "assert_near_equal(prob.get_val('y2'), 12.05848819, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most of the solvers in OpenMDAO operate hierarchically, in that you can use solvers on subgroups\n",
    "to subdivide the calculation effort. However, NewtonSolver is an exception. It does not\n",
    "call `solve_nonlinear` on its subsystems, nor does it pass data along the connections. Instead,\n",
    "the Newton solver sets all inputs in all systems and subsystems that it contains, as it follows\n",
    "the gradient, driving the residuals to convergence.  After each iteration, the iteration count and the residual norm are\n",
    "checked to see if termination has been satisfied.\n",
    "\n",
    "NewtonSolver Options\n",
    "--------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "om.show_options_table(\"openmdao.solvers.nonlinear.newton.NewtonSolver\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NewtonSolver Constructor\n",
    "\n",
    "The call signature for the `NewtonSolver` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.solvers.nonlinear.newton.NewtonSolver.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## NewtonSolver Option Examples\n",
    "\n",
    "**maxiter**\n",
    "\n",
    "  `maxiter` lets you specify the maximum number of Newton iterations to apply. In this example, we\n",
    "  cut it back from the default, ten, down to two, so that it terminates a few iterations earlier and doesn't\n",
    "  reach the specified absolute or relative tolerance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.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=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "model.linear_solver = om.DirectSolver()\n",
    "\n",
    "newton = model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "newton.options['maxiter'] = 2\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 1.)\n",
    "prob.set_val('z', np.array([5.0, 2.0]))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('y1'))\n",
    "print(prob.get_val('y2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('y1'), 25.5878516779, .00001)\n",
    "assert_near_equal(prob.get_val('y2'), 12.0607416105, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**atol**\n",
    "\n",
    "  Here, we set the absolute tolerance to a looser value that will trigger an earlier termination. After\n",
    "  each iteration, the norm of the residuals is calculated by calling `apply_nonlinear` on all of the components.\n",
    "  If this norm value is lower than the absolute\n",
    "  tolerance `atol`, the iteration will terminate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.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=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "model.linear_solver = om.DirectSolver()\n",
    "\n",
    "newton = model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "newton.options['atol'] = 1e-4\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 1.)\n",
    "prob.set_val('z', np.array([5.0, 2.0]))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('y1'))\n",
    "print(prob.get_val('y2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('y1'), 25.5882856302, .00001)\n",
    "assert_near_equal(prob.get_val('y2'), 12.05848819, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**rtol**\n",
    "\n",
    "  Here, we set the relative tolerance to a looser value that will trigger an earlier termination. After\n",
    "  each iteration, the norm of the residuals is calculated by calling `apply_nonlinear` on all of the components.\n",
    "  If the ratio of the currently calculated norm to the\n",
    "  initial residual norm is lower than the relative tolerance `rtol`, the iteration will terminate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.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=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "model.linear_solver = om.DirectSolver()\n",
    "\n",
    "newton = model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "newton.options['rtol'] = 1e-3\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 1.)\n",
    "prob.set_val('z', np.array([5.0, 2.0]))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('y1'))\n",
    "print(prob.get_val('y2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('y1'), 25.5878516779, .00001)\n",
    "assert_near_equal(prob.get_val('y2'), 12.0607416105, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**solve_subsystems**\n",
    "\n",
    "  If you set this option to True, NewtonSolver will call `solve_nonlinear` on all of its subsystems. You can\n",
    "  use this to solve difficult multi-level problems by attaching solvers to subsystems. This assures that those\n",
    "  subsystems will already be in an internally solved state when the Newton solver goes to solve it.\n",
    "\n",
    "  This example shows two instances of the Sellar model, which we have connected together to form a larger cycle.\n",
    "  We specify a Newton solver in each Sellar subgroup as well as a top-level Newton solver, which we tell to solve\n",
    "  its subsystems."
   ]
  },
  {
   "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_src31\", get_code(\"openmdao.test_suite.components.double_sellar.DoubleSellar\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `DoubleSellar` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src31`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.double_sellar import DoubleSellar\n",
    "\n",
    "prob = om.Problem(model=DoubleSellar())\n",
    "model = prob.model\n",
    "\n",
    "g1 = model.g1\n",
    "g1.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "g1.nonlinear_solver.options['rtol'] = 1.0e-5\n",
    "g1.linear_solver = om.DirectSolver()\n",
    "\n",
    "g2 = model.g2\n",
    "g2.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "g2.nonlinear_solver.options['rtol'] = 1.0e-5\n",
    "g2.linear_solver = om.DirectSolver()\n",
    "\n",
    "model.nonlinear_solver = om.NewtonSolver()\n",
    "model.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "model.nonlinear_solver.options['solve_subsystems'] = True\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(prob.get_val('g1.y1'))\n",
    "print(prob.get_val('g1.y2'))\n",
    "print(prob.get_val('g2.y1'))\n",
    "print(prob.get_val('g2.y2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('g1.y1'), 0.64, .00001)\n",
    "assert_near_equal(prob.get_val('g1.y2'), 0.80, .00001)\n",
    "assert_near_equal(prob.get_val('g2.y1'), 0.64, .00001)\n",
    "assert_near_equal(prob.get_val('g2.y2'), 0.80, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**max_sub_solves**\n",
    "\n",
    "  This option is used in conjunction with the \"solve_subsystems\" option. It controls the number of iterations for which NewtonSolver will allow subsystems to solve themselves. When the iteration count exceeds `max_sub_solves`,  Newton returns to its default behavior.\n",
    "\n",
    "  For example, if you set `max_sub_solves` to zero, then the solvers on subsystems are executed during the initial\n",
    "  evaluation, but not during any subsequent iteration."
   ]
  },
  {
   "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_src32\", get_code(\"openmdao.test_suite.components.double_sellar.SubSellar\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `SubSellar` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src32`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.double_sellar import SubSellar\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('g1', SubSellar())\n",
    "model.add_subsystem('g2', SubSellar())\n",
    "\n",
    "model.connect('g1.y2', 'g2.x')\n",
    "model.connect('g2.y2', 'g1.x')\n",
    "\n",
    "# Converge the outer loop with Gauss Seidel, with a looser tolerance.\n",
    "model.nonlinear_solver = om.NewtonSolver()\n",
    "model.linear_solver = om.DirectSolver()\n",
    "\n",
    "g1 = model.g1\n",
    "g1.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "g1.nonlinear_solver.options['rtol'] = 1.0e-5\n",
    "g1.linear_solver = om.DirectSolver()\n",
    "\n",
    "g2 = model.g2\n",
    "g2.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "g2.nonlinear_solver.options['rtol'] = 1.0e-5\n",
    "g2.linear_solver = om.DirectSolver()\n",
    "\n",
    "model.nonlinear_solver = om.NewtonSolver()\n",
    "model.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "model.nonlinear_solver.options['solve_subsystems'] = True\n",
    "model.nonlinear_solver.options['max_sub_solves'] = 0\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**err_on_non_converge**\n",
    "\n",
    "  If you set this to True, then when the solver doesn't converge, either by hitting the iteration limit\n",
    "  without meeting the tolerance criteria, or by encountering a NaN or inf, it\n",
    "  will raise an AnalysisError exception. This is mainly important when coupled with a higher-level solver or\n",
    "  driver (e.g., `pyOptSparseDriver`)that can handle the AnalysisError by adapting the stepsize and retrying. This feature can be set on any iterative nonlinear or linear solver.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.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=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "model.linear_solver = om.DirectSolver()\n",
    "\n",
    "newton = model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "newton.options['maxiter'] = 1\n",
    "newton.options['err_on_non_converge'] = True\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 1.)\n",
    "prob.set_val('z', np.array([5.0, 2.0]))\n",
    "\n",
    "try:\n",
    "    prob.run_model()\n",
    "except om.AnalysisError:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**stall_limit, stall_tol, and stall_tol_type **\n",
    "\n",
    "  In some cases, nonlinear solvers can stall out where the norm of the residual stops changing at all. This\n",
    "  can happen for a couple of reasons. You can hit numerical noise problems and just be wandering around in\n",
    "  a circle, or you can get stuck on a bound and the line search just keeps running into the same spot no\n",
    "  matter what. Either way, if you have say 100 max iterations and you stall at 15 ... you waste a lot of\n",
    "  compute time. To remedy this, you can turn on stall detection in all nonlinear solvers by setting the\n",
    "  \"stall_limit\" option to a number greater than zero.\n",
    "\n",
    "  In this example, we set stall_limit to 3. While the solver iterates, it will compare the value of the\n",
    "  residual norm to the value computed in the previous iteration.  If the value matches for three iterations\n",
    "  in a row, then iteration will terminate due to detection of a stall. If \"err_on_non_converge\" is set\n",
    "  to True, then an ``AnalysisError`` will be raised just as if we had reached the iteration count limit.\n",
    "\n",
    "  We also set the `stall_tol` to 1e-6, which is the threshold below which a change in the relative residual\n",
    "  norm is considered to be unchanged.\n",
    "\n",
    "  The option `stall_tol_type` is used to specify whether the absolute or relative norm of the residual is\n",
    "  used to detect a stall condition. If an outer solver is used in conjunction with `solve_subsystems`, the\n",
    "  relative change in the residual norm will be small for those residuals converged by the inner solver.\n",
    "  In such cases, setting this option to 'abs' may more reliably detect a stall."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', om.ExecComp('y=3*x+1'), promotes=['*'])\n",
    "\n",
    "balance = prob.model.add_subsystem('balance', om.BalanceComp(),\n",
    "                                   promotes=['*'])\n",
    "balance.add_balance('x', lower=-.1, upper=10, rhs_val=0, lhs_name='y')\n",
    "\n",
    "newton = prob.model.nonlinear_solver = om.NewtonSolver()\n",
    "newton.options['solve_subsystems'] = True\n",
    "newton.options['stall_limit'] = 3\n",
    "newton.options['stall_tol'] = 1e-8\n",
    "newton.options['stall_tol_type'] = 'abs'\n",
    "newton.options['maxiter'] = 100\n",
    "\n",
    "prob.model.linear_solver = om.DirectSolver()\n",
    "\n",
    "prob.setup()\n",
    "prob.set_solver_print()\n",
    "\n",
    "prob.run_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying a Linear Solver\n",
    "\n",
    "We can choose a different linear solver for calculating the Newton step by setting the `linear_solver` attribute. The default is to use the linear solver that was specified on the containing system, which by default is LinearBlockGS. In the following example, we modify the model to use [DirectSolver](../../../_srcdocs/packages/solvers.linear/direct) instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.components.sellar import SellarDis1withDerivatives, SellarDis2withDerivatives\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('d1', SellarDis1withDerivatives(), promotes=['x', 'z', 'y1', 'y2'])\n",
    "model.add_subsystem('d2', SellarDis2withDerivatives(), promotes=['z', 'y1', 'y2'])\n",
    "\n",
    "model.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=['obj', 'x', 'z', 'y1', 'y2'])\n",
    "\n",
    "model.add_subsystem('con_cmp1', om.ExecComp('con1 = 3.16 - y1'), promotes=['con1', 'y1'])\n",
    "model.add_subsystem('con_cmp2', om.ExecComp('con2 = y2 - 24.0'), promotes=['con2', 'y2'])\n",
    "\n",
    "model.linear_solver = om.LinearBlockGS()\n",
    "\n",
    "newton = model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "\n",
    "newton.linear_solver = om.DirectSolver()\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', 1.)\n",
    "prob.set_val('z', np.array([5.0, 2.0]))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('y1'))\n",
    "print(prob.get_val('y2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('y1'), 25.58830273, .00001)\n",
    "assert_near_equal(prob.get_val('y2'), 12.05848819, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying a Line Search Algorithm\n",
    "\n",
    "NewtonSolver has a `linesearch` attribute, which supports specification of a supplemental algorithm that can find a better point along the Newton search direction. This is typically used for cases where we have declared upper\n",
    "or lower bounds on some of the model outputs and we want to prevent Newton from moving into this non-feasible space during iteration. An algorithm that does this is called a line search.\n",
    "\n",
    "By default, NewtonSolver does not perform a line search. We will show how to specify one. First,\n",
    "let's set up a problem that has implicit bounds on one of its states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import atan\n",
    "\n",
    "class CompAtan(om.ImplicitComponent):\n",
    "    \"\"\"\n",
    "    A simple implicit component with the following equation:\n",
    "\n",
    "    F(x, y) = 33.0 * atan(y-20)**2 + x\n",
    "\n",
    "    x is an input, y is the state to be solved.\n",
    "    for x = -100, y should be 19.68734033\n",
    "\n",
    "    This equation poses a challenge because a guess that is far from the solution yields large\n",
    "    gradients and divergence. Additionally, the jacobian becomes singular at y = 20. To address\n",
    "    this, a lower and upper bound are added on y so that a solver with a BoundsEnforceLS does not\n",
    "    allow it to stray into problematic regions.\n",
    "    \"\"\"\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', 1.0)\n",
    "        self.add_output('y', 1.0, lower=1.0, upper=19.9)\n",
    "\n",
    "        self.declare_partials(of='y', wrt='x')\n",
    "        self.declare_partials(of='y', wrt='y')\n",
    "\n",
    "    def apply_nonlinear(self, inputs, outputs, residuals):\n",
    "        x = inputs['x']\n",
    "        y = outputs['y']\n",
    "\n",
    "        residuals['y'] = (33.0 * atan(y-20.0))**2 + x\n",
    "\n",
    "    def linearize(self, inputs, outputs, jacobian):\n",
    "        # x = inputs['x']  # x is not needed\n",
    "        y = outputs['y']\n",
    "\n",
    "        jacobian['y', 'y'] = 2178.0*atan(y-20.0) / (y**2 - 40.0*y + 401.0)\n",
    "        jacobian['y', 'x'] = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This equation poses a challenge because a guess that is far from the solution yields large\n",
    "gradients and the solution will diverge. Additionally, the jacobian becomes singular at y = 20. To address\n",
    "both of these problems, a lower and upper bound are added on y so that a solver with a BoundsEnforceLS does not\n",
    "allow it to stray into problematic regions. Without the linsearch, Newton is unable to solve this problem unless you start very close to the solution.\n",
    "\n",
    "Here, we specify [BoundsEnforceLS](../../../_srcdocs/packages/solvers.linesearch/backtracking)\n",
    "as our line search algorithm, and we get the expected solution for \"y\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.solvers.linesearch.tests.test_backtracking import CompAtan\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', CompAtan(), promotes_inputs=['x'])\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('x', -100.0)\n",
    "\n",
    "# Initial value for the state:\n",
    "prob.set_val('comp.y', 12.0)\n",
    "\n",
    "# You can change the om.NewtonSolver settings after setup is called\n",
    "newton = prob.model.nonlinear_solver = om.NewtonSolver()\n",
    "prob.model.linear_solver = om.DirectSolver()\n",
    "newton.options['iprint'] = 2\n",
    "newton.options['rtol'] = 1e-8\n",
    "newton.options['solve_subsystems'] = True\n",
    "\n",
    "newton.linesearch = om.BoundsEnforceLS()\n",
    "newton.linesearch.options['iprint'] = 2\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('comp.y'), 19.68734033, 1e-6)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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,
  "vscode": {
   "interpreter": {
    "hash": "9c92b2045504bea896bda7910cf76807d15460aca4268ad16ab9067dbd37dd98"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
