{
 "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": [
    "# LinearBlockGS\n",
    "\n",
    "LinearBlockGS uses Block Gauss-Seidel to solve the linear system. LinearBlockGS iterates until the linear\n",
    "residual is below a tolerance, or the maximum number of iterations has been exceeded. As such,\n",
    "it is generally usable for any system topology, and can handle cycles and implicit states\n",
    "alike. It is not always the best solver to choose, however, and is known to diverge or plateau\n",
    "on some problems. In such a case, you may need to use a solver such as\n",
    "[ScipyKrylov](../../../_srcdocs/packages/solvers.linear/scipy_iter_solver).\n",
    "\n",
    "LinearBlockGS is a block solver, so you can specify different linear solvers in the subsystems and they\n",
    "will be utilized to solve the subsystem linear problem.\n",
    "\n",
    "Note that systems without cycles or implicit states will converge in one iteration of Block Gauss-Seidel.\n",
    "\n",
    "Here, we calculate the total derivatives across the Sellar system."
   ]
  },
  {
   "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_src24\", 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_src24`\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_src25\", 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_src25`\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.LinearBlockGS()\n",
    "model.nonlinear_solver = om.NonlinearBlockGS()\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",
    "wrt = ['z']\n",
    "of = ['obj']\n",
    "\n",
    "J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(J['obj', 'z'][0][0])\n",
    "print(J['obj', 'z'][0][1])"
   ]
  },
  {
   "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(J['obj', 'z'][0][0], 9.61001056, .00001)\n",
    "assert_near_equal(J['obj', 'z'][0][1], 1.78448534, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LinearBlockGS Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "om.show_options_table(\"openmdao.solvers.linear.linear_block_gs.LinearBlockGS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LinearBlockGS Constructor\n",
    "\n",
    "The call signature for the `LinearBlockGS` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.solvers.linear.linear_block_gs.LinearBlockGS.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## Aitken relaxation\n",
    "\n",
    "This solver implements Aitken relaxation, as described in Algorithm 1 of this paper on aerostructual design [optimization](https://mdolab.engin.umich.edu/bibliography/Kenway2014a). The relaxation is turned off by default, but it may help convergence for more tightly coupled models.\n",
    "\n",
    "\n",
    "## LinearBlockGS Option Examples\n",
    "\n",
    "**maxiter**\n",
    "\n",
    "  This lets you specify the maximum number of Gauss-Seidel 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 either of the specified absolute or relative tolerances."
   ]
  },
  {
   "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.nonlinear_solver = om.NonlinearBlockGS()\n",
    "\n",
    "model.linear_solver = om.LinearBlockGS()\n",
    "model.linear_solver.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",
    "wrt = ['z']\n",
    "of = ['obj']\n",
    "\n",
    "J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(J['obj', 'z'][0][0])\n",
    "print(J['obj', 'z'][0][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(J['obj', 'z'][0][0], 9.60230118004, .00001)\n",
    "assert_near_equal(J['obj', 'z'][0][1], 1.78022500547, .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 linear residuals is calculated by calling `apply_linear`. If this norm value is lower than the absolute 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.nonlinear_solver = om.NonlinearBlockGS()\n",
    "\n",
    "model.linear_solver = om.LinearBlockGS()\n",
    "model.linear_solver.options['atol'] = 1.0e-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",
    "wrt = ['z']\n",
    "of = ['obj']\n",
    "\n",
    "J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(J['obj', 'z'][0][0])\n",
    "print(J['obj', 'z'][0][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(J['obj', 'z'][0][0], 9.61016296175, .00001)\n",
    "assert_near_equal(J['obj', 'z'][0][1], 1.78456955704, .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 linear residuals is calculated by calling `apply_linear`. 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.nonlinear_solver = om.NonlinearBlockGS()\n",
    "\n",
    "model.linear_solver = om.LinearBlockGS()\n",
    "model.linear_solver.options['rtol'] = 1.0e-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",
    "wrt = ['z']\n",
    "of = ['obj']\n",
    "\n",
    "J = prob.compute_totals(of=of, wrt=wrt, return_format='flat_dict')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(J['obj', 'z'][0][0])\n",
    "print(J['obj', 'z'][0][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(J['obj', 'z'][0][0], 9.61016296175, .00001)\n",
    "assert_near_equal(J['obj', 'z'][0][1], 1.78456955704, .00001)"
   ]
  }
 ],
 "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.1"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
