{
 "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  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function Metadata API\n",
    "\n",
    "Using [ExplicitFuncComp](components/explicit_func_comp.ipynb), you can turn a python function \n",
    "into a fully functioning OpenMDAO component.  However, in order to do that it's sometimes necessary\n",
    "to attach additional metadata to the function so that OpenMDAO can be informed of things like\n",
    "variable units and shapes, and partial derivative information.  Metadata can be\n",
    "attached to a function using the function metadata API.  It works by wrapping the function in a \n",
    "callable object that can store the metadata appropriately.\n",
    "\n",
    "## Function wrapping\n",
    "\n",
    "We wrap a function using the `omf.wrap` function, for example:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.func_api as omf\n",
    "\n",
    "def func(a):\n",
    "    x = a * 2.\n",
    "    return x\n",
    "\n",
    "f = omf.wrap(func) "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`omf.wrap` returns an instance of the `OMWrappedFunc` class that can store various metadata needed by\n",
    "OpenMDAO.  All of the metadata setting functions called on that instance return the instance itself\n",
    "so they can be stacked together.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = omf.wrap(func).add_input('a', shape=5).add_output('x', shape=5, units='m')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, if you need to make many calls to set metadata on the wrapped function, you can stack the calls\n",
    "vertically, but this will only work if you wrap the entire righ-hand-side expression in parentheses so\n",
    "that python will treat it all as a single expression.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = (omf.wrap(func)\n",
    "        .defaults(units='m')\n",
    "        .add_input('a', shape=5)\n",
    "        .add_output('x', shape=5))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If stacking isn't desired, the methods can just be called in the usual way, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = omf.wrap(func)\n",
    "f.defaults(units='m')\n",
    "f.add_input('a', shape=5)\n",
    "f.add_output('x', shape=5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variable metadata\n",
    "\n",
    "### Setting the metadata for a single variable\n",
    "\n",
    "OpenMDAO needs to know a variable's shape, initial value, and optionally other things like units.  \n",
    "This information can be specified using the `add_input` and `add_output` methods.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(x):\n",
    "    y = x.dot(np.random.random(2))\n",
    "    return y\n",
    "\n",
    "f = (omf.wrap(func)\n",
    "        .add_input('x', shape=(2,2))\n",
    "        .add_output('y', shape=2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting metadata for option variables\n",
    "\n",
    "A function may have additional non-float or non-float ndarray arguments that, at least in the\n",
    "OpenMDAO context, will be treated as component options that don't change during a given model\n",
    "execution.  These can be specified using the `declare_option` method.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(x, opt):\n",
    "    if opt == 1:\n",
    "        y = x.dot(np.random.random(2))\n",
    "    elif opt == 2:\n",
    "        y = x[:, 1] * 2.\n",
    "    elif opt == 3:\n",
    "        y = x[1, :] * 3.\n",
    "    return y\n",
    "\n",
    "f = (omf.wrap(func)\n",
    "        .add_input('x', shape=(2,2))\n",
    "        .declare_option('opt', values=[1, 2, 3])\n",
    "        .add_output('y', shape=2))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The arguments that are passable to `declare_option` are the same as those that are allowed when\n",
    "declaring option variables in an OpenMDAO component using the [OptionsDictionary](../../_srcdocs/packages/utils/options_dictionary) `declare` method.\n",
    "\n",
    "### Setting metadata for multiple variables\n",
    "\n",
    "Using the `add_inputs` and `add_outputs` methods you can specify metadata for multiple variables\n",
    "in the same call.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a, b):\n",
    "    return a.dot(b), a[:,0] * b * b\n",
    "\n",
    "f = (omf.wrap(func)\n",
    "        .add_inputs(a={'shape': (2,2), 'units': 'm'}, b={'shape': 2, 'units': 'm'})\n",
    "        .add_outputs(x={'shape': 2, 'units': 'm**2'}, y={'shape': 2, 'units': 'm**3'}))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting the metadata\n",
    "\n",
    "Variable metadata is retrieved from the wrapped function by calling the \n",
    "`get_input_meta` and `get_output_meta` methods. Each function returns an iterator over (name, \n",
    "metadata_dict) tuples, one for each input or output variable respectively.  For example, the \n",
    "following code snippet will print the name and shape of each output variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, meta in f.get_output_meta():\n",
    "    print(name, meta['shape'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting function default metadata\n",
    "\n",
    "Some metadata will be the same for all, or at least most of the variables within a given function,\n",
    "so we want to be able to specify those defaults easily without too much boilerplate.  That's the\n",
    "purpose of the `defaults` method.  For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a, b, c):\n",
    "    d = a * b * c\n",
    "    return d\n",
    "\n",
    "f = omf.wrap(func).defaults(shape=4, units='m')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Any metadata that is specific to a particular variable will override any defaults specified in\n",
    "`defaults`. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def func(a, b, c=np.ones(3)):  # shape of c is 3 which overrides the `defaults` shape of 4\n",
    "    d = a * b\n",
    "    e = c * 1.5\n",
    "    return d, e\n",
    "\n",
    "f = omf.wrap(func).defaults(shape=4, units='m')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Assumed default values\n",
    "\n",
    "In order to stay consistent with OpenMDAO's default value policy, we assume the same default\n",
    "behavior for functions, so if no shape or default value is supplied for a function variable, we\n",
    "assume that is has the value 1.0.  If the `shape` is provided and either the default value is\n",
    "not provided or is provided as a scalar value, then the assumed default value will be\n",
    "`np.ones(shape) * scalar_value`, where `scalar_value` is 1.0 if not specified.\n",
    "If `shape` is provided along with a non-scalar default value that has a different shape, then\n",
    "an exception will be raised.\n",
    "\n",
    "\n",
    "## Variable names\n",
    "\n",
    "### Setting variable names\n",
    "\n",
    "We don't need to set input names because the function can always be inspected for those, but\n",
    "we do need to associate output names with function return values. Those return values, if they are \n",
    "simple variables, for example, `return x, y`, will give us the output variable names we need.  \n",
    "But in those cases where the function returns expressions rather than simple variables, we need \n",
    "another way to specify what the names of those output variables should be.  The `output_names` \n",
    "method provides a concise way to do this, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a, b, c):\n",
    "    return a * b * c, a * b -c  # two return values that don't have simple names\n",
    "\n",
    "f = omf.wrap(func).output_names('d', 'e')  # name of return values are 'd' and 'e'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we have metadata we need to supply for the outputs, we could instead just use the\n",
    "`add_outputs` method mentioned earlier, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a, b, c):\n",
    "    return a * b * c, a * b -c  # two return values that don't have simple names\n",
    "\n",
    "# names of return values are 'd' and 'e'. \n",
    "f = omf.wrap(func).add_outputs(d={'units': 'm'}, e={'units': 'ft'})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As mentioned above, if the function's return values are simple variable names, we don't need to\n",
    "specify the output names because we can determine them by inspecting the function, e.g., "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(a, b, c):\n",
    "    d = a * b * c\n",
    "    e = a * b -c\n",
    "    return d, e  # we know from inspection that the output names are 'd' and 'e'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that in the function above, we didn't have to wrap it at all.  This is possible because we can \n",
    "inspect the source code to determine the output names and we assume the default value of all inputs\n",
    "and outputs is 1.0.  If any inputs or outputs of a function have any non-default metadata, e.g.,\n",
    "val, units, shape, etc., then that function would have to be wrapped and those metadata values\n",
    "would have to be specified. Also, if we plan to compute derivatives for the function then we would\n",
    "need to specify which partials are nonzero using the `declare_partials` method.\n",
    "\n",
    "If one or more output names are not specified and cannot be determined by inspection, then they \n",
    "must be specified using `add_output` calls. The number of `add_output` calls corresponding to unnamed\n",
    "return values must match the total number of unnamed return values, and they will be matched to those \n",
    "return values in the order that they are called.  Any call to `add_output` with an output name that \n",
    "corresponds to one already specified can occur in any order.  In the example below, there\n",
    "are two return values and neither output name is specified, so two calls to `add_output` are needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def func(x):\n",
    "    return x.dot(np.random.random(2)), x*1.5  # 2 return values and we can't infer the names\n",
    "f = (omf.wrap(func)\n",
    "        .add_input('x', shape=(2,2))\n",
    "        .add_output('y', shape=2)       # 'y' is the name of the first return value\n",
    "        .add_output('z', shape=(2,2)))  # 'z' is the name of the second return value"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the example above, the output names would be assumed to be `['y', 'z']`.\n",
    "\n",
    "### Getting variable names\n",
    "\n",
    "Lists of input names and output names can be retrieved by calling `get_input_names` and \n",
    "`get_output_names` respectively, e.g., "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('input names =', list(f.get_input_names()))\n",
    "print('output names = ', list(f.get_output_names()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partial derivatives\n",
    "\n",
    "### Setting partial derivative information\n",
    "\n",
    "Metadata that will help OpenMDAO to compute partial derivatives\n",
    "for the function can be defined using the `declare_partials` and `declare_coloring` methods.\n",
    "For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def func(x, y, z=3): \n",
    "    foo = np.log(z)/(3*x+2*y)\n",
    "    bar = 2*x+y\n",
    "    return foo, bar\n",
    "\n",
    "f = (omf.wrap(func)\n",
    "        .declare_partials(of='*', wrt='*', method='cs')\n",
    "        .declare_coloring(wrt='*', method='cs')\n",
    "        .defaults(shape=4))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The arguments for the `declare_partials` and `declare_coloring` methods match those\n",
    "of the same methods on [Component](../../_srcdocs/packages/core/component).  Multiple calls\n",
    "can be made to `declare_partials` to set up different partials, but `declare_coloring` may only\n",
    "be called once.\n",
    "\n",
    "Note that all nonzero partial derivatives *must* be declared or OpenMDAO will assume they are zero.\n",
    "\n",
    "### Getting partial derivative information\n",
    "\n",
    "The arguments passed to the `declare_partials` and `declare_coloring` methods can be retrieved \n",
    "using the `get_declare_partials` and `get_declare_coloring` methods respectively.  Each of these\n",
    "returns a list where each entry is the keyword args dict from each call, in the order that they\n",
    "where called."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f.get_declare_partials())  # returns a list of args dicts for multiple calls\n",
    "print(f.get_declare_coloring())   # returns args dict for a single call to declare_coloring"
   ]
  }
 ],
 "metadata": {
  "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.10"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
