{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "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": [
    "# Connecting Variables\n",
    "\n",
    "To cause data to flow between two systems in a model, we must connect at\n",
    "least one output variable from one system to at least one input variable\n",
    "from the other.  If the variables have units defined, then the framework\n",
    "will automatically perform the conversion.  We can also connect only part\n",
    "of an array output to an input by specifying the indices of the entries\n",
    "that we want.\n",
    "\n",
    "To connect two variables within a model, use the `connect` function.\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.group.Group.connect\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## Usage\n",
    "\n",
    "1: Connect an output variable to an input variable, with an automatic unit conversion.  "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "p.model.set_input_defaults('x', np.ones(5), units='ft')\n",
    "\n",
    "exec_comp = om.ExecComp('y=sum(x)',\n",
    "                        x={'val': np.zeros(5), 'units': 'inch'},\n",
    "                        y={'units': 'inch'})\n",
    "\n",
    "p.model.add_subsystem('comp1', exec_comp, promotes_inputs=['x'])\n",
    "\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "print(p.get_val('x', units='ft'))\n",
    "print(p.get_val('comp1.x'))\n",
    "print(p.get_val('comp1.y'))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(p.get_val('x', units='ft'), np.ones(5))\n",
    "assert_near_equal(p.get_val('comp1.x'), np.ones(5)*12.)\n",
    "assert_near_equal(p.get_val('comp1.y'), 60.)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "2: Connect one output to many inputs.  "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('C1', om.ExecComp('y=sum(x)*2.0', x=np.zeros(5)), promotes_inputs=['x'])\n",
    "p.model.add_subsystem('C2', om.ExecComp('y=sum(x)*4.0', x=np.zeros(5)), promotes_inputs=['x'])\n",
    "p.model.add_subsystem('C3', om.ExecComp('y=sum(x)*6.0', x=np.zeros(5)), promotes_inputs=['x'])\n",
    "\n",
    "p.setup()\n",
    "p.set_val('x', np.ones(5))\n",
    "p.run_model()\n",
    "\n",
    "print(p.get_val('C1.y'))\n",
    "print(p.get_val('C2.y'))\n",
    "print(p.get_val('C3.y'))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "assert_near_equal(p.get_val('C1.y'), 10.)\n",
    "assert_near_equal(p.get_val('C2.y'), 20.)\n",
    "assert_near_equal(p.get_val('C3.y'), 30.)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "(connect-with-src-indices)=\n",
    "3: Connect only part of an array output to an input of a smaller size."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('indep', om.IndepVarComp('x', np.ones(5)))\n",
    "p.model.add_subsystem('C1', om.ExecComp('y=sum(x)*2.0', x=np.zeros(3)))\n",
    "p.model.add_subsystem('C2', om.ExecComp('y=sum(x)*4.0', x=np.zeros(2)))\n",
    "\n",
    "# connect C1.x to the first 3 entries of indep.x\n",
    "p.model.connect('indep.x', 'C1.x', src_indices=[0, 1, 2])\n",
    "\n",
    "# connect C2.x to the last 2 entries of indep.x\n",
    "# use -2 (same as 3 in this case) to show that negative indices work.\n",
    "p.model.connect('indep.x', 'C2.x', src_indices=[-2, 4])\n",
    "\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "print(p['C1.x'])\n",
    "print(p['C1.y'])\n",
    "print(p['C2.x'])\n",
    "print(p['C2.y'])"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "assert_near_equal(p['C1.x'], np.ones(3))\n",
    "assert_near_equal(p['C1.y'], 6.)\n",
    "assert_near_equal(p['C2.x'], np.ones(2))\n",
    "assert_near_equal(p['C2.y'], 8.)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "4: Connect only part of a non-flat array output to a non-flat array input."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('indep', om.IndepVarComp('x', np.arange(12).reshape((4, 3))))\n",
    "p.model.add_subsystem('C1', om.ExecComp('y=sum(x)*2.0', x=np.zeros((2, 2))))\n",
    "\n",
    "# Connect C1.x to entries (0,0), (-1,1), (2,1), (1,1) of indep.x and give then a 2x2 shape.\n",
    "# To do this, create 2x2 shaped arrays representing rows and columns. Note that the final\n",
    "# src_indices is a tuple (not a list or array) containing the rows and cols arrays.\n",
    "rows = [[0, -1],[2, 1]]\n",
    "cols = [[0, 1], [1, 1]]\n",
    "p.model.connect('indep.x', 'C1.x', src_indices=(rows, cols), flat_src_indices=False)\n",
    "\n",
    "p.setup()\n",
    "p.run_model()\n",
    "\n",
    "print(p.get_val('indep.x'))\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('indep.x'), np.array([[0., 1., 2.],\n",
    "                                                  [3., 4., 5.],\n",
    "                                                  [6., 7., 8.],\n",
    "                                                  [9., 10., 11.]]))\n",
    "assert_near_equal(p.get_val('C1.x'), np.array([[0., 10.],\n",
    "                                               [7., 4.]]))\n",
    "assert_near_equal(p.get_val('C1.y'), 42.)"
   ],
   "outputs": [],
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   }
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.10 64-bit ('py38mpich_new2': conda)"
  },
  "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,
  "interpreter": {
   "hash": "49f71f8b029ccfc892068b15e7b6668703dcc2d714d079581384c7016c112935"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}