{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true,
    "tags": [
     "active-ipynb",
     "remove-input",
     "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": [
    "# Converging an Implicit Model: Nonlinear circuit analysis\n",
    "",
    "\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",
    "![diagram of a simple circuit with two resistors and one diode](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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "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",
    "        # partial derivs are constant, so we can assign their values in setup\n",
    "        R = self.options['R']\n",
    "        self.declare_partials('I', 'V_in', val=1 / R)\n",
    "        self.declare_partials('I', 'V_out', val=-1 / R)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        deltaV = inputs['V_in'] - inputs['V_out']\n",
    "        outputs['I'] = deltaV / self.options['R']\n",
    "\n",
    "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",
    "        # non-linear component, so we'll declare the partials here but compute them in compute_partials\n",
    "        self.declare_partials('I', 'V_in')\n",
    "        self.declare_partials('I', 'V_out')\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)\n",
    "\n",
    "    def compute_partials(self, inputs, J):\n",
    "        deltaV = inputs['V_in'] - inputs['V_out']\n",
    "        Is = self.options['Is']\n",
    "        Vt = self.options['Vt']\n",
    "        I = Is * np.exp(deltaV / Vt)\n",
    "\n",
    "        J['I', 'V_in'] = I/Vt\n",
    "        J['I', 'V_out'] = -I/Vt\n",
    "\n",
    "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",
    "            self.declare_partials('V', i_name, val=1)\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",
    "            self.declare_partials('V', i_name, val=-1)\n",
    "\n",
    "            # note: we don't declare any partials wrt `V` here,\n",
    "            #      because the residual doesn't directly depend on it\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)]\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()\n",
    "        self.linear_solver = om.DirectSolver()\n",
    "\n",
    "        self.nonlinear_solver.options['iprint'] = 2\n",
    "        self.nonlinear_solver.options['maxiter'] = 10\n",
    "        self.nonlinear_solver.options['solve_subsystems'] = True\n",
    "        self.nonlinear_solver.linesearch = om.ArmijoGoldsteinLS()\n",
    "        self.nonlinear_solver.linesearch.options['maxiter'] = 10\n",
    "        self.nonlinear_solver.linesearch.options['iprint'] = 2\n",
    "\n",
    "\n",
    "p = om.Problem()\n",
    "model = p.model\n",
    "\n",
    "model.add_subsystem('circuit', Circuit())\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('circuit.I_in', 0.1)\n",
    "p.set_val('circuit.Vg', 0.)\n",
    "\n",
    "# set some initial guesses\n",
    "p.set_val('circuit.n1.V', 10.)\n",
    "p.set_val('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": [
    "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)"
   ]
  }
 ],
 "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
}
