{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "upset-transaction",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "allied-means",
   "metadata": {},
   "source": [
    "# Determining Variable Units at Runtime\n",
    "\n",
    "It's sometimes useful to create a component where the units of its inputs and/or outputs are\n",
    "determined by their connections.  This allows us to create components representing general\n",
    "purpose vector or matrix operations such as norms, summations, integrators, etc., that set their\n",
    "units appropriately based on the model that they're added to.\n",
    "\n",
    "Turning on dynamic unit computation is straightforward.  You just specify `units_by_conn`, `copy_units`\n",
    "and/or `compute_units` in your `add_input` or `add_output` calls when you add variables\n",
    "to your component.\n",
    "\n",
    "Setting `units_by_conn=True` when adding an input or output variable will allow the units\n",
    "of that variable to be determined at runtime based on the variable that connects to it.\n",
    "\n",
    "Setting `copy_units=<var_name>`, where `<var_name>` is the local name of another variable in your\n",
    "component, will take the units of the variable specified in `<var_name>` and use those units\n",
    "for the variable you're adding.\n",
    "\n",
    "Setting `compute_units=<func>`, where `<func>` is a function taking a dict arg that maps variable\n",
    "names to PhysicalUnits and returning the computed units, will set the units of the variable you're adding\n",
    "as a function of the other variables in the same component of the opposite io type.  For example,\n",
    "setting `compute_units` for an output `z` on a component with inputs `x` and `y`, would cause the\n",
    "supplied function to be called with a dict of the form {`x`: x_units, `y`: y_units}, so\n",
    "the computed units of `z` could be a function of `x_units` and `y_units`.  Note that the \n",
    "compute_units function is not called until all units of the opposite io type are known for that \n",
    "component.\n",
    "\n",
    "PhysicalUnits can be combined into expressions, and the result of the expression will be a \n",
    "PhysicalUnit with the correct combined units.  This is generally the simplest way to compute the\n",
    "units inside of `compute_units`.  To get the actual unit string, you can use the `name()` method\n",
    "of the PhysicalUnit, but generally you won't need to do that since you can just return the PhysicalUnit\n",
    "directly and the framework will convert it to a unit string automatically.\n",
    "\n",
    "Note that `units_by_conn` can be specified for outputs as well as for inputs, as can `copy_units`\n",
    "and `compute_units`.\n",
    "This means that units information can propagate through the model in either forward or reverse. If\n",
    "you specify both `units_by_conn` and either `copy_units` or `compute_units` for your component's \n",
    "variables, it will allow their units to be resolved whether known units have \n",
    "been defined upstream or downstream of your component in the model.\n",
    "\n",
    "The following component with input `x` and output `y` can have its units set by known units \n",
    "that are either upstream or downstream. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acute-wellington",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class DynPartialsComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('x', units_by_conn=True, copy_units='y')\n",
    "        self.add_output('y', units_by_conn=True, copy_units='x')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = inputs['x'] * 3."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "involved-advantage",
   "metadata": {},
   "source": [
    "The following example demonstrates the flow of units information in the forward direction, where the IndepVarComp has known units, and the DynPartialsComp and the ExecComp have units set dynamically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "preliminary-liver",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = om.Problem()\n",
    "p.model.add_subsystem('indeps', om.IndepVarComp('x', units='ft'))\n",
    "p.model.add_subsystem('comp', DynPartialsComp())\n",
    "sink = p.model.add_subsystem('sink', om.ExecComp('y=x',\n",
    "                                                 x={'units_by_conn': True, 'copy_units': 'y'},\n",
    "                                                 y={'units_by_conn': True, 'copy_units': 'x'}))\n",
    "p.model.connect('indeps.x', 'comp.x')\n",
    "p.model.connect('comp.y', 'sink.x')\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "print(sink._get_var_meta('y', 'units'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5ae0ae9",
   "metadata": {
    "hide_input": true,
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert sink._get_var_meta('y', 'units') == 'ft'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "welcome-white",
   "metadata": {},
   "source": [
    "And the following shows units information flowing in reverse, from the known units of `sink.x` to the unknown units of the output `comp.y`, then to the input `comp.x`, then on to the connected auto-IndepVarComp output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "modified-allowance",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = om.Problem()\n",
    "comp = p.model.add_subsystem('comp', DynPartialsComp())\n",
    "p.model.add_subsystem('sink', om.ExecComp('y=x', units='m'))\n",
    "p.model.connect('comp.y', 'sink.x')\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "print(comp._get_var_meta('x', 'units'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd15ab26",
   "metadata": {
    "hide_input": true,
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "assert comp._get_var_meta('x', 'units') == 'm'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1187cc1c",
   "metadata": {},
   "source": [
    "Finally, an example use of `compute_units` is shown below.  We have a component with dynamic units that multiplies two matrices, so the output `O` units are a combination of the units of both inputs, `M` and `N`.  In this case we use a lambda function to compute the output units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af9ef572",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DynComputeComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        self.add_input('M', units_by_conn=True)\n",
    "        self.add_input('N', units_by_conn=True)\n",
    "\n",
    "        # use a lambda function to compute the output units based on the input units\n",
    "        self.add_output('O', compute_units=lambda units: units['M'] * units['N'])\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['O'] = inputs['M'] @ inputs['N']\n",
    "\n",
    "p = om.Problem()\n",
    "indeps = p.model.add_subsystem('indeps', om.IndepVarComp())\n",
    "indeps.add_output('M', units='ft')\n",
    "indeps.add_output('N', units='lbf')\n",
    "comp = p.model.add_subsystem('comp', DynComputeComp())\n",
    "p.model.connect('indeps.M', 'comp.M')\n",
    "p.model.connect('indeps.N', 'comp.N')\n",
    "p.setup()\n",
    "p.run_model()\n",
    "print('input units:', comp._get_var_meta('M', 'units'), 'and', comp._get_var_meta('N', 'units'))\n",
    "print('output units:', comp._get_var_meta('O', 'units'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tight-mining",
   "metadata": {},
   "source": [
    "## Debugging\n",
    "\n",
    "Sometimes, when the units of some variables are unresolvable, it can be difficult to understand\n",
    "why.  There is an OpenMDAO command line tool, `openmdao view_dyn_units`, that can be used to\n",
    "show a graph of the variables with dynamic units and any variables with known units that\n",
    "connect directly to them.  Each node in the graph is a variable, and each edge is a connection\n",
    "between that variable and another.  Note that this connection does not have to be a\n",
    "connection in the normal OpenMDAO sense.  It could be a connection internal to a component\n",
    "created by declaring a `copy_units` in the metadata of one variable that refers to another\n",
    "variable.\n",
    "\n",
    "The nodes in the graph are colored to make it easier to locate static/dynamic/unresolved\n",
    "variable units.  Variables with 'static' known units are colored green, variables with dynamic\n",
    "units that have been resolved are colored blue, and any variables with unresolved units\n",
    "are colored red.  Each node is labeled with the units of the variable, if known, or a '?' if\n",
    "unknown, followed by the absolute pathname of the variable in the model.\n",
    "\n",
    "The plot is somewhat crude and the node labels sometimes overlap, but it's possible to zoom\n",
    "in to part of the graph to make it more readable using the button that looks like a magnifying glass.\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "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.10"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
