{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a551df3",
   "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": "b26341c6",
   "metadata": {},
   "source": [
    "# Speeding up Derivative Approximations with Parallel Finite Difference and Parallel Complex Step\n",
    "\n",
    "If you have multiple processors available to you, it’s possible to speed up the calculation of approximated partials or totals by computing multiple columns of the approximated Jacobian matrix simultaneously across multiple processes. Setting up *parallel* finite difference or *parallel* complex step is identical to setting up for serial execution, except for the addition of a single `__init__` arg that sets `num_par_fd` on the System(s) where you want to compute the approximate derivatives. The `num_par_fd` arg specifies the number of approximated Jacobian columns that will be computed in parallel, assuming that enough processes are provided when the problem script is executed.\n",
    "\n",
    "As a simple example, lets use parallel finite difference to compute the partial and total derivatives using a simple model containing a component that computes its outputs by multiplying its inputs by a constant matrix. The partial jacobian for that component will be identical to that constant matrix. The component we’ll use is shown below:"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7e36763f",
   "metadata": {},
   "source": [
    "```{note}\n",
    "This feature requires MPI, and may not be able to be run on Colab or Binder.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b078fd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import time\n",
    "import openmdao.api as om\n",
    "\n",
    "\n",
    "class MatMultComp(om.ExplicitComponent):\n",
    "    def __init__(self, mat, approx_method='exact', sleep_time=0.1, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.mat = mat\n",
    "        self.approx_method = approx_method\n",
    "        self.sleep_time = sleep_time\n",
    "\n",
    "    def setup(self):\n",
    "        self.add_input('x', val=np.ones(self.mat.shape[1]))\n",
    "        self.add_output('y', val=np.zeros(self.mat.shape[0]))\n",
    "        self.num_computes = 0\n",
    "\n",
    "    def setup_partials(self):\n",
    "        self.declare_partials('*', '*', method=self.approx_method)\n",
    "\n",
    "    def compute(self, inputs, outputs):\n",
    "        outputs['y'] = self.mat.dot(inputs['x'])\n",
    "        self.num_computes += 1\n",
    "        time.sleep(self.sleep_time)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1cfe11be",
   "metadata": {},
   "source": [
    "Our model will also contain an IndepVarComp that we’ll connect to our MatMultComp, and we’ll compute total derivatives of our MatMultComp outputs with respect to our IndepVarComp output.\n",
    "\n",
    "More details about setting up for partial derivative approximation can be found [here](approximating_partial_derivatives) and an explanation of total derivative approximation can be found [here](approximating_totals)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e8d96de8",
   "metadata": {},
   "source": [
    "## Allocating processes\n",
    "\n",
    "In both cases we’ll be running our Python script under MPI using 3 processes. We need 3 processes because our model requires 1 process and we’ll be using a `num_par_fd` value of 3, so the number of processes we need is 3 * 1 = 3. In general, When we set `num_par_fd`, it acts as a multiplier on the number of processes needed for any given system when running under MPI. For example, if a given system requires N processes, then that same system when using parallel finite difference will require N * `num_par_fd` processes. This is because we’re duplicating the given system `num_par_fd` times and using those duplicate systems to solve for different columns of our approximated Jacobian in parallel.\n",
    "\n",
    "As a reminder, to run our script and give it 3 processes, we would do the following:\n",
    "``` bash\n",
    "    mpirun -n 3 python my_problem.py\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "68c4dc1e",
   "metadata": {},
   "source": [
    "## Examples\n",
    "\n",
    "First, let’s compute the partial derivatives across our MatMultComp. We’ll use a matrix with 6 columns and a `num_par_fd` value on our MatMultComp component of 4, meaning that each of our 4 processes should compute 1 or 2 columns of the partial Jacobian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f804018",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "mat = np.arange(30, dtype=float).reshape(5, 6)\n",
    "\n",
    "p = om.Problem()\n",
    "model = p.model\n",
    "comp = model.add_subsystem('comp', MatMultComp(mat, approx_method='fd', num_par_fd=4))\n",
    "\n",
    "model.set_input_defaults('comp.x', val=np.ones(mat.shape[1]))\n",
    "p.setup(mode='fwd')\n",
    "p.run_model()\n",
    "\n",
    "pre_count = comp.num_computes\n",
    "\n",
    "# calling compute_totals will result in the computation of partials for comp\n",
    "p.compute_totals(of=['comp.y'], wrt=['comp.x'])\n",
    "\n",
    "# get the partial jacobian matrix\n",
    "J = comp._jacobian['y', 'x']\n",
    "\n",
    "post_count =  comp.num_computes\n",
    "\n",
    "# how many computes were used in this proc to compute the partial jacobian?\n",
    "# Each proc should be doing 2 computes.\n",
    "jac_count = post_count - pre_count\n",
    "\n",
    "print(jac_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd1ff6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "# J and mat should be the same\n",
    "print(np.linalg.norm(J - mat))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b00df6",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "assert(jac_count < 3)\n",
    "assert_near_equal(np.linalg.norm(J - mat), 0.0, tolerance=1e-7)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "dd2ac6c7",
   "metadata": {},
   "source": [
    "Next, let’s compute the total derivatives of our MatMulComp outputs with respect to its input. Again, we’ll be using a *num_par_fd* value of 2 and a matrix having 6 columns, so each process should compute either 1 or 2 columns of the total Jacobian. This time, however, we set the *num_par_fd* on our model instead of on our MatMultComp."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ac4e176",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "mat = np.arange(30, dtype=float).reshape(5, 6)\n",
    "\n",
    "p = om.Problem(model=om.Group(num_par_fd=4))\n",
    "model = p.model\n",
    "model.approx_totals(method='fd')\n",
    "comp = model.add_subsystem('comp', MatMultComp(mat))\n",
    "\n",
    "model.set_input_defaults('comp.x', val=np.ones(mat.shape[1]))\n",
    "p.setup(mode='fwd')\n",
    "p.run_model()\n",
    "\n",
    "pre_count = comp.num_computes\n",
    "\n",
    "J = p.compute_totals(of=['comp.y'], wrt=['comp.x'], return_format='array')\n",
    "\n",
    "post_count =  comp.num_computes\n",
    "\n",
    "# how many computes were used in this proc to compute the total jacobian?\n",
    "# Each proc should be doing 2 computes.\n",
    "jac_count = post_count - pre_count\n",
    "\n",
    "print(jac_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a17f8c91",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "# J and mat should be the same\n",
    "print(np.linalg.norm(J - mat))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bae08bef",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "assert(jac_count < 3)\n",
    "assert_near_equal(np.linalg.norm(J - mat), 0.0, tolerance=1e-7)"
   ]
  }
 ],
 "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.5"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
