{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%pylab inline\n",
    "from ipyparallel import Client, error\n",
    "cluster=Client(profile=\"mpi\")\n",
    "view=cluster[:]\n",
    "view.block=True\n",
    "\n",
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Using src_indices with Promoted Variables\n",
    "\n",
    "Inputs and outputs can be connected by promoting them both to the same name, but what\n",
    "if your output is an array and you only want to connect part of it to your\n",
    "input?  \n",
    "\n",
    "If you connect variables via promotion, you must set `src_indices` when you promote the input with the [promotes](group-promotion) method,   Another\n",
    "argument, `flat_src_indices` is a boolean that determines whether the entries\n",
    "of the `src_indices` array are interpreted as indices into the flattened source\n",
    "or as indices into the unflattened source.  The default\n",
    "of `flat_src_indices=False` assumes indices map into an unflattened source.  Note also that if `src_indices` are some form of mult-index, for example `om.slicer[:, :]` or `([0,1,2,3], [0,1,2,3])`, then setting \n",
    "`flat_src_indices=True` is invalid and will raise an exception.\n",
    "\n",
    "\n",
    "## Basic Example\n",
    "\n",
    "Here is a simple example showing how to connect an independent array variable to two different components where each component gets part of the array."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class MyComp1(om.ExplicitComponent):\n",
    "    \"\"\" multiplies input array by 2. \"\"\"\n",
    "    def setup(self):\n",
    "        self.add_input('x', np.ones(3))\n",
    "        self.add_output('y', 1.0)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = np.sum(inputs['x'])*2.0\n",
    "\n",
    "class MyComp2(om.ExplicitComponent):\n",
    "    \"\"\" multiplies input array by 4. \"\"\"\n",
    "    def setup(self):\n",
    "        self.add_input('x', np.ones(2))\n",
    "        self.add_output('y', 1.0)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = np.sum(inputs['x'])*4.0\n",
    "\n",
    "class MyGroup(om.Group):\n",
    "    def setup(self):\n",
    "        self.add_subsystem('comp1', MyComp1())\n",
    "        self.add_subsystem('comp2', MyComp2())\n",
    "\n",
    "    def configure(self):\n",
    "        # splits input via promotes using src_indices\n",
    "        self.promotes('comp1', inputs=['x'], src_indices=[0, 1, 2])\n",
    "        self.promotes('comp2', inputs=['x'], src_indices=[3, 4])\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "p.model.set_input_defaults('x', np.ones(5))\n",
    "p.model.add_subsystem('G1', MyGroup(), promotes_inputs=['x'])\n",
    "\n",
    "p.setup()\n",
    "inp = np.random.random(5)\n",
    "p.set_val('x', inp)\n",
    "p.run_model()\n",
    "\n",
    "print(p.get_val('G1.comp1.x'))\n",
    "print(p.get_val('G1.comp2.x'))\n",
    "print(p.get_val('G1.comp1.y'))\n",
    "print(p.get_val('G1.comp2.y'))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "assert_near_equal(p.get_val('G1.comp1.x'), inp[:3])\n",
    "assert_near_equal(p.get_val('G1.comp2.x'), inp[3:])\n",
    "assert_near_equal(p.get_val('G1.comp1.y'), np.sum(inp[:3]*2))\n",
    "assert_near_equal(p.get_val('G1.comp2.y'), np.sum(inp[3:]*4))"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Using src_indices with 2D Arrays\n",
    "\n",
    "In this example, the source array is shape (4,3) and the input array is shape (2,2)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "class MyComp(om.ExplicitComponent):\n",
    "    def setup(self):\n",
    "        # We want to pull the following 4 values out of the source:\n",
    "        # (0,0), (3,1), (2,1), and (1,1).  To do that, we use numpy style\n",
    "        # indexing like this: ([[0, 3], [2, 1]], [[0, 1], [1, 1]]).\n",
    "        # We've  split up our indexing into a row array and a column array, and we've\n",
    "        # shaped those arrays to be 2x2, so the result of applying our src_indices to\n",
    "        # the source array will be a 2x2 array.\n",
    "        self.add_input('x', np.ones((2, 2)))\n",
    "        self.add_output('y', 1.0)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = np.sum(inputs['x'])\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "# by promoting the following output and inputs to 'x', they will\n",
    "# be automatically connected\n",
    "p.model.add_subsystem('indep', om.IndepVarComp('x', np.arange(12).reshape((4, 3))),\n",
    "                      promotes_outputs=['x'])\n",
    "p.model.add_subsystem('C1', MyComp())\n",
    "p.model.promotes('C1', inputs=['x'], src_indices=([[0, 3], [2, 1]], [[0, 1], [1, 1]]))\n",
    "\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "print(p.get_val('C1.x'))\n",
    "print(p.get_val('C1.y'))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "assert_near_equal(p.get_val('C1.x'),\n",
    "                 np.array([[0., 10.],\n",
    "                           [7., 4.]]))\n",
    "assert_near_equal(p.get_val('C1.y'), 21.)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "If the source array is shape (4,3), the input is scalar, and we want to\n",
    "connect it to the (3, 1) entry of the source, then the `promotes`\n",
    "call might look like the following if we use `flat_src_indices`:\n",
    "\n",
    "```\n",
    "    p.model.promotes('C1', inputs=['x'], src_indices=[10], shape=1, flat_src_indices=True)\n",
    "```\n",
    "\n",
    "If we instead use the default setting of `flat_src_indices=False`,  we would just access the\n",
    "(3, 1) entry as expected.\n",
    "\n",
    "```\n",
    "    p.model.promotes('C1', inputs=['x'], src_indices=(3, 1), shape=1)\n",
    "\n",
    "```"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "5: If the source array is flat and the input is shape (2,2), the `promotes` call might look like this:\n",
    "```\n",
    "    p.model.promotes('C1', inputs=['x'], src_indices=[0, 10, 7, 4], shape=(2,2))\n",
    "```"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "```{note}\n",
    "If the source array is flat, we allow the use of flat src_indices even without \n",
    "setting `flat_src_indices=True`.\n",
    "```"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Distributed component example\n",
    "\n",
    "In the example, a distributed component promotes its input and receives certain entries of the source array based on its rank.  Note that negative indices are supported."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%%px\n",
    "\n",
    "import openmdao.api as om\n",
    "import numpy as np\n",
    "\n",
    "class MyComp(om.ExplicitComponent):\n",
    "    def __init__(self, idxs, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.idxs = idxs\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', np.ones(len(self.idxs)))\n",
    "        self.add_output('y', 1.0)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = np.sum(inputs['x'])*2.0\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('indep', om.IndepVarComp('x', np.arange(5, dtype=float)),\n",
    "                      promotes_outputs=['x'])\n",
    "\n",
    "# decide what parts of the array we want based on our rank\n",
    "if p.comm.rank == 0:\n",
    "    idxs = [0, 1, 2]\n",
    "else:\n",
    "    # use [3, -1] here rather than [3, 4] just to show that we\n",
    "    # can use negative indices.\n",
    "    idxs = [3, -1]\n",
    "\n",
    "p.model.add_subsystem('C1', MyComp(idxs))\n",
    "\n",
    "p.model.promotes('C1', inputs=['x'], src_indices=idxs)\n",
    "\n",
    "p.setup()\n",
    "p.set_val('x', np.arange(5, dtype=float))\n",
    "p.run_model()\n",
    "\n",
    "# each rank holds the assigned portion of the input array\n",
    "print(p.get_val('C1.x'))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%%px\n",
    "\n",
    "# the output in each rank is based on the local inputs\n",
    "print(p.get_val('C1.y'))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%%px\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(p['C1.x'],\n",
    "                  np.arange(3, dtype=float) if p.model.C1.comm.rank == 0 else np.arange(3, 5, dtype=float))\n",
    "\n",
    "# the output in each rank is based on the local inputs\n",
    "assert_near_equal(p['C1.y'], 6. if p.model.C1.comm.rank == 0 else 14.)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.5"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}