{
 "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  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MuxComp\n",
    "\n",
    "`MuxComp` works to combine\n",
    "multiple inputs into a single value.  This can be useful in situations where scalar outputs\n",
    "from multiple components need to be fed into a single vectorized component.\n",
    "\n",
    "`MuxComp` combines two or more inputs into a single output by stacking them along an axis.\n",
    "\n",
    "## MuxComp Options\n",
    "\n",
    "`MuxComp` has a single option, `vec_size`, which provides the number of inputs to be\n",
    "combined into a single output. The default value of `vec_size` is 2.\n",
    "\n",
    "## Adding Variables\n",
    "\n",
    "A single `MuxComp` can mux 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 is to occur is given via the `axis` argument.\n",
    "\n",
    "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",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.mux_comp.MuxComp.add_var\n",
    "        :noindex:\n",
    "```"
   ]
  },
  {
   "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": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "# 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": [
    "from openmdao.utils.assert_utils import assert_near_equal\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",
    "    assert_near_equal(p.get_val('vec_mag_comp.r_mag')[i], expected_i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(om_slicer)=\n",
    "## Example: Using `om.slicer` to reduce a 3-column matrix into constituent vectors\n",
    "\n",
    "The inverse functionality, that is \"demuxing\" or breaking up of inputs into multiple values, is accomplished using the `om.slicer()` utility.\n",
    "\n",
    "This example is contrived and could be achieved with a single vectorized component, but it serves to give an example of how this would be implemented."
   ]
  },
  {
   "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",
    "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])"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.0"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
