{
 "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\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MuxComp and DemuxComp\n",
    "\n",
    "```{Note}\n",
    "DemuxComp is being deprecated in favor of `om.slicer` for a safer and more robust solution. [Below](om_slicer) is an example of how to use `om.slicer` to achieve the same result.\n",
    "```\n",
    "\n",
    "`DemuxComp` and `MuxComp` work together to break up inputs into multiple values (demux) or combine\n",
    "multiple inputs into a single value (mux).  This can be useful in situations where scalar outputs\n",
    "from multiple components need to be fed into a single vectorized component.\n",
    "\n",
    "`DemuxComp` takes a single input of arbitary shape (the size of at least one axis must be equal\n",
    "to `vec_size`).  It can then be broken along that axis, resulting in `vec_size` outputs.\n",
    "\n",
    "`MuxComp` combines two or more inputs into a single output by stacking them along an axis.\n",
    "\n",
    "## MuxComp and DemuxComp Options\n",
    "\n",
    "These components have a single option, `vec_size`, which provides the number of inputs to be\n",
    "combined into a single output (for `MuxComp`) or the number of outputs into which an input is\n",
    "to be broken (for `DemuxComp`).  The default value of `vec_size` is 2.\n",
    "\n",
    "## Adding Variables\n",
    "\n",
    "A single `MuxComp` or `DemuxComp` can mux or demux multiple variables, so long as all variables\n",
    "are compatible with the given `vec_size`.  Variables are added via the `add_var` method.\n",
    "\n",
    "The axis along which the muxing/demuxing is to occur is given via the axis argument.\n",
    "For DemuxComp, the specified axis index must be the index of one of the input dimensions (you cannot demux along axis 3 of a 2D input).\n",
    "In addition, the axis on which the Demuxing is to be done must have length `vec_size`.\n",
    "\n",
    "For MuxComp, the variables are joined along a new dimension, the index of which is given by axis.\n",
    "The specified axis follows the convention used by the `numpy.stack` function.\n",
    "Giving `axis = 0` will stack the inputs along the first axis (vertically).\n",
    "Giving `axis = 1` will stack the inputs along the second axis (horizontally).\n",
    "Giving `axis = -1` will stack the inputs along the last axis, and so is dependent on the shape of the inputs.\n",
    "Due to the axis convention of `numpy.stack`, the axis index is only valid if it is less than or\n",
    "equal to the number of dimensions in the inputs.\n",
    "For example, 1D arrays can be stacked vertically (`axis = 0`) or horizontally (`axis = 1`), but not\n",
    "depth-wise (`axis = 2`).\n",
    "\n",
    "For DemuxComp, the name of the given variable is the **input**.  It is demuxed into variables whose\n",
    "names are appended with `_n` where `n` is an integer from 0 through `vec_size`-1.\n",
    "Conversely, for MuxComp, the given variable name is the output, and each input is appended with `_n`.\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.mux_comp.MuxComp.add_var\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.demux_comp.DemuxComp.add_var\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## Example: Demuxing a 3-column matrix into constituent vectors\n",
    "\n",
    "This example is contrived and could be achieved with a single vectorized component, but it serves\n",
    "to give an example to the capabilities of the Demux component.  Given a position vector in the\n",
    "Earth-centered, Earth-fixed (ECEF) frame (n x 3), extract the three (n x 1) columns from the matrix\n",
    "and use the first two to compute the longitude at the given position vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "# The number of elements to be demuxed\n",
    "n = 3\n",
    "\n",
    "# The size of each element to be demuxed\n",
    "m = 100\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "demux_comp = p.model.add_subsystem(name='demux', subsys=om.DemuxComp(vec_size=n),\n",
    "                                 promotes_inputs=['pos_ecef'])\n",
    "\n",
    "demux_comp.add_var('pos_ecef', shape=(m, n), axis=1, units='km')\n",
    "\n",
    "p.model.add_subsystem(name='longitude_comp',\n",
    "                      subsys=om.ExecComp('long = atan(y/x)',\n",
    "                                         x={'val': np.ones(m), 'units': 'km'},\n",
    "                                         y={'val': np.ones(m), 'units': 'km'},\n",
    "                                         long={'val': np.ones(m), 'units': 'rad'}))\n",
    "\n",
    "p.model.connect('demux.pos_ecef_0', 'longitude_comp.x')\n",
    "p.model.connect('demux.pos_ecef_1', 'longitude_comp.y')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('pos_ecef', 6378 * np.cos(np.linspace(0, 2*np.pi, m)), indices=om.slicer[:, 0])\n",
    "p.set_val('pos_ecef', 6378 * np.sin(np.linspace(0, 2*np.pi, m)), indices=om.slicer[:, 1])\n",
    "p.set_val('pos_ecef', 0.0, indices=om.slicer[:, 2])\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "expected = np.arctan(p.get_val('pos_ecef', indices=om.slicer[:, 1]) / p.get_val('pos_ecef', indices=om.slicer[:, 0]))\n",
    "\n",
    "print(p.get_val('longitude_comp.long'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(p.get_val('longitude_comp.long'), expected)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(om_slicer)=\n",
    "## Example: Using `om.slicer` to reduce a 3-column matrix into constituent vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The number of elements to be demuxed\n",
    "n = 3\n",
    "\n",
    "arr_5x3 = np.array([\n",
    "    [1, 2, 3],\n",
    "    [4, 5, 6],\n",
    "    [7, 8, 9],\n",
    "    [10, 11, 12],\n",
    "    [13, 14, 15],\n",
    "])\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "p.model.add_subsystem('indep', om.IndepVarComp('x', arr_5x3, units='km'), promotes=['*'])\n",
    "p.model.add_subsystem('indep2', om.IndepVarComp('y', arr_5x3, units='km'), promotes=['*'])\n",
    "\n",
    "p.model.add_subsystem(name='longitude_comp',\n",
    "                      subsys=om.ExecComp('long = atan(y/x)',\n",
    "                                         x={'val': np.ones(n), 'units': 'km'},\n",
    "                                         y={'val': np.ones(n), 'units': 'km'},\n",
    "                                         long={'val': np.ones(n), 'units': 'rad'}))\n",
    "\n",
    "# Use the src_indices arg in promotes to perform the demuxing\n",
    "p.model.promotes('longitude_comp', inputs=['x'], src_indices=om.slicer[0, :])\n",
    "p.model.promotes('longitude_comp', inputs=['y'], src_indices=om.slicer[1, :])\n",
    "\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "print(p.get_val('longitude_comp.x'))\n",
    "print(p.get_val('longitude_comp.y'))\n",
    "print(p.get_val('longitude_comp.long'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert(list(p.get_val('longitude_comp.x')) == [1, 2, 3])\n",
    "assert(list(p.get_val('longitude_comp.y')) == [4, 5, 6])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Muxing 3 (n x 1) columns into a single (n x 3) matrix\n",
    "\n",
    "In this example we start with three (n x 1) column vectors (`x`, `y`, and `z`) and\n",
    "combine them into a single position vector `r` (n x 3).  This is achieved by stacking the vectors\n",
    "along `axis = 1`.  Like the previous example, this is somewhat contrived but is intended to demonstrate\n",
    "the capabilities of the MuxComp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The number of elements to be muxed\n",
    "n = 3\n",
    "\n",
    "# The size of each element to be muxed\n",
    "m = 100\n",
    "\n",
    "p = om.Problem()\n",
    "\n",
    "mux_comp = p.model.add_subsystem(name='mux', subsys=om.MuxComp(vec_size=n))\n",
    "\n",
    "mux_comp.add_var('r', shape=(m,), axis=1, units='m')\n",
    "\n",
    "p.model.add_subsystem(name='vec_mag_comp',\n",
    "                      subsys=om.VectorMagnitudeComp(vec_size=m, length=n, in_name='r',\n",
    "                                                    mag_name='r_mag', units='m'))\n",
    "\n",
    "p.model.connect('mux.r', 'vec_mag_comp.r')\n",
    "\n",
    "p.setup()\n",
    "\n",
    "p.set_val('mux.r_0', 1 + np.random.rand(m))\n",
    "p.set_val('mux.r_1', 1 + np.random.rand(m))\n",
    "p.set_val('mux.r_2', 1 + np.random.rand(m))\n",
    "\n",
    "p.run_model()\n",
    "\n",
    "# Verify the results against numpy.dot in a for loop.\n",
    "for i in range(n):\n",
    "    r_i = [p.get_val('mux.r_0')[i], p.get_val('mux.r_1')[i], p.get_val('mux.r_2')[i]]\n",
    "    expected_i = np.sqrt(np.dot(r_i, r_i))\n",
    "    print(p.get_val('vec_mag_comp.r_mag')[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "# Verify the results against numpy.dot in a for loop.\n",
    "for i in range(n):\n",
    "    r_i = [p.get_val('mux.r_0')[i], p.get_val('mux.r_1')[i], p.get_val('mux.r_2')[i]]\n",
    "    expected_i = np.sqrt(np.dot(r_i, r_i))\n",
    "    assert_near_equal(p.get_val('vec_mag_comp.r_mag')[i], expected_i)"
   ]
  }
 ],
 "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.1"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
