{
 "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\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "allied-means",
   "metadata": {},
   "source": [
    "# Determining Variable Shapes at Runtime\n",
    "\n",
    "It's sometimes useful to create a component where the shapes 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 size\n",
    "themselves appropriately based on the model that they're added to.\n",
    "\n",
    "Turning on dynamic shape computation is straightforward.  You just specify `shape_by_conn`\n",
    "and/or `copy_shape` in your `add_input` or `add_output` calls when you add variables\n",
    "to your component.\n",
    "\n",
    "Setting `shape_by_conn=True` when adding an input or output variable will allow the shape\n",
    "of that variable to be determined at runtime based on the variable that connects to it.\n",
    "\n",
    "Setting `copy_shape=<var_name>`, where `<var_name>` is the local name of another variable in your\n",
    "component, will take the shape of the variable specified in `<var_name>` and use that\n",
    "shape for the variable you're adding.  \n",
    "\n",
    "Note that `shape_by_conn` can be specified for outputs as well as for inputs, as can `copy_shape`.\n",
    "This means that shape information can propagate through the model in either forward or reverse. If\n",
    "you specify both `shape_by_conn` and `copy_shape` for your component's dynamically shaped variables, \n",
    "it will allow your their shapes to be resolved whether known shapes have been defined upstream or \n",
    "downstream of your component in the model.\n",
    "\n",
    "For example, the following component with input `x` and output `y` can have its shapes set by known shapes \n",
    "that are either upstream or downstream. Note that this component also has sparse partials, diagonal in this case, \n",
    "and those are specified within the `setup_partials` method which is called after all shapes have been computed.\n",
    "It uses the `_get_var_meta` method to get the size of its variables in order to determine the size\n",
    "of the partials.\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', shape_by_conn=True, copy_shape='y')\n",
    "        self.add_output('y', shape_by_conn=True, copy_shape='x')\n",
    "\n",
    "    def setup_partials(self):\n",
    "        size = self._get_var_meta('x', 'size')\n",
    "        self.mat = np.eye(size) * 3.\n",
    "        rng = np.arange(size)\n",
    "        self.declare_partials('y', 'x', rows=rng, cols=rng, val=3.0)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = self.mat.dot(inputs['x'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "involved-advantage",
   "metadata": {},
   "source": [
    "The following example demonstrates the flow of shape information in the forward direction, where the IndepVarComp has a known size, and the DynPartialsComp and the ExecComp are sized dynamically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "preliminary-liver",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('indeps', om.IndepVarComp('x', val=np.ones(5)))\n",
    "p.model.add_subsystem('comp', DynPartialsComp())\n",
    "p.model.add_subsystem('sink', om.ExecComp('y=x',\n",
    "                                          x={'shape_by_conn': True, 'copy_shape': 'y'},\n",
    "                                          y={'shape_by_conn': True, 'copy_shape': '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",
    "J = p.compute_totals(of=['sink.y'], wrt=['indeps.x'])\n",
    "print(J['sink.y', 'indeps.x'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5ae0ae9",
   "metadata": {
    "hide_input": true,
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(J['sink.y', 'indeps.x'], np.eye(5) * 3.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "welcome-white",
   "metadata": {},
   "source": [
    "And the following shows shape information flowing in reverse, from the known shape of `sink.x` to the unknown shape 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": [
    "import numpy as np\n",
    "\n",
    "p = om.Problem()\n",
    "p.model.add_subsystem('comp', DynPartialsComp())\n",
    "p.model.add_subsystem('sink', om.ExecComp('y=x', shape=5))\n",
    "p.model.connect('comp.y', 'sink.x')\n",
    "p.setup()\n",
    "p.run_model()\n",
    "J = p.compute_totals(of=['sink.y'], wrt=['comp.x'])\n",
    "print(J['sink.y', 'comp.x'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd15ab26",
   "metadata": {
    "hide_input": true,
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb",
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(J['sink.y', 'comp.x'], np.eye(5) * 3.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tight-mining",
   "metadata": {},
   "source": [
    "## Debugging\n",
    "\n",
    "Sometimes, when the shapes of some variables are unresolvable, it can be difficult to understand\n",
    "why.  There is an OpenMDAO command line tool, `openmdao view_dyn_shapes`, that can be used to\n",
    "show a graph of the dynamically shaped variables and any statically shaped variables 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_shape` 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 shapes.  Statically shaped variables are colored green, dynamically shaped\n",
    "variables that have been resolved are colored blue, and any variables with unresolved shapes\n",
    "are colored red.  Each node is labeled with the shape 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",
    "\n",
    "Below is an example plot for a simple model with three components and no unresolved shapes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "instant-state",
   "metadata": {
    "hide_input": true,
    "scrolled": true,
    "tags": [
     "remove_input"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.visualization.dyn_shape_plot import view_dyn_shapes\n",
    "p = om.Problem()\n",
    "indep = p.model.add_subsystem('indep', om.IndepVarComp('x1', val=np.ones((2,3))))\n",
    "indep.add_output('x2', val=np.ones((4,2)))\n",
    "\n",
    "p.model.add_subsystem('C1', om.ExecComp('y1, y2 = x1*2, x2*2',\n",
    "                                                x1={'shape_by_conn': True, 'copy_shape': 'y1'},\n",
    "                                                x2={'shape_by_conn': True, 'copy_shape': 'y2'},\n",
    "                                                y1={'shape_by_conn': True, 'copy_shape': 'x1'},\n",
    "                                                y2={'shape_by_conn': True, 'copy_shape': 'x2'}))\n",
    "\n",
    "p.model.add_subsystem('C2', om.ExecComp('y1, y2 = x1*2, x2*2',\n",
    "                                                x1={'shape_by_conn': True, 'copy_shape': 'y1'},\n",
    "                                                x2={'shape_by_conn': True, 'copy_shape': 'y2'},\n",
    "                                                y1={'shape_by_conn': True, 'copy_shape': 'x1'},\n",
    "                                                y2={'shape_by_conn': True, 'copy_shape': 'x2'}))\n",
    "\n",
    "p.model.connect('indep.x1', 'C1.x1')\n",
    "p.model.connect('indep.x2', 'C1.x2')\n",
    "p.model.connect('C1.y1', 'C2.x1')\n",
    "p.model.connect('C1.y2', 'C2.x2')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "view_dyn_shapes(p.model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "residential-theme",
   "metadata": {},
   "source": [
    "## Connecting Non-Distributed and Distributed Variables\n",
    "\n",
    "Dynamically shaped connections between distributed outputs and non-distributed inputs are not allowed because OpenMDAO assumes data will be transferred locally when computing variable shapes.  Since non-distributed variables must be identical in size and value on all ranks where they exist, the distributed output would have to also be identical in size and value on all ranks.  If that is the case, then the output should just be non-distributed as well.\n",
    "\n",
    "Dynamically shaped connections between non-distributed outputs and distributed inputs are currently allowed, though their use is not recommended. Such connections require that all src_indices in all ranks of the distributed input are identical."
   ]
  }
 ],
 "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
}
