{
 "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": [
    "# ImplicitFuncComp\n",
    "\n",
    "`ImplicitFuncComp` is a component that provides a shortcut for building an ImplicitComponent based on a python function. That function takes inputs and states as arguments and returns residual values. The function may take some inputs that are non-differentable and are assumed to be static during the computation of derivatives.  These static values may be of any hashable type.  All other arguments and return values must be either floats or numpy arrays. The mapping between a state argument and its residual output must be specified in the metadata when the output (state) is added by setting 'resid' to the name of the residual.\n",
    "\n",
    "It may seem confusing to use `add_output` to specify state variables since the state variables\n",
    "are actually input arguments to the function, but in OpenMDAO's view of the world, states are outputs so we use `add_output` to specify them.  Also, using the metadata to specify which input arguments are actually states gives more flexibility in terms of how the function arguments are ordered. For example, if it's desirable for a function to be passable to `scipy.optimize.newton`, then the function's arguments can be ordered with the states first, followed by the inputs, in order to match the order expected by `scipy.optimize.newton`.\n",
    "\n",
    "The `add_output` function is part of the [Function Metadata API](../func_api.ipynb).  You use this API to specify various metadata that OpenMDAO needs in order to properly configure a fully functional implicit component. You should read and understand the [Function Metadata API](../func_api.ipynb) before you continue with this section.\n",
    "\n",
    "## ImplicitFuncComp Options\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "def func(a):\n",
    "    y = a * 2.\n",
    "    return y\n",
    "om.show_options_table(om.ImplicitFuncComp(func))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ImplicitFuncComp Constructor\n",
    "\n",
    "The call signature for the `ImplicitFuncComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.implicit_func_comp.ImplicitFuncComp.__init__\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ImplicitFuncComp Example: A simple implicit function component\n",
    "\n",
    "The simplest implicit function component requires the definition of a function that takes\n",
    "inputs and states as arguments and returns residual values.  This function maps to the `apply_nonlinear`\n",
    "method in the OpenMDAO component API. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "import openmdao.func_api as omf\n",
    "\n",
    "def apply_nl(a, b, c, x):  # inputs a, b, c and state x\n",
    "    R_x = a * x ** 2 + b * x + c\n",
    "    return R_x\n",
    "\n",
    "f = (omf.wrap(apply_nl)\n",
    "        .add_output('x', resid='R_x', val=0.0)\n",
    "        .declare_partials(of='*', wrt='*', method='cs')\n",
    "        )\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', om.ImplicitFuncComp(f))\n",
    "\n",
    "p.model.nonlinear_solver = om.NewtonSolver(solve_subsystems=False, iprint=0)\n",
    "\n",
    "# need this since comp is implicit and doesn't have a solve_linear\n",
    "p.model.linear_solver = om.DirectSolver()\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('comp.a', 2.)\n",
    "p.set_val('comp.b', -8.)\n",
    "p.set_val('comp.c', 6.)\n",
    "p.run_model()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_check_partials, assert_check_totals\n",
    "\n",
    "assert_check_partials(p.check_partials(includes=['comp'], out_stream=None), atol=1e-5)\n",
    "assert_check_totals(p.check_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'], out_stream=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ImplicitFuncComp Example: Partial Derivatives\n",
    "\n",
    "All nonzero partial derivatives for an ImplicitFuncComp must be declared when the function is wrapped. Otherwise, OpenMDAO will assume that all partial derivatives for that component are zero.  For example, in the previous model we didn't bother to declare partials because we weren't computing derivatives.  Now we'll compute the total derivatives, so we need to declare the partials.  Also, because our implicit function component does not define its own `linearize` function, we'll specify a `method` of `fd` or `cs` when we declare our partials, so they'll be computed using finite differencing or complex step.  Finally, because our implicit component doesn't define its own `solve_linear` function, we have to specify a linear solver for our component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_nl(a, b, c, x):\n",
    "    R_x = a * x ** 2 + b * x + c\n",
    "    return R_x\n",
    "\n",
    "def solve_nonlinear(a, b, c, x):\n",
    "    x = (-b + (b ** 2 - 4 * a * c) ** 0.5) / (2 * a)\n",
    "    return x\n",
    "\n",
    "f = (omf.wrap(apply_nl)\n",
    "        .add_output('x', resid='R_x', val=0.0)\n",
    "        .declare_partials(of='*', wrt='*', method='cs')\n",
    "        )\n",
    "\n",
    "p = om.Problem()\n",
    "comp = p.model.add_subsystem('comp', om.ImplicitFuncComp(f, solve_nonlinear=solve_nonlinear))\n",
    "\n",
    "# need this since comp is implicit and doesn't have a solve_linear\n",
    "comp.linear_solver = om.DirectSolver()\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('comp.a', 2.)\n",
    "p.set_val('comp.b', -8.)\n",
    "p.set_val('comp.c', 6.)\n",
    "p.run_model()\n",
    "\n",
    "\n",
    "J = p.compute_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'])\n",
    "print(('comp.x', 'comp.a'), J['comp.x', 'comp.a'])\n",
    "print(('comp.x', 'comp.b'), J['comp.x', 'comp.b'])\n",
    "print(('comp.x', 'comp.c'), J['comp.x', 'comp.c'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_check_totals\n",
    "\n",
    "assert_check_partials(p.check_partials(includes=['comp'], out_stream=None), atol=1e-5)\n",
    "assert_check_totals(p.check_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'], out_stream=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ImplicitFuncComp Example: Specifying linearize, solve_linear, and solve_nonlinear functions\n",
    "\n",
    "The following implicit function component specifies `linearize`, `solve_linear`, and `solve_nonlinear`\n",
    "functions, so no external linear or nonlinear solvers are required to compute outputs or derivatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def apply_nl(a, b, c, x):\n",
    "    R_x = a * x ** 2 + b * x + c\n",
    "    return R_x\n",
    "\n",
    "def solve_nonlinear(a, b, c, x):\n",
    "    x = (-b + (b ** 2 - 4 * a * c) ** 0.5) / (2 * a)\n",
    "    return x\n",
    "\n",
    "def linearize(a, b, c, x, partials):\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",
    "    inv_jac = 1.0 / (2 * a * x + b)\n",
    "    return inv_jac\n",
    "\n",
    "def solve_linear(d_x, mode, inv_jac):\n",
    "    if mode == 'fwd':\n",
    "        d_x = inv_jac * d_x\n",
    "        return d_x\n",
    "    elif mode == 'rev':\n",
    "        dR_x = inv_jac * d_x\n",
    "        return dR_x\n",
    "\n",
    "f = (omf.wrap(apply_nl)\n",
    "        .add_output('x', resid='R_x', val=0.0)\n",
    "        .declare_partials(of='*', wrt='*')\n",
    "        )\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', om.ImplicitFuncComp(f,\n",
    "                                                  solve_nonlinear=solve_nonlinear,\n",
    "                                                  solve_linear=solve_linear, \n",
    "                                                  linearize=linearize))\n",
    "p.setup()\n",
    "\n",
    "p.set_val('comp.a', 2.)\n",
    "p.set_val('comp.b', -8.)\n",
    "p.set_val('comp.c', 6.)\n",
    "p.run_model()\n",
    "\n",
    "J = p.compute_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'])\n",
    "print(('comp.x', 'comp.a'), J['comp.x', 'comp.a'])\n",
    "print(('comp.x', 'comp.b'), J['comp.x', 'comp.b'])\n",
    "print(('comp.x', 'comp.c'), J['comp.x', 'comp.c'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_check_partials(p.check_partials(includes=['comp'], out_stream=None), atol=1e-5)\n",
    "assert_check_totals(p.check_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'], out_stream=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using jax to Compute Partial Derivatives\n",
    "\n",
    "If the function used to instantiate the ExplicitFuncComp declares partials or coloring that use `method='jax'`, or if the component's `use_jax` option is set, then the jax AD package will be used to compute all of the component's derivatives.  Currently it's not possible to mix jax with finite difference methods ('cs' and 'fd') in the same component.\n",
    "\n",
    "Note that `jax` is not an OpenMDAO dependency, so you'll have to install it manually.\n",
    "\n",
    "```python\n",
    "pip install jax\n",
    "pip install jaxlib\n",
    "```\n",
    "\n",
    "should work in most cases.\n",
    "\n",
    "To activate jax's just-in-time compilation capability, set the `use_jit` option on the component. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', om.ImplicitFuncComp(f,\n",
    "                                                  solve_nonlinear=solve_nonlinear,\n",
    "                                                  solve_linear=solve_linear, \n",
    "                                                  linearize=linearize,\n",
    "                                                  use_jit=True))\n",
    "p.setup()\n",
    "\n",
    "p.set_val('comp.a', 2.)\n",
    "p.set_val('comp.b', -8.)\n",
    "p.set_val('comp.c', 6.)\n",
    "p.run_model()\n",
    "\n",
    "J = p.compute_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'])\n",
    "print(('comp.x', 'comp.a'), J['comp.x', 'comp.a'])\n",
    "print(('comp.x', 'comp.b'), J['comp.x', 'comp.b'])\n",
    "print(('comp.x', 'comp.c'), J['comp.x', 'comp.c'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_check_partials(p.check_partials(includes=['comp'], out_stream=None), atol=1e-5)\n",
    "assert_check_totals(p.check_totals(of=['comp.x'], wrt=['comp.a', 'comp.b', 'comp.c'], out_stream=None))"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "interpreter": {
   "hash": "245bb6672fbc289f90037d9f00b5ee20de7d921e65d14dbc4c07ab973781223d"
  },
  "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.8.12"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
