{
 "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": [
    "# BalanceComp\n",
    "\n",
    "`BalanceComp` is a specialized implementation of [ImplicitComponent](../../core_features/working_with_components/implicit_component.ipynb) that is intended to provide a simple way to implement most implicit equations without the need to define your own residuals.\n",
    "\n",
    "## BalanceComp Options\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(\"openmdao.components.balance_comp.BalanceComp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BalanceComp Constructor\n",
    "\n",
    "The call signature for the `BalanceComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.balance_comp.BalanceComp.__init__()\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the BalanceComp\n",
    "\n",
    "`BalanceComp` allows you to add one or more state variables and its associated\n",
    "implicit equations.  For each ``balance`` added to the component it\n",
    "solves the following equation:\n",
    "\n",
    "$$\n",
    "  \\begin{align}\n",
    "  \\mathcal{R}_{name} =\n",
    "  \\frac{f_{mult}(x,...) \\times f_{lhs}(x,...) - f_{rhs}(x,...)}{f_{norm}(f_{rhs}(x,...))}\n",
    "  \\end{align}\n",
    "$$\n",
    "\n",
    "The optional normalization function $f_{norm}(f_{rhs})$ is computed as:\n",
    "\n",
    "$$\n",
    "  \\begin{align}\n",
    "  f_{norm}(f_{rhs}(x,...)) =\n",
    "  \\begin{cases}\n",
    "   \\left| f_{rhs} \\right|, & \\text{if normalize and } \\left| f_{rhs} \\right| \\geq 2 \\\\\n",
    "   0.25 f_{rhs}^2 + 1,     & \\text{if normalize and } \\left| f_{rhs} \\right| < 2 \\\\\n",
    "   1,                      & \\text{if not normalize}\n",
    "  \\end{cases}\n",
    "  \\end{align}\n",
    "$$\n",
    "\n",
    "The following inputs and outputs are associated with each implicit state.\n",
    "\n",
    "```{eval-rst}\n",
    "=========== ======= ====================================================\n",
    "Name        I/O     Description\n",
    "=========== ======= ====================================================\n",
    "{name}      output  implicit state variable\n",
    "lhs:{name}  input   left-hand side of equation to be balanced\n",
    "rhs:{name}  input   right-hand side of equation to be balanced\n",
    "mult:{name} input   left-hand side multiplier of equation to be balanced\n",
    "=========== ======= ====================================================\n",
    "```\n",
    "\n",
    "The default value for the `rhs:{name}` input can be set to via the\n",
    "`rhs_val` argument (see arguments below). If the rhs value is fixed (e.g. 0),\n",
    "then the input can be left unconnected. The `lhs:{name}` input must always have\n",
    "something connected to it and should be dependent upon the value of the implicit state variable.\n",
    "\n",
    "The multiplier is optional and will default to 1.0 if not connected.\n",
    "\n",
    "`BalanceComp` supports vectorized implicit states. Simply provide a default\n",
    "value or shape when adding the balance that reflects the correct shape.\n",
    "\n",
    "You can provide the arguments to create a balance when instantiating a `BalanceComp`\n",
    "or you can use the ``add_balance`` method to create one or more state variables after\n",
    "instantiation.  The constructor accepts all the same arguments as the ``add_balance``\n",
    "method:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.balance_comp.BalanceComp.add_balance\n",
    "       :noindex:\n",
    "```\n",
    "\n",
    "Note that the `kwargs` arguments can include any of the keyword arguments normally available\n",
    "when creating an output variable with the\n",
    "[add_output](../../../_srcdocs/packages/core/component.html#openmdao.core.component.Component.add_output) method of a `Component`.\n",
    "\n",
    "## Example:  Scalar Root Finding\n",
    "\n",
    "The following example uses a BalanceComp to implicitly solve the\n",
    "equation:\n",
    "\n",
    "$$\n",
    "    2 \\cdot x^2 = 4\n",
    "$$\n",
    "\n",
    "Here, our LHS is connected to a computed value for $x^2$, the multiplier is 2, and the RHS\n",
    "is 4.  The expected solution is $x=\\sqrt{2}$.  We initialize $x$ with a value of 1 so that\n",
    "it finds the positive root."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "bal = om.BalanceComp()\n",
    "bal.add_balance('x', use_mult=True)\n",
    "\n",
    "exec_comp = om.ExecComp('y=x**2', x={'val': 1}, y={'val': 1})\n",
    "\n",
    "prob.model.add_subsystem(name='exec', subsys=exec_comp)\n",
    "prob.model.add_subsystem(name='balance', subsys=bal)\n",
    "\n",
    "prob.model.connect('balance.x', 'exec.x')\n",
    "prob.model.connect('exec.y', 'balance.lhs:x')\n",
    "\n",
    "prob.model.linear_solver = om.DirectSolver(assemble_jac=True)\n",
    "prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=0)\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('balance.rhs:x', 4)\n",
    "prob.set_val('balance.mult:x', 2.)\n",
    "\n",
    "# A reasonable initial guess to find the positive root.\n",
    "prob['balance.x'] = 1.0\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('balance.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.testing import assert_almost_equal\n",
    "\n",
    "assert_almost_equal(prob.get_val('balance.x'), np.sqrt(2), decimal=7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, we could simplify the code by using the `mult_val` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = om.Problem()\n",
    "\n",
    "bal = om.BalanceComp()\n",
    "bal.add_balance('x', use_mult=True, mult_val=2.0)\n",
    "\n",
    "exec_comp = om.ExecComp('y=x**2', x={'val': 1}, y={'val': 1})\n",
    "\n",
    "prob.model.add_subsystem(name='exec', subsys=exec_comp)\n",
    "prob.model.add_subsystem(name='balance', subsys=bal)\n",
    "\n",
    "prob.model.connect('balance.x', 'exec.x')\n",
    "prob.model.connect('exec.y', 'balance.lhs:x')\n",
    "\n",
    "prob.model.linear_solver = om.DirectSolver(assemble_jac=True)\n",
    "prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=0)\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('balance.rhs:x', 4)\n",
    "\n",
    "# A reasonable initial guess to find the positive root.\n",
    "prob.set_val('balance.x', 1.0)\n",
    "\n",
    "prob.run_model()\n",
    "print(prob.get_val('balance.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_almost_equal(prob.get_val('balance.x'), np.sqrt(2), decimal=7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example:  Vectorized Root Finding\n",
    "\n",
    "The following example uses a BalanceComp to implicitly solve the equation:\n",
    "\n",
    "$$\n",
    "    b \\cdot x + c  = 0\n",
    "$$\n",
    "\n",
    "for various values of $b$, and $c$.  Here, our LHS is connected to a computed value of\n",
    "the linear equation.  The multiplier is one and the RHS is zero (the defaults), and thus\n",
    "they need not be connected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 100\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "exec_comp = om.ExecComp('y=b*x+c',\n",
    "                        b={'val': np.random.uniform(0.01, 100, size=n)},\n",
    "                        c={'val': np.random.rand(n)},\n",
    "                        x={'val': np.zeros(n)},\n",
    "                        y={'val': np.ones(n)})\n",
    "\n",
    "prob.model.add_subsystem(name='exec', subsys=exec_comp)\n",
    "prob.model.add_subsystem(name='balance', subsys=om.BalanceComp('x', val=np.ones(n)))\n",
    "\n",
    "prob.model.connect('balance.x', 'exec.x')\n",
    "prob.model.connect('exec.y', 'balance.lhs:x')\n",
    "\n",
    "prob.model.linear_solver = om.DirectSolver(assemble_jac=True)\n",
    "prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=0)\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('balance.x', np.random.rand(n))\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "b = prob.get_val('exec.b')\n",
    "c = prob.get_val('exec.c')\n",
    "\n",
    "print(prob.get_val('balance.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(-c/b)  # expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_almost_equal(prob.get_val('balance.x'), -c/b, decimal=6)\n",
    "assert_almost_equal(-c/b, prob.get_val('balance.x'), decimal=6)  # expected"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example:  Providing an Initial Guess for a State Variable\n",
    "\n",
    "`BalanceComp` has a `guess_func` option that can be used to supply an initial guess\n",
    "value for the state variables.  This option provides the same functionality as the\n",
    "[guess_nonlinear](../../../_srcdocs/packages/core/implicitcomponent.html#openmdao.core.implicitcomponent.ImplicitComponent.guess_nonlinear) \n",
    "method of `ImplicitComponent`.\n",
    "\n",
    "The Kepler example script shows how `guess_func` can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = om.Problem()\n",
    "\n",
    "bal = om.BalanceComp()\n",
    "\n",
    "bal.add_balance(name='E', val=0.0, units='rad', eq_units='rad', rhs_name='M')\n",
    "\n",
    "# Use M (mean anomaly) as the initial guess for E (eccentric anomaly)\n",
    "def guess_function(inputs, outputs, residuals):\n",
    "    if np.abs(residuals['E']) > 1.0E-2:\n",
    "        outputs['E'] = inputs['M']\n",
    "\n",
    "bal.options['guess_func'] = guess_function\n",
    "\n",
    "# ExecComp used to compute the LHS of Kepler's equation.\n",
    "lhs_comp = om.ExecComp('lhs=E - ecc * sin(E)',\n",
    "                       lhs={'val': 0.0, 'units': 'rad'},\n",
    "                       E={'val': 0.0, 'units': 'rad'},\n",
    "                       ecc={'val': 0.0})\n",
    "\n",
    "prob.model.add_subsystem(name='balance', subsys=bal,\n",
    "                         promotes_inputs=['M'],\n",
    "                         promotes_outputs=['E'])\n",
    "\n",
    "prob.model.set_input_defaults('M', 85.0, units='deg')\n",
    "\n",
    "prob.model.add_subsystem(name='lhs_comp', subsys=lhs_comp,\n",
    "                         promotes_inputs=['E', 'ecc'])\n",
    "\n",
    "# Explicit connections\n",
    "prob.model.connect('lhs_comp.lhs', 'balance.lhs:E')\n",
    "\n",
    "# Set up solvers\n",
    "prob.model.linear_solver = om.DirectSolver()\n",
    "prob.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, maxiter=100, iprint=2)\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('ecc', 0.6)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(np.degrees(prob.get_val('E')))"
   ]
  },
  {
   "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_almost_equal(np.degrees(prob.get_val('E')), 115.9, decimal=1)"
   ]
  }
 ],
 "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
}
