{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from ipyparallel import Client, error  # noqa: F401\n",
    "cluster=Client(profile=\"mpi\")\n",
    "view=cluster[:]\n",
    "view.block=True\n",
    "\n",
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Variables\n",
    "\n",
    "At times when you need to perform a computation using large input arrays, you may want to perform that computation in multiple processes, where each process operates on some subset of the input values. This may be done purely for performance reasons, or it may be necessary because the entire input will not fit in the memory of a single machine. In any case, this can be accomplished in OpenMDAO by declaring those inputs and outputs as distributed. By definition, a `distributed variable` is an input or output where each process contains only a part of the whole variable. Distributed variables are declared by setting the optional \"distributed\" argument to True when adding the variable to a component. A component that has at least one distributed variable can also be called a distributed component.\n",
    "\n",
    "Any variable that is not distributed is called a `non-distributed variable`. When the model is run under MPI, every process contains a copy of the entire non-distributed variable.  We also call these duplicated variables.\n",
    "\n",
    "We’ve already seen that by using [src_indices](connect-with-src-indices), we can connect an input to only a subset of an output variable. By giving different values for src_indices in each MPI process, we can distribute computations on a distributed output across the processes. All of the scenarios that involve connecting distributed and non-distributed variables are detailed in [Connections involving distributed variables](../working_with_groups/dist_serial.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "This feature requires MPI, and may not be able to be run on Colab or Binder.\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Example: Simple Component with Distributed Input and Output\n",
    "\n",
    "The following example shows how to create a simple component, *SimpleDistrib*, that takes a distributed variable as an input and computes a distributed output. The calculation is divided across the available processes, but the details of that division are not contained in the component.  In fact, the input is sized based on its connected source using the \"shape_by_conn\" argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class SimpleDistrib(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Distributed Input\n",
    "        self.add_input('in_dist', shape_by_conn=True, distributed=True)\n",
    "\n",
    "        # Distributed Output\n",
    "        self.add_output('out_dist', copy_shape='in_dist', distributed=True)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        x = inputs['in_dist']\n",
    "\n",
    "        # \"Computationally Intensive\" operation that we wish to parallelize.\n",
    "        f_x = x**2 - 2.0*x + 4.0\n",
    "\n",
    "        outputs['out_dist'] = f_x"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next part of the example, we take the `SimpleDistrib` component, place it into a model, and run it. Suppose the vector of data we want to process has 7 elements. We have 4 processors available for computation, so if we distribute them as evenly as we can, 3 procs can handle 2 elements each, and the 4th processor can pick up the last one.  OpenMDAO's utilities includes the `evenly_distrib_idxs` function which computes the sizes and offsets for all ranks. The sizes are used to determine how much of the array to allocate on any specific rank. The offsets are used to figure out where the local portion of the array starts, and in this example, is used to set the initial value properly.  In this case, the initial value for the full distributed input \"in_dist\" is a vector of 7 values between 3.0 and 9.0, and each processor has a 1 or 2 element piece of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "size = 7\n",
    "\n",
    "if MPI:\n",
    "    comm = MPI.COMM_WORLD\n",
    "    rank = comm.rank\n",
    "    sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "else:\n",
    "    # When running without MPI, the entire variable is on one proc.\n",
    "    rank = 0\n",
    "    sizes = {rank : size}\n",
    "    offsets = {rank : 0}\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "# Create a distributed source for the distributed input.\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", SimpleDistrib())\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "# Values on each rank.\n",
    "for var in ['indep.x_dist', 'D1.out_dist']:\n",
    "    print(var, prob.get_val(var))\n",
    "\n",
    "# Full gathered values.\n",
    "for var in ['indep.x_dist', 'D1.out_dist']:\n",
    "    print(var, prob.get_val(var, get_remote=True))\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val(var, get_remote=True), np.array([7., 12., 19., 28., 39., 52., 67.]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that we created a connection source 'x_dist' that passes its value to 'D1.in_dist'.  OpenMDAO requires a source for non-constant inputs, and usually creates one automatically as an output of a component referred to as an 'Auto-IVC'. However, the automatic creation is not supported for distributed variables.  We must manually create an `IndepVarComp` and connect it to our input. \n",
    "\n",
    "When using distributed variables, OpenMDAO can't always size the component inputs based on the shape of the connected source. In this example, the component determines its own split using `evenly_distrib_idxs`. This requires that the component know the full vector size, which is passed in via the option 'vec_size'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "class SimpleDistrib(om.ExplicitComponent):\n",
    "\n",
    "    def initialize(self):\n",
    "        self.options.declare('vec_size', types=int, default=1,\n",
    "                             desc=\"Total size of vector.\")\n",
    "\n",
    "    def setup(self):\n",
    "        comm = self.comm\n",
    "        rank = comm.rank\n",
    "\n",
    "        size = self.options['vec_size']\n",
    "        sizes, _ = evenly_distrib_idxs(comm.size, size)\n",
    "        mysize = sizes[rank]\n",
    "\n",
    "        # Distributed Input\n",
    "        self.add_input('in_dist', np.ones(mysize, float), distributed=True)\n",
    "\n",
    "        # Distributed Output\n",
    "        self.add_output('out_dist', np.ones(mysize, float), distributed=True)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        x = inputs['in_dist']\n",
    "\n",
    "        # \"Computationally Intensive\" operation that we wish to parallelize.\n",
    "        f_x = x**2 - 2.0*x + 4.0\n",
    "\n",
    "        outputs['out_dist'] = f_x\n",
    "\n",
    "\n",
    "size = 7\n",
    "\n",
    "if MPI:\n",
    "    comm = MPI.COMM_WORLD\n",
    "    rank = comm.rank\n",
    "    sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "else:\n",
    "    # When running without MPI, the entire variable is on one proc.\n",
    "    rank = 0\n",
    "    sizes = {rank : size}\n",
    "    offsets = {rank : 0}\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "# Create a distributed source for the distributed input.\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", SimpleDistrib(vec_size=size))\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "# Values on each rank.\n",
    "for var in ['indep.x_dist', 'D1.out_dist']:\n",
    "    print(var, prob.get_val(var))\n",
    "\n",
    "# Full gathered values.\n",
    "for var in ['indep.x_dist', 'D1.out_dist']:\n",
    "    print(var, prob.get_val(var, get_remote=True))\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val(var, get_remote=True), np.array([7., 12., 19., 28., 39., 52., 67.]))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Example: Distributed I/O and a Non-Distributed Input\n",
    "\n",
    "OpenMDAO supports both non-distributed and distributed I/O on the same component, so in this example, we expand the problem to include a non-distributed input. In this case, the non-distributed input also has a vector width of 7, but those values will be the same on each processor. This non-distributed input is included in the computation by taking the vector sum and adding it to the distributed output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "\n",
    "class MixedDistrib1(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Distributed Input\n",
    "        self.add_input('in_dist', shape_by_conn=True, distributed=True)\n",
    "\n",
    "        # Non-Distributed Input\n",
    "        self.add_input('in_nd', shape_by_conn=True)\n",
    "\n",
    "        # Distributed Output\n",
    "        self.add_output('out_dist', copy_shape='in_dist', distributed=True)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        Id = inputs['in_dist']\n",
    "        Ind = inputs['in_nd']\n",
    "\n",
    "        # \"Computationally Intensive\" operation that we wish to parallelize.\n",
    "        f_Id = Id**2 - 2.0*Id + 4.0\n",
    "\n",
    "        # This operation is repeated on all procs.\n",
    "        f_Ind = Ind ** 0.5\n",
    "\n",
    "        outputs['out_dist'] = f_Id + np.sum(f_Ind)\n",
    "\n",
    "size = 7\n",
    "\n",
    "if MPI:\n",
    "    comm = MPI.COMM_WORLD\n",
    "    rank = comm.rank\n",
    "    sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "else:\n",
    "    # When running without MPI, the entire variable is on one proc.\n",
    "    rank = 0\n",
    "    sizes = {rank : size}\n",
    "    offsets = {rank : 0}\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "# Create a distributed source for the distributed input.\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "ivc.add_output('x_nd', np.zeros(size))\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", MixedDistrib1())\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "model.connect('indep.x_nd', 'D1.in_nd')\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "# Set initial values of non-distributed variable.\n",
    "x_nd_init = 1.0 + 2.0*np.arange(size)\n",
    "prob.set_val('indep.x_nd', x_nd_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "# Values on each rank.\n",
    "for var in ['indep.x_dist', 'indep.x_nd', 'D1.out_dist']:\n",
    "    print(var, prob.get_val(var))\n",
    "\n",
    "# Full gathered values.\n",
    "for var in ['indep.x_dist', 'indep.x_nd', 'D1.out_dist']:\n",
    "    print(var, prob.get_val(var, get_remote=True))\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "assert_near_equal(prob.get_val(var, get_remote=True), np.array([24.53604616, 29.53604616, 36.53604616, 45.53604616, 56.53604616, 69.53604616, 84.53604616]), 1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: Distributed I/O and a Non-Distributed Output\n",
    "\n",
    "You can also create a component with a non-distributed output and distributed outputs and inputs. This situation tends to be more tricky and usually requires you to perform some MPI operations in your component's `compute` method. If the non-distributed output is only a function of the non-distributed inputs, then you can handle that variable just like you do on any other component. However, this example extends the previous component to include a non-distributed output that is a function of both the non-distributed and distributed inputs. In this case, it's a function of the sum of the square root of each element in the full distributed vector. Since the data is not all on any local processor, we use an MPI operation, in this case `Allreduce`, to make a summation across the distributed vector, and gather the answer back to each processor. The MPI operation and your implementation will vary, but consider this to be a general example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-output",
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "\n",
    "class MixedDistrib2(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Distributed Input\n",
    "        self.add_input('in_dist', shape_by_conn=True, distributed=True)\n",
    "\n",
    "        # Non-Distributed Input\n",
    "        self.add_input('in_nd', shape_by_conn=True)\n",
    "\n",
    "        # Distributed Output\n",
    "        self.add_output('out_dist', copy_shape='in_dist', distributed=True)\n",
    "\n",
    "        # Non-Distributed Output\n",
    "        self.add_output('out_nd', copy_shape='in_nd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        Id = inputs['in_dist']\n",
    "        Ind = inputs['in_nd']\n",
    "\n",
    "        # \"Computationally Intensive\" operation that we wish to parallelize.\n",
    "        f_Id = Id**2 - 2.0*Id + 4.0\n",
    "\n",
    "        # These operations are repeated on all procs.\n",
    "        f_Ind = Ind ** 0.5\n",
    "        g_Ind = Ind**2 + 3.0*Ind - 5.0\n",
    "\n",
    "        # Compute square root of our portion of the distributed input.\n",
    "        g_Id = Id ** 0.5\n",
    "\n",
    "        # Distributed output\n",
    "        outputs['out_dist'] = f_Id + np.sum(f_Ind)\n",
    "\n",
    "        # Non-Distributed output\n",
    "        if MPI and comm.size > 1:\n",
    "\n",
    "            # We need to gather the summed values to compute the total sum over all procs.\n",
    "            local_sum = np.array(np.sum(g_Id))\n",
    "            total_sum = local_sum.copy()\n",
    "            self.comm.Allreduce(local_sum, total_sum, op=MPI.SUM)\n",
    "\n",
    "            outputs['out_nd'] = g_Ind + total_sum\n",
    "        else:\n",
    "            # Recommended to make sure your code can run without MPI too, for testing.\n",
    "            outputs['out_nd'] = g_Ind + np.sum(g_Id)\n",
    "\n",
    "size = 7\n",
    "\n",
    "if MPI:\n",
    "    comm = MPI.COMM_WORLD\n",
    "    rank = comm.rank\n",
    "    sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "else:\n",
    "    # When running without MPI, the entire variable is on one proc.\n",
    "    rank = 0\n",
    "    sizes = {rank : size}\n",
    "    offsets = {rank : 0}\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "# Create a distributed source for the distributed input.\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "ivc.add_output('x_nd', np.zeros(size))\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", MixedDistrib2())\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "model.connect('indep.x_nd', 'D1.in_nd')\n",
    "\n",
    "prob.setup()\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "# Set initial values of non-distributed variable.\n",
    "x_nd_init = 1.0 + 2.0*np.arange(size)\n",
    "prob.set_val('indep.x_nd', x_nd_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "# Values on each rank.\n",
    "for var in ['indep.x_dist', 'indep.x_nd', 'D1.out_dist', 'D1.out_nd']:\n",
    "    print(var, prob.get_val(var))\n",
    "\n",
    "# Full gathered values.\n",
    "for var in ['indep.x_dist', 'indep.x_nd', 'D1.out_dist', 'D1.out_nd']:\n",
    "    print(var, prob.get_val(var, get_remote=True))\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "assert_near_equal(prob.get_val(var, get_remote=True), np.array([15.89178696, 29.89178696, 51.89178696, 81.89178696, 119.89178696, 165.89178696, 219.89178696]), 1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "In this example, we introduce a new component called an [IndepVarComp](indepvarcomp.ipynb). If you used OpenMDAO prior to version 3.2, then you are familiar with this component. It is used to define an independent variable.\n",
    "\n",
    "You usually do not have to define these because OpenMDAO defines and uses them automatically for all unconnected inputs in your model. This automatically-created `IndepVarComp` is called an Auto-IVC.\n",
    "\n",
    "However, when we define a distributed input, it sometimes isn’t possible to determine the full size of the corresponding independent variable, and the *IndepVarComp* cannot be created automatically. So, for unconnected inputs on a distributed component, you must manually create one, as we did in this example.\n",
    "```\n",
    "\n",
    "# Derivatives with Distributed Variables\n",
    "\n",
    "In the following examples, we show how to add analytic derivatives to the distributed examples given above. In most cases it is straighforward, but there is a special case when a non-distributed output depends on a distributed input or vice versa:\n",
    "\n",
    "> When you have a distributed output depending on a non-distributed input, or a non-distributed output depending on a distributed input, the [matrix-free](matrix-free-api) format is required, meaning that you will need to define a `compute_jacvec_product` method if your component is an `ExplicitComponent`, or `apply_linear` if your component is an `ImplicitComponent`.  The reasons for this are a bit subtle, but it boils down to OpenMDAO not knowing what kind of distributed operations are being done in the `compute` and having no way to manage those details when you propagate things in reverse.\n",
    "\n",
    "An important thing to keep in mind when deciding whether or not you need to some sort of reduce operation when computing derivatives using the matrix-free API with mixed distributed/non-distributed derivatives is that derivatives of any non-distributed variable in that case must be the same across all ranks.\n",
    "\n",
    "## Derivatives: Distributed I/O and a Non-Distributed Input\n",
    "\n",
    "In this example, we have a distributed input, a distributed output, and a non-distributed input. Because we have a distributed output, 'out_dist' depending on a non-distributed input 'in_nd', we have to use the matrix-free API and define a `compute_jacvec_product` method for our component. It is described in the feature document for [ExplicitComponent](explicit_component.ipynb). Note that for this component we only have to do an allreduce when computing the jacobian vector product in reverse mode. That's because in reverse mode we have to sum up the contribution of the entire distributed output 'out_dist' to the non-distributed input 'in_nd' so that the derivatives for 'in_nd' will be the same across all ranks. \n",
    "We verify that the derivatives are correct using `check_totals` with complex step since our component is complex-safe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "\n",
    "class MixedDistrib1(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Distributed Input\n",
    "        self.add_input('in_dist', shape_by_conn=True, distributed=True)\n",
    "\n",
    "        # Non-Distributed Input\n",
    "        self.add_input('in_nd', shape_by_conn=True)\n",
    "\n",
    "        # Distributed Output\n",
    "        self.add_output('out_dist', copy_shape='in_dist', distributed=True)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        Id = inputs['in_dist']\n",
    "        Ind = inputs['in_nd']\n",
    "\n",
    "        # \"Computationally Intensive\" operation that we wish to parallelize.\n",
    "        f_Id = Id**2 - 2.0*Id + 4.0\n",
    "\n",
    "        # This operation is repeated on all procs.\n",
    "        f_Ind = Ind ** 0.5\n",
    "\n",
    "        outputs['out_dist'] = f_Id + np.sum(f_Ind)\n",
    "\n",
    "    def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):\n",
    "        Id = inputs['in_dist']\n",
    "        Ind = inputs['in_nd']\n",
    "\n",
    "        df_dId = 2.0 * Id - 2.0\n",
    "        df_dInd = 0.5 / Ind ** 0.5\n",
    "\n",
    "        nId = len(Id)\n",
    "        nInd = len(Ind)\n",
    "\n",
    "        if mode == 'fwd':\n",
    "            if 'out_dist' in d_outputs:\n",
    "                if 'in_nd' in d_inputs:\n",
    "                    d_outputs['out_dist'] += np.tile(df_dInd, nId).reshape((nId, nInd)).dot(d_inputs['in_nd'])\n",
    "                if 'in_dist' in d_inputs:\n",
    "                    d_outputs['out_dist'] += df_dId * d_inputs['in_dist']\n",
    "\n",
    "        else:  # rev\n",
    "            if 'out_dist' in d_outputs:\n",
    "                if 'in_nd' in d_inputs:\n",
    "                    d_inputs['in_nd'] += self.comm.allreduce(np.tile(df_dInd, nId).reshape((nId, nInd)).T.dot(d_outputs['out_dist']))\n",
    "\n",
    "                if 'in_dist' in d_inputs:\n",
    "                    d_inputs['in_dist'] += df_dId * d_outputs['out_dist']\n",
    "\n",
    "\n",
    "size = 7\n",
    "\n",
    "if MPI:\n",
    "    comm = MPI.COMM_WORLD\n",
    "    rank = comm.rank\n",
    "    sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "else:\n",
    "    # When running without MPI, the entire variable is on one proc.\n",
    "    rank = 0\n",
    "    sizes = {rank : size}\n",
    "    offsets = {rank : 0}\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "# Create a distributed source for the distributed input.\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "ivc.add_output('x_nd', np.zeros(size))\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", MixedDistrib1())\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "model.connect('indep.x_nd', 'D1.in_nd')\n",
    "\n",
    "model.add_design_var('indep.x_nd')\n",
    "model.add_design_var('indep.x_dist')\n",
    "model.add_objective('D1.out_dist')\n",
    "\n",
    "prob.setup(force_alloc_complex=True)\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "# Set initial values of non-distributed variable.\n",
    "x_nd_init = 1.0 + 2.0*np.arange(size)\n",
    "prob.set_val('indep.x_nd', x_nd_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "if rank > 0:\n",
    "    prob.check_totals(method='cs', out_stream=None)\n",
    "else:\n",
    "    prob.check_totals(method='cs')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_check_totals\n",
    "assert_check_totals(prob.check_totals(method='cs', out_stream=None))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Derivatives: Mixed Distributed and Nondistributed I/O\n",
    "\n",
    "The following example shows how to implement derivatives on the earlier `MixedDistrib2` component, which is similar to `MixDistrib1` but adds a non-distributed output, 'out_nd' to the mix.  This component uses the `allreduce` operation to combine contributions to its non-distributed input `in_nd` from its distributed output `out_dist` when in reverse mode. In forward mode it uses `allreduce` to combine the contributions to its non-distributed output `out_nd` from its distributed input `in_dist`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "\n",
    "class MixedDistrib2(om.ExplicitComponent):\n",
    "\n",
    "    def setup(self):\n",
    "\n",
    "        # Distributed Input\n",
    "        self.add_input('in_dist', shape_by_conn=True, distributed=True)\n",
    "\n",
    "        # Non-Distributed Input\n",
    "        self.add_input('in_nd', shape_by_conn=True)\n",
    "\n",
    "        # Distributed Output\n",
    "        self.add_output('out_dist', copy_shape='in_dist', distributed=True)\n",
    "\n",
    "        # Non-Distributed Output\n",
    "        self.add_output('out_nd', copy_shape='in_nd')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        Id = inputs['in_dist']\n",
    "        Ind = inputs['in_nd']\n",
    "\n",
    "        # \"Computationally Intensive\" operation that we wish to parallelize.\n",
    "        f_Id = Id**2 - 2.0*Id + 4.0\n",
    "\n",
    "        # These operations are repeated on all procs.\n",
    "        f_Ind = Ind ** 0.5\n",
    "        g_Ind = Ind**2 + 3.0*Ind - 5.0\n",
    "\n",
    "        # Compute square root of our portion of the distributed input.\n",
    "        g_Id = Id ** 0.5\n",
    "\n",
    "        # Distributed output\n",
    "        outputs['out_dist'] = f_Id + np.sum(f_Ind)\n",
    "\n",
    "        # Non-Distributed output\n",
    "        if MPI and comm.size > 1:\n",
    "\n",
    "            # We need to gather the summed values to compute the total sum over all procs.\n",
    "            local_sum = np.array(np.sum(g_Id))\n",
    "            total_sum = local_sum.copy()\n",
    "            self.comm.Allreduce(local_sum, total_sum, op=MPI.SUM)\n",
    "            outputs['out_nd'] = g_Ind + total_sum\n",
    "        else:\n",
    "            # Recommended to make sure your code can run without MPI too, for testing.\n",
    "            outputs['out_nd'] = g_Ind + np.sum(g_Id)\n",
    "\n",
    "    def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):\n",
    "        Id = inputs['in_dist']\n",
    "        Ind = inputs['in_nd']\n",
    "\n",
    "        df_dId = 2.0 * Id - 2.0\n",
    "        df_dInd = 0.5 / Ind ** 0.5\n",
    "        dg_dId = 0.5 / Id ** 0.5\n",
    "        dg_dInd = 2.0 * Ind + 3.0\n",
    "\n",
    "        nId = len(Id)\n",
    "        nInd = len(Ind)\n",
    "\n",
    "        if mode == 'fwd':\n",
    "            if 'out_dist' in d_outputs:\n",
    "                if 'in_dist' in d_inputs:\n",
    "                    d_outputs['out_dist'] += df_dId * d_inputs['in_dist']\n",
    "                if 'in_nd' in d_inputs:\n",
    "                    d_outputs['out_dist'] += np.tile(df_dInd, nId).reshape((nId, nInd)).dot(d_inputs['in_nd'])\n",
    "            if 'out_nd' in d_outputs:\n",
    "                if 'in_dist' in d_inputs:\n",
    "                    d_outputs['out_nd'] += self.comm.allreduce(np.tile(dg_dId, nInd).reshape((nInd, nId)).dot(d_inputs['in_dist']))\n",
    "                if 'in_nd' in d_inputs:\n",
    "                    d_outputs['out_nd'] += dg_dInd * d_inputs['in_nd']\n",
    "\n",
    "        else:  # rev\n",
    "            if 'out_dist' in d_outputs:\n",
    "                if 'in_dist' in d_inputs:\n",
    "                    d_inputs['in_dist'] += df_dId * d_outputs['out_dist']\n",
    "                if 'in_nd' in d_inputs:\n",
    "                    d_inputs['in_nd'] += self.comm.allreduce(np.tile(df_dInd, nId).reshape((nId, nInd)).T.dot(d_outputs['out_dist']))\n",
    "\n",
    "            if 'out_nd' in d_outputs:\n",
    "                if 'in_dist' in d_inputs:\n",
    "                    d_inputs['in_dist'] += np.tile(dg_dId, nInd).reshape((nInd, nId)).T.dot(d_outputs['out_nd'])\n",
    "                if 'in_nd' in d_inputs:\n",
    "                    d_inputs['in_nd'] += dg_dInd * d_outputs['out_nd']\n",
    "\n",
    "\n",
    "size = 7\n",
    "\n",
    "if MPI:\n",
    "    comm = MPI.COMM_WORLD\n",
    "    rank = comm.rank\n",
    "    sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "else:\n",
    "    # When running without MPI, the entire variable is on one proc.\n",
    "    rank = 0\n",
    "    sizes = {rank : size}\n",
    "    offsets = {rank : 0}\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "# Create a distributed source for the distributed input.\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "ivc.add_output('x_nd', np.zeros(size))\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", MixedDistrib2())\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "model.connect('indep.x_nd', 'D1.in_nd')\n",
    "\n",
    "model.add_design_var('indep.x_nd')\n",
    "model.add_design_var('indep.x_dist')\n",
    "model.add_constraint('D1.out_dist', lower=0.0)\n",
    "model.add_constraint('D1.out_nd', lower=0.0)\n",
    "\n",
    "prob.setup(force_alloc_complex=True)\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "# Set initial values of non-distributed variable.\n",
    "x_nd_init = 1.0 + 2.0*np.arange(size)\n",
    "prob.set_val('indep.x_nd', x_nd_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "if rank > 0:\n",
    "    prob.check_totals(method='cs', out_stream=None)\n",
    "else:\n",
    "    prob.check_totals(method='cs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "# We need this complicated test to be sure there are no errors in the allreduces.\n",
    "from openmdao.utils.array_utils import evenly_distrib_idxs\n",
    "from openmdao.utils.assert_utils import assert_check_totals\n",
    "\n",
    "size = 5\n",
    "comm = MPI.COMM_WORLD\n",
    "rank = comm.rank\n",
    "sizes, offsets = evenly_distrib_idxs(comm.size, size)\n",
    "\n",
    "prob = om.Problem()\n",
    "model = prob.model\n",
    "\n",
    "ivc = om.IndepVarComp()\n",
    "ivc.add_output('x_dist', np.zeros(sizes[rank]), distributed=True)\n",
    "ivc.add_output('x_nd', np.zeros(size))\n",
    "\n",
    "model.add_subsystem(\"indep\", ivc)\n",
    "model.add_subsystem(\"D1\", MixedDistrib2())\n",
    "model.add_subsystem(\"D2\", MixedDistrib2())\n",
    "model.add_subsystem(\"D3\", MixedDistrib2())\n",
    "model.add_subsystem(\"D4\", MixedDistrib2())\n",
    "\n",
    "model.connect('indep.x_dist', 'D1.in_dist')\n",
    "model.connect('indep.x_nd', 'D1.in_nd')\n",
    "model.connect('D1.out_dist', 'D2.in_dist')\n",
    "model.connect('D1.out_nd', 'D2.in_nd')\n",
    "model.connect('D2.out_dist', 'D3.in_dist')\n",
    "model.connect('D2.out_nd', 'D3.in_nd')\n",
    "model.connect('D3.out_dist', 'D4.in_dist')\n",
    "model.connect('D3.out_nd', 'D4.in_nd')\n",
    "\n",
    "model.add_design_var('indep.x_nd')\n",
    "model.add_design_var('indep.x_dist')\n",
    "model.add_constraint('D4.out_dist', lower=0.0)\n",
    "model.add_constraint('D4.out_nd', lower=0.0)\n",
    "\n",
    "prob.setup(force_alloc_complex=True, mode='rev')\n",
    "\n",
    "# Set initial values of distributed variable.\n",
    "x_dist_init = 3.0 + np.arange(size)[offsets[rank]:offsets[rank] + sizes[rank]]\n",
    "prob.set_val('indep.x_dist', x_dist_init)\n",
    "\n",
    "# Set initial values of non-distributed variable.\n",
    "x_nd_init = 1.0 + 2.0*np.arange(size)\n",
    "prob.set_val('indep.x_nd', x_nd_init)\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "assert_check_totals(prob.check_totals(method='cs', out_stream=None))\n"
   ]
  }
 ],
 "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.3"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
