{
 "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": [
    "# ImplicitComponent\n",
    "\n",
    "Implicit variables are those that are computed as an implicit function of other variables. For instance, $y$ would be an implicit variable, given that it is computed by solving $cos(x⋅y)−z⋅y=0$. In OpenMDAO, implicit variables are defined as the outputs of components that inherit from [ImplicitComponent](../../../_srcdocs/packages/core/implicitcomponent).\n",
    "\n",
    "In the above implicit expression, $y$ is the implicit variable while $x$ and $z$ would be considered inputs.\n",
    "\n",
    "## ImplicitComponent Methods\n",
    "\n",
    "The implementation of each method will be illustrated using a simple implicit component that computes the output $x$\n",
    "implicitly via a quadratic equation, $ax^2+bx+c=0$, where $a$, $b$, and $c$ are inputs to the component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class QuadraticComp(om.ImplicitComponent):\n",
    "    \"\"\"\n",
    "    A Simple Implicit Component representing a Quadratic Equation.\n",
    "\n",
    "    R(a, b, c, x) = ax^2 + bx + c\n",
    "\n",
    "    Solution via Quadratic Formula:\n",
    "    x = (-b + sqrt(b^2 - 4ac)) / 2a\n",
    "    \"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `setup()` :\n",
    "\n",
    "Declare input and output variables via `add_input` and `add_output`. Information like variable names, sizes, units, and bounds are declared. Also, declare partial derivatives that this component provides. Here we use the wild card to say that this component provides derivatives of all implicit residuals with respect to all inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup(self):\n",
    "    self.add_input('a', val=1., tags=['tag_a'])\n",
    "    self.add_input('b', val=1.)\n",
    "    self.add_input('c', val=1.)\n",
    "    self.add_output('x', val=0., tags=['tag_x'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `setup_partials()`:\n",
    "\n",
    "Declare partial derivatives that this component provides. Here we use the wild card to say that this component provides derivatives of all implicit residuals with respect to all inputs and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def setup_partials(self):\n",
    "    self.declare_partials(of='*', wrt='*')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `apply_nonlinear(inputs, outputs, residuals)` :\n",
    "\n",
    "Compute the `residuals`, given the `inputs` and `outputs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_nonlinear(self, inputs, outputs, residuals):\n",
    "    a = inputs['a']\n",
    "    b = inputs['b']\n",
    "    c = inputs['c']\n",
    "    x = outputs['x']\n",
    "    residuals['x'] = a * x ** 2 + b * x + c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `solve_nonlinear(inputs, outputs)` :\n",
    "\n",
    "[Optional] Compute the `outputs`, given the `inputs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_nonlinear(self, inputs, outputs):\n",
    "    a = inputs['a']\n",
    "    b = inputs['b']\n",
    "    c = inputs['c']\n",
    "    outputs['x'] = (-b + (b ** 2 - 4 * a * c) ** 0.5) / (2 * a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `linearize(inputs, outputs, partials)` :\n",
    "\n",
    "[Optional] The component’s partial derivatives of interest are those of the `residuals` with respect to the `inputs` and the `outputs`. If the user computes the partial derivatives explicitly, they are provided here. If the user wants to implement partial derivatives in a matrix-free way, this method provides a place to perform any necessary assembly or pre-processing for the matrix-vector products. Regardless of how the partial derivatives are computed, this method provides a place to perform any relevant factorizations for directly solving or preconditioning the linear system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def linearize(self, inputs, outputs, partials):\n",
    "    a = inputs['a']\n",
    "    b = inputs['b']\n",
    "    # c = inputs['c']  # not needed\n",
    "    x = outputs['x']\n",
    "\n",
    "    partials['x', 'a'] = x ** 2\n",
    "    partials['x', 'b'] = x\n",
    "    partials['x', 'c'] = 1.0\n",
    "    partials['x', 'x'] = 2 * a * x + b\n",
    "\n",
    "    self.inv_jac = 1.0 / (2 * a * x + b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `apply_linear(inputs, outputs, d_inputs, d_outputs, d_residuals, mode)` :\n",
    "\n",
    "[Optional] If the user wants to implement partial derivatives in a matrix-free way, this method performs the matrix-vector product. If `mode` is ‘fwd’, this method computes $d\\_residuals=J⋅[d\\_inputs \\ d\\_outputs]^T$\n",
    ". If `mode` is ‘rev’, this method computes $[d\\_inputs \\ d\\_outputs]^T=J^T⋅d\\_residuals$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_linear(self, inputs, outputs,\n",
    "                 d_inputs, d_outputs, d_residuals, mode):\n",
    "    a = inputs['a']\n",
    "    b = inputs['b']\n",
    "    # c = inputs['c']  # not needed\n",
    "    x = outputs['x']\n",
    "    if mode == 'fwd':\n",
    "        if 'x' in d_residuals:\n",
    "            if 'x' in d_outputs:\n",
    "                d_residuals['x'] += (2 * a * x + b) * d_outputs['x']\n",
    "            if 'a' in d_inputs:\n",
    "                d_residuals['x'] += x ** 2 * d_inputs['a']\n",
    "            if 'b' in d_inputs:\n",
    "                d_residuals['x'] += x * d_inputs['b']\n",
    "            if 'c' in d_inputs:\n",
    "                d_residuals['x'] += d_inputs['c']\n",
    "    elif mode == 'rev':\n",
    "        if 'x' in d_residuals:\n",
    "            if 'x' in d_outputs:\n",
    "                d_outputs['x'] += (2 * a * x + b) * d_residuals['x']\n",
    "            if 'a' in d_inputs:\n",
    "                d_inputs['a'] += x ** 2 * d_residuals['x']\n",
    "            if 'b' in d_inputs:\n",
    "                d_inputs['b'] += x * d_residuals['x']\n",
    "            if 'c' in d_inputs:\n",
    "                d_inputs['c'] += d_residuals['x']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `solve_linear(d_outputs, d_residuals, mode)` :\n",
    "\n",
    "[Optional] Solves a linear system where the matrix is $d\\_residuals/d\\_outputs$\n",
    "or its transpose. If `mode` is ‘fwd’, the right-hand side vector is $d\\_residuals$ and the solution vector is $d\\_outputs$. If `mode` is ‘rev’, the right-hand side vector is $d\\_outputs$ and the solution vector is $d\\_residuals$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_linear(self, d_outputs, d_residuals, mode):\n",
    "    if mode == 'fwd':\n",
    "        d_outputs['x'] = self.inv_jac * d_residuals['x']\n",
    "    elif mode == 'rev':\n",
    "        d_residuals['x'] = self.inv_jac * d_outputs['x']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `guess_nonlinear(self, inputs, outputs, residuals)` :\n",
    "\n",
    "[Optional] This method allows the user to calculate and specify an initial guess for implicit states. It is called at the start of the solve loop from certain nonlinear solvers (i.e. NewtonSolver and BroydenSolver), so it is useful for when you would like to “reset” the initial conditions on an inner-nested solve whenever an outer loop solver or driver changes other values. However if you include this function without checking the residuals, you will be starting the solvers from scratch during each *run_model* during an optimization. Add a check of your residuals and you can skip it when the system is converged.\n",
    "\n",
    "Since it is a hook for custom code, you could also use it to monitor variables in the input, output, or residual vectors and change the initial guess when some condition is met.\n",
    "\n",
    "Here is a simple example where we use NewtonSolver to find one of the roots of a second-order quadratic equation. Which root you get depends on the initial guess."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{Note}\n",
    "When a solver is restarted from the last successful solve due to having the `restart_from_sucessful` option set to `True`, the `guess_nonlinear` function will not be called.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "import numpy as np\n",
    "\n",
    "class ImpWithInitial(om.ImplicitComponent):\n",
    "    \"\"\"\n",
    "    An implicit component to solve the quadratic equation: x^2 - 4x + 3\n",
    "    (solutions at x=1 and x=3)\n",
    "    \"\"\"\n",
    "    def setup(self):\n",
    "        self.add_input('a', val=1.)\n",
    "        self.add_input('b', val=-4.)\n",
    "        self.add_input('c', val=3.)\n",
    "\n",
    "        self.add_output('x', val=0.)\n",
    "\n",
    "        self.declare_partials(of='*', wrt='*')\n",
    "\n",
    "    def apply_nonlinear(self, inputs, outputs, residuals):\n",
    "        a = inputs['a']\n",
    "        b = inputs['b']\n",
    "        c = inputs['c']\n",
    "        x = outputs['x']\n",
    "        residuals['x'] = a * x ** 2 + b * x + c\n",
    "\n",
    "    def linearize(self, inputs, outputs, partials):\n",
    "        a = inputs['a']  # not needed\n",
    "        b = inputs['b']\n",
    "        # c = inputs['c']  # not needed\n",
    "        x = outputs['x']\n",
    "\n",
    "        partials['x', 'a'] = x ** 2\n",
    "        partials['x', 'b'] = x\n",
    "        partials['x', 'c'] = 1.0\n",
    "        partials['x', 'x'] = 2 * a * x + b\n",
    "\n",
    "    def guess_nonlinear(self, inputs, outputs, resids):\n",
    "        # Check residuals\n",
    "        if np.abs(resids['x']) > 1.0E-2:\n",
    "            # Default initial state of zero for x takes us to x=1 solution.\n",
    "            # Here we set it to a value that will take us to the x=3 solution.\n",
    "            outputs['x'] = 5.0\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "model.add_subsystem('comp', ImpWithInitial())\n",
    "\n",
    "model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "model.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "prob.setup()\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.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",
    "assert_near_equal(prob['comp.x'], 3.)"
   ]
  }
 ],
 "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.3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
