{
 "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": [
    "# Building Models with Solvers and Implicit Components\n",
    "This tutorial will show you how to define implicit components and build models with them.\n",
    "We'll use a nonlinear circuit analysis example problem.\n",
    "\n",
    "## Circuit analysis\n",
    "Consider a simple electrical circuit made up from two resistors, a diode, and a constant current source.\n",
    "Our goal is to solve for the steady-state voltages at node 1 and node 2.\n",
    "\n",
    "![circuit_diagram](images/circuit_diagram.png)\n",
    "\n",
    "In order to find the voltages, we'll employ [Kirchoff's current law](https://en.wikipedia.org/wiki/Kirchhoff%27s_circuit_laws),\n",
    "and solve for the voltages needed at each node to drive the net current to 0.\n",
    "\n",
    "This means that the voltages at each node are *state variables* for the analysis.\n",
    "In other words, V1 and V2 are defined implicitly by the following residual equation:\n",
    "\n",
    "${R_{node_j}} = \\sum_k I_{k}^{in} - \\sum_k I_{k}^{out} = 0 .$\n",
    "\n",
    "To build this model we're going to define three different components:\n",
    "\n",
    "1. Resistor (Explicit)\n",
    "2. Diode (Explicit)\n",
    "3. Node (Implicit)\n",
    "\n",
    "ExplicitComponents - Resistor and Diode\n",
    "\n",
    "The `Resistor` and `Diode` components will each compute their current, given the voltages on either side.\n",
    "These calculations are analytic functions, so we'll inherit from [ExplicitComponent](../../features/core_features/working_with_components/explicit_component.ipynb). These components will each declare some options to allow you to pass in the relevant physical constants, and to allow you to give some reasonable default values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class Resistor(om.ExplicitComponent):\n",
    "    \"\"\"Computes current across a resistor using Ohm's law.\"\"\"\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('R', default=1., desc='Resistance in Ohms')\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('V_in', units='V')\n",
    "        self.add_input('V_out', units='V')\n",
    "        self.add_output('I', units='A')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('I', 'V_in', method='fd')\n",
    "        self.declare_partials('I', 'V_out', method='fd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        deltaV = inputs['V_in'] - inputs['V_out']\n",
    "        outputs['I'] = deltaV / self.options['R']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Diode(om.ExplicitComponent):\n",
    "    \"\"\"Computes current across a diode using the Shockley diode equation.\"\"\"\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('Is', default=1e-15, desc='Saturation current in Amps')\n",
    "        self.options.declare('Vt', default=.025875, desc='Thermal voltage in Volts')\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('V_in', units='V')\n",
    "        self.add_input('V_out', units='V')\n",
    "        self.add_output('I', units='A')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('I', 'V_in', method='fd')\n",
    "        self.declare_partials('I', 'V_out', method='fd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        deltaV = inputs['V_in'] - inputs['V_out']\n",
    "        Is = self.options['Is']\n",
    "        Vt = self.options['Vt']\n",
    "        outputs['I'] = Is * (np.exp(deltaV / Vt) - 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{important}\n",
    "Since we've provided default values for the options, they won't be required arguments when instantiating `Resistor` or `Diode`. Check out the [Features](../../features/features.ipynb) section for more details on how to use [component options](../../features/core_features/options/options.ipynb).\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ImplicitComponent - Node\n",
    "\n",
    "The `Node` component inherits from [ImplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb), which has a different interface than [ExplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb).\n",
    "Rather than compute the values of its outputs, it computes residuals via the `apply_nonlinear` method.\n",
    "When those residuals have been driven to zero, the values of the outputs will be implicitly known.\n",
    "`apply_nonlinear` computes the `residuals` using values from `inputs` and `outputs`.\n",
    "Notice that we still define *V* as an output of the `Node` component, albeit one that is implicitly defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node(om.ImplicitComponent):\n",
    "    \"\"\"Computes voltage residual across a node based on incoming and outgoing current.\"\"\"\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('n_in', default=1, types=int, desc='number of connections with + assumed in')\n",
    "        self.options.declare('n_out', default=1, types=int, desc='number of current connections + assumed out')\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_output('V', val=5., units='V')\n",
    "\n",
    "        for i in range(self.options['n_in']):\n",
    "            i_name = 'I_in:{}'.format(i)\n",
    "            self.add_input(i_name, units='A')\n",
    "\n",
    "        for i in range(self.options['n_out']):\n",
    "            i_name = 'I_out:{}'.format(i)\n",
    "            self.add_input(i_name, units='A')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        #note: we don't declare any partials wrt `V` here,\n",
    "        #      because the residual doesn't directly depend on it\n",
    "        self.declare_partials('V', 'I*', method='fd')\n",
    "\n",
    "    def apply_nonlinear(self, inputs, outputs, residuals):\n",
    "        residuals['V'] = 0.\n",
    "        for i_conn in range(self.options['n_in']):\n",
    "            residuals['V'] += inputs['I_in:{}'.format(i_conn)]\n",
    "        for i_conn in range(self.options['n_out']):\n",
    "            residuals['V'] -= inputs['I_out:{}'.format(i_conn)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Every state variable must have exactly one corresponding residual which is defined in the `apply_nonlinear` method. The `residuals` equations in an implicit component are not analogous to the `outputs` equations in the `compute` method of an explicit component. Instead of defining an explicit equation for the output, `residuals['example_output']` defines an equation for the residual *associated with* the output (state variable) `example_output`. In our example, `residuals['V']` defines the equation of the residual associated with the state variable *V*. There will be no explicit\n",
    "equation defining *V*, instead, the residual equation sums the currents associated with *V* so the sum can be driven to zero.\n",
    "\n",
    "An implicit component varies its outputs (state variables, in this case *V*) to drive the residual equation to zero. In our model, *V* does not show up directly in the residual equation. Instead, our explicit components `Resistor` and `Diode` create a dependence of the currents on *V*, so by using a solver on a higher level of the model hierarchy, we can vary *V* to have an effect on current, and we can drive the residuals to zero.\n",
    "\n",
    "All implicit components must define the `apply_nonlinear` method, but it is not a requirement that every [ImplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb)  define the `solve_nonlinear` method. (The `solve_nonlinear` method provides a way to explicitly define an output within an implicit component.) In fact, for the `Node` component, it is not even possible to define a `solve_nonlinear` because *V* does not show up directly in the residual function. So the implicit function represented by instances of the `Node` component must be converged at a higher level in the model hierarchy.\n",
    "\n",
    "There are cases where it is possible, and even advantageous, to define the `solve_nonlinear` method.\n",
    "For example, when a component is performing an engineering analysis with its own specialized nonlinear solver routines (e.g. CFD or FEM), then it makes sense to expose those to OpenMDAO via `solve_nonlinear` so OpenMDAO can make use of them. Just remember that `apply_nonlinear` must be defined, regardless of whether you also define `solve_nonlinear`.\n",
    "\n",
    "```{important}\n",
    "In this case, the residual equation is not a direct function of the state variable *V*. Often, however, the residual might be a direct function of one or more output variables. If that is the case, you can access the values via `outputs['V']`. See the [ImplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb) documentation for an example of this.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Circuit Group and Solving It with NewtonSolver\n",
    "We can combine the `Resistor`, `Diode`, and `Node` into the circuit pictured above using a [Group](../../features/core_features/working_with_groups/main.ipynb). Adding components and connecting their variables is the same as what you've seen before in the [Sellar - Two Discipline](../../basic_user_guide/multidisciplinary_optimization/sellar.ipynb) tutorial. What is new here is the additional use of the nonlinear [NewtonSolver](../../features/building_blocks/solvers/newton.ipynb) and linear [DirectSolver](../../features/building_blocks/solvers/direct_solver.ipynb) to converge the system.\n",
    "\n",
    "In previous tutorials, we used a gradient-free [NonlinearBlockGaussSeidel](../../features/building_blocks/solvers/nonlinear_block_gs.ipynb) solver, but that won't work here. Just above, we discussed that the `Node` class does not, and in fact can not, define its own `solve_nonlinear` method. Hence, there would be no calculations for the GaussSeidel solver to iterate on.\n",
    "Instead we use the Newton solver at the `Circuit` level, which uses Jacobian information to compute group level updates for all the variables simultaneously. The Newton solver's use of that Jacobian information is why we need to declare a linear solver in this case.\n",
    "\n",
    "```{important}\n",
    "OpenMDAO provides a library of [linear solvers](linear-solvers) that are useful in different advanced scenarios. For many problems, especially problems built from components with mostly scalar variables, the [DirectSolver](../../features/building_blocks/solvers/direct_solver.ipynb) will be both the most efficient and the easiest to use. We recommend you stick with [DirectSolver](../../features/building_blocks/solvers/direct_solver.ipynb) unless you have a good reason to switch.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "class Circuit(om.Group):\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_subsystem('n1', Node(n_in=1, n_out=2), promotes_inputs=[('I_in:0', 'I_in')])\n",
    "        self.add_subsystem('n2', Node())  # leaving defaults\n",
    "\n",
    "        self.add_subsystem('R1', Resistor(R=100.), promotes_inputs=[('V_out', 'Vg')])\n",
    "        self.add_subsystem('R2', Resistor(R=10000.))\n",
    "        self.add_subsystem('D1', Diode(), promotes_inputs=[('V_out', 'Vg')])\n",
    "\n",
    "        self.connect('n1.V', ['R1.V_in', 'R2.V_in'])\n",
    "        self.connect('R1.I', 'n1.I_out:0')\n",
    "        self.connect('R2.I', 'n1.I_out:1')\n",
    "\n",
    "        self.connect('n2.V', ['R2.V_out', 'D1.V_in'])\n",
    "        self.connect('R2.I', 'n2.I_in:0')\n",
    "        self.connect('D1.I', 'n2.I_out:0')\n",
    "\n",
    "        self.nonlinear_solver = om.NewtonSolver(solve_subsystems=False)\n",
    "        self.nonlinear_solver.options['iprint'] = 2\n",
    "        self.nonlinear_solver.options['maxiter'] = 20\n",
    "        self.linear_solver = om.DirectSolver()\n",
    "\n",
    "p = om.Problem()\n",
    "model = p.model\n",
    "\n",
    "model.add_subsystem('ground', om.IndepVarComp('V', 0., units='V'))\n",
    "model.add_subsystem('source', om.IndepVarComp('I', 0.1, units='A'))\n",
    "model.add_subsystem('circuit', Circuit())\n",
    "\n",
    "model.connect('source.I', 'circuit.I_in')\n",
    "model.connect('ground.V', 'circuit.Vg')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "# set some initial guesses\n",
    "p['circuit.n1.V'] = 10.\n",
    "p['circuit.n2.V'] = 1.\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "print(p['circuit.n1.V'])\n",
    "print(p['circuit.n2.V'])\n",
    "print(p['circuit.R1.I'])\n",
    "print(p['circuit.R2.I'])\n",
    "print(p['circuit.D1.I'])\n",
    "\n",
    "# sanity check: should sum to .1 Amps\n",
    "print(p['circuit.R1.I'] + p['circuit.D1.I'])"
   ]
  },
  {
   "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(p['circuit.n1.V'], 9.90804735, 1e-5)\n",
    "assert_near_equal(p['circuit.n2.V'], 0.71278185, 1e-5)\n",
    "assert_near_equal(p['circuit.R1.I'], 0.09908047, 1e-5)\n",
    "assert_near_equal(p['circuit.R2.I'], 0.00091953, 1e-5)\n",
    "assert_near_equal(p['circuit.D1.I'], 0.00091953, 1e-5)\n",
    "\n",
    "# sanity check: should sum to .1 Amps\n",
    "assert_near_equal(p['circuit.R1.I'] + p['circuit.D1.I'], .1, 1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modifying Solver Settings in Your Run Script\n",
    "\n",
    "In the above run script, we set some initial guess values: `prob['n1.V']=10` and `prob['n2.V']=1`.\n",
    "If you try to play around with those initial guesses a bit, you will see that convergence is really sensitive to\n",
    "the initial guess you used for *n2.V*.Below we provide a second run script that uses the same `Circuit` group we defined previously, but which additionally modifies some solver settings and initial guesses. If we set the initial guess for `prob['n2.V']=1e-3`, then the model starts out with a massive residual. It also converges much more slowly, so although we gave it more than twice the number of iterations, it doesn't even get close to a converged answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmdao.test_suite.scripts.circuit_analysis import Circuit\n",
    "\n",
    "p = om.Problem()\n",
    "model = p.model\n",
    "\n",
    "model.add_subsystem('ground', om.IndepVarComp('V', 0., units='V'))\n",
    "model.add_subsystem('source', om.IndepVarComp('I', 0.1, units='A'))\n",
    "model.add_subsystem('circuit', Circuit())\n",
    "\n",
    "model.connect('source.I', 'circuit.I_in')\n",
    "model.connect('ground.V', 'circuit.Vg')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "# you can change the NewtonSolver settings in circuit after setup is called\n",
    "newton = p.model.circuit.nonlinear_solver\n",
    "newton.options['maxiter'] = 50\n",
    "\n",
    "# set some initial guesses\n",
    "p['circuit.n1.V'] = 10.\n",
    "p['circuit.n2.V'] = 1e-3\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "print(p['circuit.n1.V'])\n",
    "print(p['circuit.n2.V'])\n",
    "\n",
    "# sanity check: should sum to .1 Amps\n",
    "print(p['circuit.R1.I'] + p['circuit.D1.I'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(p['circuit.n1.V'], 9.98744708, 1e-5)\n",
    "assert_near_equal(p['circuit.n2.V'], 8.73215484, 1e-5)\n",
    "assert_near_equal(p['circuit.R1.I'] + p['circuit.D1.I'], 0.09987447, 1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{important}\n",
    "You actually *can* get this model to converge. But you have to set the options for `maxiter=400` and `rtol=1e-100`. (The `rtol` value needs to be very low to prevent premature termination.)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tweaking Newton Solver Settings to Get More Robust Convergence\n",
    "\n",
    "The [NewtonSolver](../../features/building_blocks/solvers/newton.ipynb) has a lot of features that allow you to modify its behavior to handle more challenging problems. We're going to look at two of the most important ones here:\n",
    "\n",
    "1. [Line searches](linesearch-section)\n",
    "2. The *solve_subsystems* option\n",
    "\n",
    "If we use both of these in combination, we can dramatically improve the robustness of the solver for this problem.\n",
    "The *linesearch* attribute makes sure that the solver doesn't take too big of a step. The *solve_subsystems* option allows the `Resistor` and `Diode` components (the two `ExplicitComponents`) to help the convergence by updating their own output values given their inputs. When you use [NewtonSolver](../../features/building_blocks/solvers/newton.ipynb) on models with a lot of `ExplicitComponents`, you may find that turning on *solve_subsystems* helps convergence, but you need to be careful about the [execution order](../../features/core_features/working_with_groups/set_order.ipynb) when you try this.\n",
    "\n",
    "```{important}\n",
    "For this case, we used the [ArmijoGoldsteinLS](../../features/building_blocks/solvers/armijo_goldstein.ipynb), which basically limits step sizes so that the residual always goes down. For many problems you might want to use [BoundsEnforceLS](../../features/building_blocks/solvers/bounds_enforce.ipynb) instead, which only activates the line search to enforce upper and lower bounds on the outputs in the model.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = om.Problem()\n",
    "model = p.model\n",
    "\n",
    "model.add_subsystem('ground', om.IndepVarComp('V', 0., units='V'))\n",
    "model.add_subsystem('source', om.IndepVarComp('I', 0.1, units='A'))\n",
    "model.add_subsystem('circuit', Circuit())\n",
    "\n",
    "model.connect('source.I', 'circuit.I_in')\n",
    "model.connect('ground.V', 'circuit.Vg')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "# you can change the NewtonSolver settings in circuit after setup is called\n",
    "newton = p.model.circuit.nonlinear_solver\n",
    "newton.options['iprint'] = 2\n",
    "newton.options['maxiter'] = 10\n",
    "newton.options['solve_subsystems'] = True\n",
    "newton.linesearch = om.ArmijoGoldsteinLS()\n",
    "newton.linesearch.options['maxiter'] = 10\n",
    "newton.linesearch.options['iprint'] = 2\n",
    "\n",
    "# set some initial guesses\n",
    "p['circuit.n1.V'] = 10.\n",
    "p['circuit.n2.V'] = 1e-3\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "print(p['circuit.n1.V'])\n",
    "print(p['circuit.n2.V'])\n",
    "print(p['circuit.R1.I'])\n",
    "print(p['circuit.R2.I'])\n",
    "print(p['circuit.D1.I'])\n",
    "\n",
    "# sanity check: should sum to .1 Amps\n",
    "print(p['circuit.R1.I'] + p['circuit.D1.I'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(p['circuit.n1.V'], 9.90804735, 1e-5)\n",
    "assert_near_equal(p['circuit.n2.V'], 0.71278185, 1e-5)\n",
    "assert_near_equal(p['circuit.R1.I'], 0.09908047, 1e-5)\n",
    "assert_near_equal(p['circuit.R2.I'], 0.00091953, 1e-5)\n",
    "assert_near_equal(p['circuit.D1.I'], 0.00091953, 1e-5)\n",
    "\n",
    "# sanity check: should sum to .1 Amps\n",
    "assert_near_equal(p['circuit.R1.I'] + p['circuit.D1.I'], .1, 1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{important}\n",
    "This tutorial used finite difference to approximate the partial derivatives for all the components. Check out [this example](../../examples/circuit_analysis_examples.ipynb) if you want to see the same problem solved with analytic derivatives.\n",
    "```"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
