{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94d7a6cc",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "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",
   "id": "45f5bd8c",
   "metadata": {},
   "source": [
    "# Parallel Coloring for Multipoint or Fan-Out Problems\n",
    "\n",
    "In many models, there is an opportunity to parallelize across multiple points (e.g. multiple load cases for a structural optimization, multiple flight conditions for an aerodynamic optimization). Executing the nonlinear solve for this model in parallel offers a large potential speed-up, but when computing total derivatives, achieving that same parallel performance may require the use of OpenMDAO’s parallel coloring algorithm.\n",
    "\n",
    "```{Note}\n",
    "Parallel coloring is appropriate when you have some inexpensive serial data path in your model, before the parallel points. For more details on when a model calls for parallel coloring see the [Theory Manual entry on the fan-out model structures](../../../theory_manual/advanced_linear_solvers_special_cases/fan_out).\n",
    "```\n",
    "\n",
    "Parallel coloring is specified via the `parallel_deriv_color` argument to the [add_constraint()](../adding_desvars_cons_objs/adding_constraint) method. The color specified can be any hashable object (e.g. string, int, tuple). Two constraints, pointing to variables from different components on different processors, given the same `parallel_deriv_color` argument will be solved for in parallel with each other."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "415a3416",
   "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",
   "id": "945d0dee",
   "metadata": {},
   "source": [
    "## Usage Example\n",
    "\n",
    "Here is a toy problem that runs on two processors showing how to use this feature\n",
    "\n",
    "### Class definitions for a simple problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "860b0d19",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class SumComp(om.ExplicitComponent):\n",
    "    def __init__(self, size):\n",
    "        super().__init__()\n",
    "        self.size = size\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=np.zeros(self.size))\n",
    "        self.add_output('y', val=0.0)\n",
    "\n",
    "        self.declare_partials(of='*', wrt='*')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = np.sum(inputs['x'])\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        partials['y', 'x'] = np.ones(inputs['x'].size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3923e6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "class SlowComp(om.ExplicitComponent):\n",
    "    \"\"\"\n",
    "    Component with a delay that multiplies the input by a multiplier.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, delay=1.0, size=3, mult=2.0):\n",
    "        super().__init__()\n",
    "        self.delay = delay\n",
    "        self.size = size\n",
    "        self.mult = mult\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=0.0)\n",
    "        self.add_output('y', val=np.zeros(self.size))\n",
    "\n",
    "        self.declare_partials(of='*', wrt='*')\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = inputs['x'] * self.mult\n",
    "\n",
    "    def compute_partials(self, inputs, partials):\n",
    "        partials['y', 'x'] = self.mult\n",
    "\n",
    "    def _apply_linear(self, jac, vec_names, rel_systems, mode, scope_out=None, scope_in=None):\n",
    "        time.sleep(self.delay)\n",
    "        super()._apply_linear(jac, vec_names, rel_systems, mode, scope_out, scope_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9cb1619",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PartialDependGroup(om.Group):\n",
    "    def setup(self):\n",
    "        size = 4\n",
    "\n",
    "        self.add_subsystem('Comp1', SumComp(size))\n",
    "        pargroup = self.add_subsystem('ParallelGroup1', om.ParallelGroup())\n",
    "\n",
    "        self.set_input_defaults('Comp1.x', val=np.arange(size, dtype=float)+1.0)\n",
    "\n",
    "        self.linear_solver = om.LinearBlockGS()\n",
    "        self.linear_solver.options['iprint'] = -1\n",
    "        pargroup.linear_solver = om.LinearBlockGS()\n",
    "        pargroup.linear_solver.options['iprint'] = -1\n",
    "\n",
    "        delay = .1\n",
    "        pargroup.add_subsystem('Con1', SlowComp(delay=delay, size=2, mult=2.0))\n",
    "        pargroup.add_subsystem('Con2', SlowComp(delay=delay, size=2, mult=-3.0))\n",
    "\n",
    "        self.connect('Comp1.y', 'ParallelGroup1.Con1.x')\n",
    "        self.connect('Comp1.y', 'ParallelGroup1.Con2.x')\n",
    "\n",
    "        color = 'parcon'\n",
    "        self.add_design_var('Comp1.x')\n",
    "        self.add_constraint('ParallelGroup1.Con1.y', lower=0.0, parallel_deriv_color=color)\n",
    "        self.add_constraint('ParallelGroup1.Con2.y', upper=0.0, parallel_deriv_color=color)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7ac18fde",
   "metadata": {},
   "source": [
    "### Run script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6f84fbf",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.notebook_utils import get_code\n",
    "from myst_nb import glue\n",
    "glue(\"code_src65\", get_code(\"openmdao.core.tests.test_parallel_derivatives.PartialDependGroup\"), display=False)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b8d6ae8a",
   "metadata": {},
   "source": [
    ":::{Admonition} `PartialDependGroup` class definition \n",
    ":class: dropdown\n",
    "\n",
    "{glue:}`code_src65`\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae0530ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.core.tests.test_parallel_derivatives import PartialDependGroup\n",
    "\n",
    "size = 4\n",
    "\n",
    "of = ['ParallelGroup1.Con1.y', 'ParallelGroup1.Con2.y']\n",
    "wrt = ['Comp1.x']\n",
    "\n",
    "p = om.Problem(model=PartialDependGroup())\n",
    "p.setup(mode='rev')\n",
    "p.run_model()\n",
    "\n",
    "J = p.compute_totals(of, wrt, return_format='dict')\n",
    "\n",
    "print(J['ParallelGroup1.Con1.y']['Comp1.x'][0])\n",
    "print(J['ParallelGroup1.Con2.y']['Comp1.x'][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89a0ad26",
   "metadata": {
    "scrolled": true,
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "import numpy as np\n",
    "\n",
    "expected = np.array([[2., 2., 2., 2.],[-3., -3., -3., -3.]])\n",
    "assert_near_equal(J['ParallelGroup1.Con1.y']['Comp1.x'][0], expected[0])\n",
    "assert_near_equal(J['ParallelGroup1.Con2.y']['Comp1.x'][0], expected[1])"
   ]
  }
 ],
 "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": 5
}
