{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21150be1",
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "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": "1d49901a",
   "metadata": {},
   "source": [
    "# Simultaneous Total Derivative Coloring For Separable Problems\n",
    "\n",
    "When OpenMDAO solves for total derivatives, it loops over either design variables in ‘fwd’ mode or responses in ‘rev’ mode. For each of those variables, it performs a linear solve for each member of that variable, so for a scalar variable there would be only a single linear solve, and there would be N solves for an array variable of size N.\n",
    "\n",
    "Certain problems have a special kind of sparsity structure in the total derivative Jacobian that allows OpenMDAO to solve for multiple derivatives simultaneously. This can result in far fewer linear solves and much-improved performance. These problems are said to have separable variables. The concept of separability is explained in the [Theory Manual](../../../theory_manual/advanced_linear_solvers_special_cases/separable).\n",
    "\n",
    "Simultaneous derivative coloring in OpenMDAO can be performed either statically or dynamically.\n",
    "\n",
    "When mode is set to ‘fwd’ or ‘rev’, a unidirectional coloring algorithm is used to group columns or rows, respectively, for simultaneous derivative calculation. The algorithm used in this case is the greedy algorithm with ordering by incidence degree found in T. F. Coleman and J. J. More, *Estimation of sparse Jacobian matrices and graph coloring problems*, SIAM J. Numer. Anal., 20 (1983), pp. 187–209.\n",
    "\n",
    "When using simultaneous derivatives, setting `mode=’auto’` in `Problem.setup` will indicate that bidirectional coloring should be used. Bidirectional coloring can significantly decrease the number of linear solves needed to generate the total Jacobian relative to coloring only in fwd or rev mode.\n",
    "\n",
    "For more information on the bidirectional coloring algorithm, see T. F. Coleman and A. Verma, *The efficient computation of sparse Jacobian matrices using automatic differentiation*, SIAM J. Sci. Comput., 19 (1998), pp. 1210–1233."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58a03898",
   "metadata": {},
   "source": [
    "```{Note}\n",
    "Bidirectional coloring is a new feature and should be considered experimental at this point.\n",
    "\n",
    "``` "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ace1838a",
   "metadata": {},
   "source": [
    "The OpenMDAO algorithms use the sparsity patterns for the partial derivatives given by the [declare_partials](sparse_partials) calls from all of the components in your model. So you should [specify the sparsity of the partial derivatives](sparse_partials) of your components in order to make it possible to find a better automatic coloring.\n",
    "\n",
    "## Dynamic Coloring\n",
    "\n",
    "Dynamic coloring computes the derivative colors at runtime, shortly after the driver begins the optimization. This has the advantage of simplicity and robustness to changes in the model, but adds the cost of the coloring computation to the run time of the optimization. Generally, however, this cost will be small relative to the cost of the full optimization unless your total jacobian is very large.\n",
    "\n",
    "Activating dynamic coloring is simple. Just call the `declare_coloring` function on the driver. For example:\n",
    "\n",
    "    prob.driver.declare_coloring()\n",
    "\n",
    "If you want to change the number of `compute_totals` calls that the coloring algorithm uses to compute the Jacobian sparsity (default is 3), the tolerance used to determine nonzeros (default is 1e-25), you can pass the `num_full_jacs`, and `tol` args. You can also pass the `min_improve_pct` arg, which specifies how much the coloring must reduce the number of linear solves required to generate the total Jacobian else coloring will be deactivated.\n",
    "\n",
    "For example:\n",
    "\n",
    "    prob.driver.declare_coloring(num_full_jacs=2, tol=1e-20, min_improve_pct=10.)\n",
    "\n",
    "If you want to perform a tolerance sweep, trying out a range of tolerances when determining what values of the sparsity matrix should be considered zero, then you can set the `orders` arg. OpenMDAO will then sweep over tolerances from the given `tol` plus and minus `orders` orders of magnitude. For each tolerance, it checks the number of nonzero values, and chooses a tolerance from the largest group of tolerances that share the same number of nonzeros. If it can’t find at least two tolerances that result in the same number of nonzeros, an exception is raised.\n",
    "\n",
    "Here’s an example:\n",
    "\n",
    "    prob.driver.declare_coloring(tol=1e-35, orders=20)\n",
    "\n",
    "Finally, when using bidirectional coloring, one can use the `direct` method or the `substitution` method to construct the column/row adjacency matrix during the coloring of the forward/reverse partitions of the jacobian.  The `direct` method is the default.  The `substitution` method sometimes results in a better coloring, i.e., a coloring with fewer linear solves, but it does this by allowing some jacobian values in the forward and reverse partitions to overlap.  This overlapping is corrected for in the full jacobian by subtracting some jacobian values from others, which can introduce round-off error.\n",
    "\n",
    "To cause bidirectional coloring to use the `substitution` method, call `declare_coloring` with `direct=False`, for example:\n",
    "\n",
    "    prob.driver.declare_coloring(direct=False)\n",
    "\n",
    "Because of the possibility of round-off errors for `substitution` method, the best advice is to start off using the default `direct` method, and once everything is working, try `substitution` method to see if it improves performance without causing convergence problems.  There are detailed descriptions of the `direct` and `substitution` methods in the Coleman and Verma paper mentioned above.\n",
    "    \n",
    "Whenever a dynamic coloring is computed, the coloring is written to a file called *total_coloring.pkl* for later ‘static’ use. The file will be written in the `{prob_name}_out/coloring_files` directory.\n",
    "\n",
    "You can see a more complete example of setting up an optimization with simultaneous derivatives in the [Simple Optimization using Simultaneous Derivatives](../../../examples/simul_deriv_example) example."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83a2cd72",
   "metadata": {},
   "source": [
    "(total-coloring-report)=\n",
    "## Total coloring as a report\n",
    "\n",
    "As part of the dynamic coloring process, OpenMDAO will automatically generate a visualization of the coloring that can be found in `reports/{problem_name}/total_coloring.html`. If the `bokeh` python plotting package is available, this report will be an interactive visualization, otherwise an ascii representation of the total coloring will be provided. More information about these visualizations is given below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76e6ec9c",
   "metadata": {},
   "source": [
    "(static-coloring)=\n",
    "## Static Coloring\n",
    "\n",
    "To get rid of the runtime cost of computing the coloring, you can precompute it and tell the driver to use its precomputed coloring by calling the `use_fixed_coloring` method on the driver. Note that this call should be made after calling `declare_coloring`.\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.core.driver.Driver.use_fixed_coloring\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "You don’t need to tell `use_fixed_coloring` the name of the coloring file to use, because it uses a fixed name, *total_coloring.pkl*, and knows what directory to look in based on the directory specified in `problem.options['coloring_dir']`. However, you can pass the name of a coloring file to `use_fixed_coloring` if you want to use a specific coloring file that doesn’t follow the standard naming convention.\n",
    "\n",
    "While using a precomputed coloring has the advantage of removing the runtime cost of computing the coloring, it should be used with care, because any changes in the model, design variables, or responses can make the existing coloring invalid. If *any* configuration changes have been made to the optimization, it’s recommended to regenerate the coloring before re-running the optimization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "657db0f1",
   "metadata": {},
   "source": [
    "(total_coloring-command)=\n",
    "## Updating total coloring from the command line\n",
    "\n",
    "The total coloring can be regenerated and written to the *total_coloring.pkl* file in a directory determined by the value of `problem.options['coloring_files']` using the following command:\n",
    "\n",
    "    openmdao total_coloring <your_script_name>\n",
    "\n",
    "Note that if you have passed a coloring filename into `use_fixed_coloring` instead of letting the framework determine the filename automatically, the framework will not update the contents of the coloring file even if you run `openmdao total_coloring` from the command line.\n",
    "\n",
    "The `total_coloring` command also generates summary information that can sometimes be useful. The tolerance that was actually used to determine if an entry in the total Jacobian is considered to be non-zero is displayed, along with the number of zero entries found in this case, and how many times that number of zero entries occurred when sweeping over different tolerances between +- a number of orders of magnitude around the given tolerance. If no tolerance is given, the default is 1e-15. If the number of occurrences is only 1, an exception will be raised, and you should increase the number of total derivative computations that the algorithm uses to compute the sparsity pattern. You can do that with the *-n* option. The following, for example, will perform the total derivative computation 5 times.\n",
    "\n",
    "    openmdao total_coloring <your_script_name> -n 5\n",
    "\n",
    "Note that when multiple total jacobian computations are performed, we take the absolute values of each jacobian and add them together, then divide by the maximum value, resulting in values between 0 and 1 for each entry.\n",
    "\n",
    "If repeating the total derivative computation multiple times doesn’t work, try changing the tolerance using the *-t* option as follows:\n",
    "\n",
    "    openmdao total_coloring <your_script_name> -n 5 -t 1e-10\n",
    "\n",
    "Be careful when setting the tolerance, however, because if you make it too large then you may be zeroing out Jacobian entries that should not be ignored and your optimization may not converge.\n",
    "\n",
    "If you want to examine the sparsity structure of your total Jacobian, you can use the *–view* option as follows:\n",
    "\n",
    "    openmdao total_coloring <your_script_name> --view\n",
    "\n",
    "which will display a visualization of the sparsity structure with rows and columns labelled with the response and design variable names, respectively. See the figure here.\n",
    "\n",
    "A text-based view is also available using the *–textview* arg. For example:\n",
    "\n",
    "    openmdao total_coloring <your_script_name> --textview\n",
    "\n",
    "will display something like the following:\n",
    "\n",
    "    ....................f 0  circle.area\n",
    "    f.........f.........f 1  r_con.g\n",
    "    .f.........f........f 2  r_con.g\n",
    "    ..f.........f.......f 3  r_con.g\n",
    "    ...f.........f......f 4  r_con.g\n",
    "    ....f.........f.....f 5  r_con.g\n",
    "    .....f.........f....f 6  r_con.g\n",
    "    ......f.........f...f 7  r_con.g\n",
    "    .......f.........f..f 8  r_con.g\n",
    "    ........f.........f.f 9  r_con.g\n",
    "    .........f.........ff 10  r_con.g\n",
    "    f.........f.......... 11  theta_con.g\n",
    "    ..f.........f........ 12  theta_con.g\n",
    "    ....f.........f...... 13  theta_con.g\n",
    "    ......f.........f.... 14  theta_con.g\n",
    "    ........f.........f.. 15  theta_con.g\n",
    "    ff........ff......... 16  delta_theta_con.g\n",
    "    ..ff........ff....... 17  delta_theta_con.g\n",
    "    ....ff........ff..... 18  delta_theta_con.g\n",
    "    ......ff........ff... 19  delta_theta_con.g\n",
    "    ........ff........ff. 20  delta_theta_con.g\n",
    "    f.................... 21  l_conx.g\n",
    "    |indeps.x\n",
    "              |indeps.y\n",
    "                        |indeps.r\n",
    "\n",
    "Note that the design variables are displayed along the bottom of the matrix, with a pipe symbol (|) that lines up with the starting column for that variable. Also, an ‘f’ indicates a nonzero value that is colored in ‘fwd’ mode, while an ‘r’ indicates a nonzero value colored in ‘rev’ mode. A ‘.’ indicates a zero value.\n",
    "\n",
    "The coloring file will be written in pickle format to the standard location and will be loaded using the `use_fixed_coloring` function like this:\n",
    "\n",
    "    prob.driver.use_fixed_coloring()\n",
    "\n",
    "Note that there are two ways to generate files that can be loaded using `use_fixed_coloring`. You can either run the `openmdao total_coloring` command line tool, or you can just run your model, and as long as you’ve called `declare_coloring` on your driver, it will automatically generate a coloring file that you can ‘lock in’ at some later point by adding a call to `use_fixed_coloring`, after you’re done making changes to your model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b07c8e38",
   "metadata": {},
   "source": [
    "(view_coloring-command)=\n",
    "## Viewing total coloring from the command line\n",
    "\n",
    "If you have a coloring file that was generated earlier and you want to view its statistics, you can use the `openmdao view_coloring` command to generate a small report.\n",
    "\n",
    "    openmdao view_coloring <your_coloring_file> -m\n",
    "\n",
    "will show metadata associated with the creation of the coloring along with a short summary. For example:\n",
    "\n",
    "    Coloring metadata:\n",
    "    {'orders': 20, 'num_full_jacs': 3, 'tol': 1e-15}\n",
    "\n",
    "    Jacobian shape: (22, 21)  (13.42% nonzero)\n",
    "\n",
    "    FWD solves: 5   REV solves: 0\n",
    "\n",
    "    Total colors vs. total size: 5 vs 21  (76.2% improvement)\n",
    "\n",
    "    Time to compute sparsity: 0.024192 sec.\n",
    "    Time to compute coloring: 0.001076 sec.\n",
    "\n",
    "Adding a *–view* arg will pop up an interactive plot showing the coloring of the jacobian. Forward colors will be colored shades of blue and reverse colors will be colored in shades of red. Hovering over a particular cell of the Jacobian will display the location, the color number, the coloring direction, and the particular ‘of’ and ‘wrt’ variables for that particular sub-jacobian. Note that this viewer requires the installation of bokeh. See the figure below.\n",
    "\n",
    "If bokeh is not available, then the resulting html viewer will contain an ascii version of the sparsity pattern similar to that in the previous section, along with some statistics about the coloring.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb6402ac",
   "metadata": {},
   "source": [
    "![Coloring viewer example](coloring_viewer.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5ba339",
   "metadata": {},
   "source": [
    "```{Note}\n",
    "Your coloring file(s) will be found in the standard directory `problem.options[‘coloring_dir’]`. That directory may contain a total coloring file, *total_coloring.pkl*, in additon to files containing partial derivative colorings for particular component classes or instances, as well as semi-total derivative coloring files for particular groups.\n",
    "``` "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b6551cd",
   "metadata": {},
   "source": [
    "If you run `openmdao total_coloring` and it turns out there is no simultaneous total coloring available, or that you don’t gain very much by coloring, don’t be surprised. Not all total Jacobians are sparse enough to benefit significantly from simultaneous derivatives."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "954218f7",
   "metadata": {},
   "source": [
    "(display_coloring-api)=\n",
    "## Viewing total coloring from a script\n",
    "\n",
    "The `openmdao.api` includes a `display_coloring` function to allow viewing of coloring data from within a script.\n",
    "This function will generate an html file, but will render the coloring as plain text if bokeh is unavailable or if the `as_text` argument is True.\n",
    "\n",
    "```{eval-rst}\n",
    "    .. autofunction:: openmdao.utils.coloring.display_coloring\n",
    "        :noindex:\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "924ffcb6",
   "metadata": {},
   "source": [
    "## Checking that it works\n",
    "\n",
    "After activating simultaneous derivatives, you should check your total derivatives using the [check_totals](check_total_derivatives) function. The algorithm that we use still has a small chance of computing an incorrect coloring due to the possibility that the total Jacobian being analyzed by the algorithm contained one or more zero values that are only incidentally zero. Using `check_totals` is the way to be sure that something hasn’t gone wrong.\n",
    "\n",
    "If you used the automatic coloring algorithm, and you find that `check_totals` is reporting incorrect total derivatives, then you should try using the *-n* and *-t* options mentioned earlier until you get the correct total derivatives."
   ]
  }
 ],
 "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": 5
}
