{
 "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": [
    "In the [previous tutorial](../models_implicit_components/models_with_solvers_implicit.ipynb) , we built up a model of an electrical circuit using a combination of [ImplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb) and [ExplicitComponent](../../features/core_features/working_with_components/explicit_component.ipynb) instances.In that tutorial, all of the implicit relationships in the model came directly from the physics of the model itself.\n",
    "\n",
    "However, you often need to add implicit relationships to models by driving the values of two separate variables to be equal to each other. In this tutorial, we'll show you how to do that using the [BalanceComp](../../features/building_blocks/components/balance_comp.ipynb).\n",
    "\n",
    "# Using BalanceComp to Create Implicit Relationships in Groups\n",
    "\n",
    "The electrical circuit model from the [previous tutorial](../models_implicit_components/models_with_solvers_implicit.ipynb) represents a very basic\n",
    "circuit with a current source of .1 Amps. Here is a reminder of what that circuit looked like:\n",
    "\n",
    "![circuit_diagram](images/circuit_diagram.png)\n",
    "\n",
    "When we solved that circuit, the resulting voltage at *node 1* was 9.9 Volts. Let's say you wanted to power this circuit with a 1.5-Volt battery, instead of using a current source. We can make a small modification to our original model to capture this new setup.\n",
    "\n",
    "Given any value for `source.I`, this model outputs the value for `n1.V` that balances the model. The voltage at the ground is also known via `ground.V`. So the voltage across the current source is\n",
    "\n",
    "$$V_{source} = V1 - V0$$\n",
    "\n",
    "To represent a voltage source with a specific voltage, we can add an additional state variable and residual equation to our model:\n",
    "\n",
    "$${R}_{batt} = V1 - V0 - V_{source}^{*}$$\n",
    "\n",
    "where $V_{source}^{*}$, the desired source voltage, is given by the user as parameter to the model.\n",
    "\n",
    "We could write a new component, inheriting from [ImplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb), to include this new relationship into the model, but OpenMDAO provides [BalanceComp](../../features/building_blocks/components/balance_comp.ipynb), a general utility component that is designed specifically for this type of situation.\n",
    "\n",
    "What we're going to do is add a [BalanceComp](../../features/building_blocks/components/balance_comp.ipynb) to the top level of the model.\n",
    "The `BalanceComp` will define a residual that will drive the source current to force the delta-V across the battery to be what we want. We'll also add an [ExecComp](../../features/building_blocks/components/exec_comp.ipynb) to compute that delta-V from the ground voltage and the voltage at node 1 and then connect everything up. Lastly, since we added an [ImplicitComponent](../../features/core_features/working_with_components/implicit_component.ipynb) at the top level of the model, we'll also move the [NewtonSolver](../../features/building_blocks/solvers/newton.ipynb) up to the top level of the model, too.\n",
    "\n",
    "```{important}\n",
    "BalanceComp can handle more than just $lhs-rhs=0$. It has a number of inputs that let you tweak that behavior. It can support multiple residuals and array variables as well. Check out the [documentation](../../features/building_blocks/components/balance_comp.ipynb) on it for details.\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(\"circuit_src\", get_code(\"openmdao.test_suite.scripts.circuit_analysis.Circuit\"), display=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{Admonition} `Circuit` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`circuit_src`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "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",
    "\n",
    "# replacing the fixed current source with a BalanceComp to represent a fixed Voltage source\n",
    "# model.add_subsystem('source', om.IndepVarComp('I', 0.1, units='A'))\n",
    "model.add_subsystem('batt', om.IndepVarComp('V', 1.5, units='V'))\n",
    "bal = model.add_subsystem('batt_balance', om.BalanceComp())\n",
    "bal.add_balance('I', units='A', eq_units='V')\n",
    "\n",
    "model.add_subsystem('circuit', Circuit())\n",
    "model.add_subsystem('batt_deltaV', om.ExecComp('dV = V1 - V2', V1={'units':'V'},\n",
    "                                               V2={'units':'V'}, dV={'units':'V'}))\n",
    "\n",
    "# current into the circuit is now the output state from the batt_balance comp\n",
    "model.connect('batt_balance.I', 'circuit.I_in')\n",
    "model.connect('ground.V', ['circuit.Vg','batt_deltaV.V2'])\n",
    "model.connect('circuit.n1.V', 'batt_deltaV.V1')\n",
    "\n",
    "# set the lhs and rhs for the battery residual\n",
    "model.connect('batt.V', 'batt_balance.rhs:I')\n",
    "model.connect('batt_deltaV.dV', 'batt_balance.lhs:I')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "###################\n",
    "# Solver Setup\n",
    "###################\n",
    "\n",
    "# change the circuit solver to RunOnce because we're\n",
    "# going to converge at the top level of the model with newton instead\n",
    "p.model.circuit.nonlinear_solver = om.NonlinearRunOnce()\n",
    "p.model.circuit.linear_solver = om.LinearRunOnce()\n",
    "\n",
    "# Put Newton at the top so it can also converge the new BalanceComp residual\n",
    "newton = p.model.nonlinear_solver = om.NewtonSolver()\n",
    "p.model.linear_solver = om.DirectSolver()\n",
    "newton.options['iprint'] = 2\n",
    "newton.options['maxiter'] = 20\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 initial guesses from the current source problem\n",
    "p['circuit.n1.V'] = 9.8\n",
    "p['circuit.n2.V'] = .7\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'])"
   ]
  },
  {
   "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'], 1.5, 1e-5)\n",
    "assert_near_equal(p['circuit.n2.V'], 0.65113362, 1e-5)\n",
    "assert_near_equal(p['circuit.R1.I'], 0.015, 1e-5)\n",
    "assert_near_equal(p['circuit.R2.I'], 8.48866375e-05, 1e-5)\n",
    "assert_near_equal(p['circuit.D1.I'], 8.48866375e-05, 1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Understanding How Everything Is Connected in This Model\n",
    "\n",
    "There are a number of connections in this model, and several different residuals being converged. Trying to keep track of all the connections in your head can be a bit challenging, but OpenMDAO offers some visualization tools to help see what's going on.\n",
    "\n",
    "The `openmdao n2` command can be used to view an $N^2$ diagram of the model.  It can be used\n",
    "as follows:\n",
    "```\n",
    "openmdao n2 <your_python_script>\n",
    "```\n",
    "\n",
    "You can do the same thing programmatically by calling the `n2` function in your python script (after setup):\n",
    "```\n",
    "p.setup()\n",
    "\n",
    "om.n2(p)\n",
    "```\n",
    "\n",
    "Here is what the resulting visualization would look like for the above model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "om.n2(p, outfile=\"circuit_analysis.html\")"
   ]
  }
 ],
 "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
}
