{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "eace0044",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:40.768205Z",
     "iopub.status.busy": "2026-10-02T14:41:40.768041Z",
     "iopub.status.idle": "2026-10-02T14:41:40.772178Z",
     "shell.execute_reply": "2026-10-02T14:41:40.771617Z"
    },
    "papermill": {
     "duration": 0.006499,
     "end_time": "2026-10-02T14:41:40.772799+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:40.766300+00:00",
     "status": "completed"
    },
    "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",
   "id": "9db1af23",
   "metadata": {
    "papermill": {
     "duration": 0.041823,
     "end_time": "2026-10-02T14:41:40.815650+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:40.773827+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "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",
   "id": "620c43f1",
   "metadata": {
    "papermill": {
     "duration": 0.0007,
     "end_time": "2026-10-02T14:41:40.817232+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:40.816532+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "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": 2,
   "id": "35c93bcd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:40.820347Z",
     "iopub.status.busy": "2026-10-02T14:41:40.820213Z",
     "iopub.status.idle": "2026-10-02T14:41:43.187645Z",
     "shell.execute_reply": "2026-10-02T14:41:43.187154Z"
    },
    "papermill": {
     "duration": 2.370425,
     "end_time": "2026-10-02T14:41:43.188391+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:40.817966+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1790952103.170866] [runnervm8df0l:6767 :0]        ib_iface.c:1269 UCX  ERROR mana_0: iface 0x5595e7acc040 failed to create UD QP TX wr:256 sge:6 inl:64 resp:0 RX wr:4096 sge:1 resp:0 failed: Operation not supported\n",
      "[1790952103.171140] [runnervm8df0l:6767 :0]      ucp_worker.c:1412 UCX  ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n",
      "2.7220229687070043\n",
      "2.7219026493045364\n",
      "2.765173358222124\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[runnervm8df0l:06767] pml_ucx.c:313  Error: Failed to create UCP worker\n"
     ]
    }
   ],
   "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": 3,
   "id": "f89ecc19",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:43.195162Z",
     "iopub.status.busy": "2026-10-02T14:41:43.194864Z",
     "iopub.status.idle": "2026-10-02T14:41:43.198778Z",
     "shell.execute_reply": "2026-10-02T14:41:43.198329Z"
    },
    "papermill": {
     "duration": 0.009481,
     "end_time": "2026-10-02T14:41:43.199260+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:43.189779+00:00",
     "status": "completed"
    },
    "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",
   "id": "da7a57e2",
   "metadata": {
    "papermill": {
     "duration": 0.001113,
     "end_time": "2026-10-02T14:41:43.201552+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:43.200439+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "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": 4,
   "id": "7b612ba5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:43.204629Z",
     "iopub.status.busy": "2026-10-02T14:41:43.204390Z",
     "iopub.status.idle": "2026-10-02T14:41:43.213095Z",
     "shell.execute_reply": "2026-10-02T14:41:43.212665Z"
    },
    "papermill": {
     "duration": 0.010786,
     "end_time": "2026-10-02T14:41:43.213461+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:43.202675+00:00",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1. 2. 3.]\n",
      "[4. 5. 6.]\n",
      "[1.32581766 1.19028995 1.10714872]\n"
     ]
    }
   ],
   "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": 5,
   "id": "9d3b3f70",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-10-02T14:41:43.217473Z",
     "iopub.status.busy": "2026-10-02T14:41:43.217318Z",
     "iopub.status.idle": "2026-10-02T14:41:43.219725Z",
     "shell.execute_reply": "2026-10-02T14:41:43.219272Z"
    },
    "papermill": {
     "duration": 0.005498,
     "end_time": "2026-10-02T14:41:43.220163+00:00",
     "exception": false,
     "start_time": "2026-10-02T14:41:43.214665+00:00",
     "status": "completed"
    },
    "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.13.14"
  },
  "orphan": true,
  "papermill": {
   "default_parameters": {},
   "duration": 3.751094,
   "end_time": "2026-10-02T14:41:43.839857+00:00",
   "environment_variables": {},
   "exception": null,
   "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/features/building_blocks/components/mux_comp.ipynb",
   "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/features/building_blocks/components/mux_comp.ipynb",
   "parameters": {},
   "start_time": "2026-10-02T14:41:40.088763+00:00",
   "version": "2.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}