{
 "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": [
    "# ExplicitFuncComp\n",
    "\n",
    "\n",
    "`ExplicitFuncComp` is a component that provides a shortcut for building an ExplicitComponent based\n",
    "on a python function. The function must have one or more differentiable arguments, which can be either\n",
    "floats or numpy arrays, and zero or more arguments that are treated as non-differentiable and are \n",
    "assumed to remain static during the computation of derivatives. If jax is used, then any non-differentiable arguments must be hashable.  The function must also return either a single float or numpy array or a tuple of such values, where each entry in the tuple represents a different output of the OpenMDAO component. The ExplicitFuncComp automatically takes care of all of the component API methods, so you just need to instantiate it with a function.  In most cases that function will need additional metadata that you can add using the [Function Metadata API](../func_api.ipynb).  You should read and understand the [Function Metadata API](../func_api.ipynb) before you continue with this section.\n",
    "\n",
    "## ExplicitFuncComp 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.ExplicitFuncComp(func))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExplicitFuncComp Constructor\n",
    "\n",
    "The call signature for the `ExplicitFuncComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.explicit_func_comp.ExplicitFuncComp.__init__\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExplicitFuncComp Example: Simple\n",
    "\n",
    "For example, here is a simple component that takes the input and adds one to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "def func(x=2.0):\n",
    "    y = x + 1.\n",
    "    return y\n",
    "\n",
    "model.add_subsystem('comp', om.ExplicitFuncComp(func))\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.y'))"
   ]
  },
  {
   "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(prob.get_val('comp.y'), 3.0, 0.00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExplicitFuncComp Example: Arrays\n",
    "\n",
    "You can declare an ExplicitFuncComp with arrays for inputs and outputs.  In the case of inputs,\n",
    "you must either provide default array values or you must set their 'shape' metadata correctly using\n",
    "the [Function Metadata API](../func_api.ipynb).  For outputs you must provide 'shape' \n",
    "metadata as well. In the example below, the input shape of `x` is set via the function default value \n",
    "and the output `y` has its shape set via the `add_output` method of the [Function Metadata API](../func_api.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.func_api as omf\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "def func(x=np.array([1., 2., 3.])):\n",
    "    y = x[:2]\n",
    "    return y\n",
    "\n",
    "f = omf.wrap(func).add_output('y', shape=2)\n",
    "\n",
    "model.add_subsystem('comp', om.ExplicitFuncComp(f))\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('comp.y'), np.array([1.0, 2.0]), 0.00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExplicitFuncComp Example: Variable Properties\n",
    "\n",
    "You can also declare properties like 'units' on the inputs and outputs. In this\n",
    "example we declare all our inputs to be inches to trigger conversion from a variable expressed in feet\n",
    "in one connection source."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "def func(x, y):\n",
    "    z = x + y\n",
    "    return z\n",
    "\n",
    "f = omf.wrap(func).defaults(units='inch')\n",
    "\n",
    "model.add_subsystem('comp', om.ExplicitFuncComp(f))\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('comp.x', 12.0, units='inch')\n",
    "prob.set_val('comp.y', 1.0, units='ft')\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.z'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('comp.z'), 24.0, 0.00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExplicitFuncComp Example: Partial Derivatives\n",
    "\n",
    "All nonzero partial derivatives for an ExplicitFuncComp must be declared when the function is wrapped.\n",
    "Otherwise, OpenMDAO will assume that all partial derivatives for that component are zero.  For\n",
    "example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a=2.0, b=3.0):\n",
    "    x = 2. * a\n",
    "    y = b - 1.0 / 3.0\n",
    "    return x, y\n",
    "\n",
    "f = (omf.wrap(func)\n",
    "        .defaults(method='cs')\n",
    "        .declare_partials(of='x', wrt='a')\n",
    "        .declare_partials(of='y', wrt='b'))\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', om.ExplicitFuncComp(f))\n",
    "p.setup()\n",
    "p.run_model()\n",
    "J = p.compute_totals(of=['comp.x', 'comp.y'], wrt=['comp.a', 'comp.b'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(J['comp.x', 'comp.a'][0][0])\n",
    "print(J['comp.y', 'comp.b'][0][0])"
   ]
  },
  {
   "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(J['comp.x', 'comp.a'][0][0], 2.0, .00001)\n",
    "assert_near_equal(J['comp.y', 'comp.b'][0][0], 1.0, .00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "## ExplicitFuncComp Example: Sparse Partials\n",
    "\n",
    "If you know that some of the partials are sparse, then you should declare them as sparse\n",
    "in order to get the best possible performance when computing derivatives for your component. Here's an example of a function with sparse, in this case diagonal, partials:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(x=np.ones(5), y=np.ones(5)):\n",
    "    z = x * y\n",
    "    return z\n",
    "\n",
    "f = (omf.wrap(func)\n",
    "        .add_output('z', shape=5)\n",
    "        .declare_partials(of='z', wrt=['x', 'y'], method='cs', rows=np.arange(5), cols=np.arange(5)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ExplicitFuncComp Example: Default metadata values\n",
    "\n",
    "Metadata that can apply to all the variables in the component are shape and units.\n",
    "These can be set via the `defaults` method of the [Function Metadata API](../func_api.ipynb). \n",
    "In the following example the variables all share the same shape and units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def func(x=np.ones(5)):\n",
    "    y = 2. * x[2]\n",
    "    return y\n",
    "\n",
    "f = omf.wrap(func).defaults(shape=1, units='m')\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', om.ExplicitFuncComp(f))\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('comp.x', [100., 200., 300., 400., 500.], units='cm')\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('comp.y'), [6.], 0.00001)"
   ]
  },
  {
   "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": [
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', om.ExplicitFuncComp(f, use_jit=True))\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('comp.x', [100., 200., 300., 400., 500.], units='cm')\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('comp.y'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('comp.y'), [6.], 0.00001)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "interpreter": {
   "hash": "16ed3b9d19aea4081222d83efe886ecdef9297b336d0dad7bfd68d57c9797627"
  },
  "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
}
