{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "described-buyer", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:37.287208Z", "iopub.status.busy": "2026-10-02T14:44:37.286935Z", "iopub.status.idle": "2026-10-02T14:44:37.292851Z", "shell.execute_reply": "2026-10-02T14:44:37.292244Z" }, "papermill": { "duration": 0.010086, "end_time": "2026-10-02T14:44:37.294121+00:00", "exception": false, "start_time": "2026-10-02T14:44:37.284035+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": "sensitive-mining", "metadata": { "papermill": { "duration": 0.001414, "end_time": "2026-10-02T14:44:37.297350+00:00", "exception": false, "start_time": "2026-10-02T14:44:37.295936+00:00", "status": "completed" }, "tags": [] }, "source": [ "# Simultaneous Coloring of Approximated Derivatives\n", "\n", "In OpenMDAO, partial derivatives for components can be approximated using either finite difference or complex step. Sometimes the partial or jacobians in these cases are sparse, and the computing of these jacobians can be made more efficient using simultaneous derivative coloring. For an explanation of a similar coloring for *total* derivatives, see [Simultaneous Coloring For Separable Problems](../core_features/working_with_derivatives/simul_derivs.ipynb). Finite difference and complex step only work in forward mode, so only a forward mode coloring is possible when using them, but depending on the sparsity pattern of the jacobian, it may still be possible to get significant efficiency gains.\n", "\n", "\n", "## Dynamic Coloring\n", "\n", "Setting up a problem to use dynamic coloring of approximated derivatives requires a call to the `declare_coloring` function.\n", "\n", "```{eval-rst}\n", ".. automethod:: openmdao.core.system.System.declare_coloring\n", " :noindex:\n", "```\n", "\n", "For example, the code below sets up coloring for partial derivatives of outputs of `comp` with respect to inputs of `comp` starting with 'x'. Let's assume here that `MyComp` is an `ExplicitComponent`. If it were an `ImplicitComponent`, then the wildcard pattern 'x*' would be applied to all inputs *and* outputs (states) of `comp`.\n", "\n", "```python\n", " comp = prob.model.add_subsystem('comp', MyComp())\n", " comp.declare_coloring('x*', method='cs', num_full_jacs=2, min_improve_pct=10.)\n", "```\n", "\n", "\n", "Note that in the call to `declare_coloring`, we also set `num_full_jacs` to 2. This means\n", "that the first 2 times that a partial jacobian is computed for 'comp', it's nonzero values will be computed\n", "without coloring and stored. Just prior to the 3rd time, the jacobian's sparsity pattern will be computed, which then allows the coloring to be computed and used for the rest of the run. We also set `min_improve_pct` to 10, meaning that if the computed coloring does not reduce the number of nonlinear solves required to compute `comp's` partial jacobian by 10 percent, then `comp` will not use coloring at all.\n", "\n", "The purpose of `declare_coloring` is to provide all of the necessary information to allow\n", "OpenMDAO to generate a coloring, either dynamically or manually using `openmdao partial_coloring`.\n", "\n", "Coloring files that are generated dynamically will be placed in the directory specified in `problem.options['coloring_dir']` and will be named based on the value of the `per_instance` arg passed to `declare_coloring`. If `per_instance` is True, the file will be named based on the full pathname of the component being colored. If False, the name will be based on the full module pathname of the class of the given\n", "component.\n", "\n", "`declare_coloring` should generally be called in the `setup` function of the component.\n", "\n", "Note that computing a partial jacobian when the jacobian is very large can be quite expensive, even if the jacobian is sparse, because a solution must be computed for every column of the jacobian, so you should set `num_full_jacs` only as high as is necessary to ensure that non-constant computed zeros in the jacobian are unlikely. OpenMDAO injects random noise into the inputs when solving for the columns of the jacobian, which should make non-constant computed zeros fairly unlikely even for `num_full_jacs=1`.\n", "\n", "\n", "Here's a modified version of our total coloring example, where we replace one of our components with one of type DynamicPartialsComp that computes a dynamic partial coloring. A total coloring is also performed, as in the previous example, but this time the total coloring uses sparsity information computed by our component during its dynamic partial coloring.\n", "\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "dedicated-vatican", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:37.411641Z", "iopub.status.busy": "2026-10-02T14:44:37.411377Z", "iopub.status.idle": "2026-10-02T14:44:38.708462Z", "shell.execute_reply": "2026-10-02T14:44:38.707888Z" }, "papermill": { "duration": 1.322399, "end_time": "2026-10-02T14:44:38.709236+00:00", "exception": false, "start_time": "2026-10-02T14:44:37.386837+00:00", "status": "completed" }, "tags": [] }, "outputs": [], "source": [ "import numpy as np\n", "import openmdao.api as om\n", "\n", "\n", "class DynamicPartialsComp(om.ExplicitComponent):\n", " def __init__(self, size):\n", " super().__init__()\n", " self.size = size\n", " self.num_computes = 0\n", "\n", " def setup(self):\n", " self.add_input('y', np.ones(self.size))\n", " self.add_input('x', np.ones(self.size))\n", " self.add_output('g', np.ones(self.size))\n", "\n", " self.declare_partials('*', '*', method='cs')\n", "\n", " # turn on dynamic partial coloring\n", " self.declare_coloring(wrt='*', method='cs', perturb_size=1e-5, num_full_jacs=1, tol=1e-20,\n", " show_summary=True, show_sparsity=False)\n", "\n", " def compute(self, inputs, outputs):\n", " outputs['g'] = np.arctan(inputs['y'] / inputs['x'])\n", " self.num_computes += 1\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "realistic-terrorism", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:38.833173Z", "iopub.status.busy": "2026-10-02T14:44:38.832757Z", "iopub.status.idle": "2026-10-02T14:44:40.043545Z", "shell.execute_reply": "2026-10-02T14:44:40.042607Z" }, "papermill": { "duration": 1.333619, "end_time": "2026-10-02T14:44:40.044124+00:00", "exception": false, "start_time": "2026-10-02T14:44:38.710505+00:00", "status": "completed" }, "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[1790952280.030966] [runnervm8df0l:8720 :0] ib_iface.c:1269 UCX ERROR mana_0: iface 0x562090a1e550 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", "[1790952280.031262] [runnervm8df0l:8720 :0] ucp_worker.c:1412 UCX ERROR uct_iface_open(ud_verbs/mana_0:1) failed: Input/output error\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "[runnervm8df0l:08720] pml_ucx.c:313 Error: Failed to create UCP worker\n" ] } ], "source": [ "import openmdao.api as om\n", "\n", "SIZE = 10\n", "\n", "p = om.Problem()\n", "model = p.model\n", "\n", "# DynamicPartialsComp is set up to do dynamic partial coloring\n", "arctan_yox = model.add_subsystem('arctan_yox', DynamicPartialsComp(SIZE), promotes_inputs=['x', 'y'])\n", "\n", "model.add_subsystem('circle', om.ExecComp('area=pi*r**2'), promotes_inputs=['r'])\n", "\n", "model.add_subsystem('r_con', om.ExecComp('g=x**2 + y**2 - r', has_diag_partials=True,\n", " g=np.ones(SIZE), x=np.ones(SIZE), y=np.ones(SIZE)),\n", " promotes_inputs=['x', 'y', 'r'])\n", "\n", "thetas = np.linspace(0, np.pi/4, SIZE)\n", "model.add_subsystem('theta_con', om.ExecComp('g = x - theta', has_diag_partials=True,\n", " g=np.ones(SIZE), x=np.ones(SIZE),\n", " theta=thetas))\n", "model.add_subsystem('delta_theta_con', om.ExecComp('g = even - odd', has_diag_partials=True,\n", " g=np.ones(SIZE//2), even=np.ones(SIZE//2),\n", " odd=np.ones(SIZE//2)))\n", "\n", "model.add_subsystem('l_conx', om.ExecComp('g=x-1', has_diag_partials=True, g=np.ones(SIZE), x=np.ones(SIZE)),\n", " promotes_inputs=['x'])\n", "\n", "IND = np.arange(SIZE, dtype=int)\n", "ODD_IND = IND[1::2] # all odd indices\n", "EVEN_IND = IND[0::2] # all even indices\n", "\n", "model.connect('arctan_yox.g', 'theta_con.x')\n", "model.connect('arctan_yox.g', 'delta_theta_con.even', src_indices=EVEN_IND)\n", "model.connect('arctan_yox.g', 'delta_theta_con.odd', src_indices=ODD_IND)\n", "\n", "p.driver = om.ScipyOptimizeDriver()\n", "p.driver.options['optimizer'] = 'SLSQP'\n", "p.driver.options['disp'] = False\n", "\n", "#####################################\n", "# set up dynamic total coloring here\n", "p.driver.declare_coloring(show_summary=True, show_sparsity=False)\n", "#####################################\n", "\n", "model.add_design_var('x')\n", "model.add_design_var('y')\n", "model.add_design_var('r', lower=.5, upper=10)\n", "\n", "# nonlinear constraints\n", "model.add_constraint('r_con.g', equals=0)\n", "\n", "model.add_constraint('theta_con.g', lower=-1e-5, upper=1e-5, indices=EVEN_IND)\n", "model.add_constraint('delta_theta_con.g', lower=-1e-5, upper=1e-5)\n", "\n", "# this constrains x[0] to be 1 (see definition of l_conx)\n", "model.add_constraint('l_conx.g', equals=0, linear=False, indices=[0,])\n", "\n", "# linear constraint\n", "model.add_constraint('y', equals=0, indices=[0,], linear=True)\n", "\n", "model.add_objective('circle.area', ref=-1)\n", "\n", "p.setup(mode='fwd')\n", "\n", "# the following were randomly generated using np.random.random(10)*2-1 to randomly\n", "# disperse them within a unit circle centered at the origin.\n", "p.set_val('x', np.array([ 0.55994437, -0.95923447, 0.21798656, -0.02158783, 0.62183717,\n", " 0.04007379, 0.46044942, -0.10129622, 0.27720413, -0.37107886]))\n", "p.set_val('y', np.array([ 0.52577864, 0.30894559, 0.8420792 , 0.35039912, -0.67290778,\n", " -0.86236787, -0.97500023, 0.47739414, 0.51174103, 0.10052582]))\n", "p.set_val('r', .7)" ] }, { "cell_type": "markdown", "id": "approximate-immigration", "metadata": { "papermill": { "duration": 0.074768, "end_time": "2026-10-02T14:44:40.120370+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.045602+00:00", "status": "completed" }, "tags": [] }, "source": [ "Coloring info will be displayed during run_driver. The number of colors in the partial coloring of arctan_yox should be 2 and the number of colors in the total coloring should be 5." ] }, { "cell_type": "code", "execution_count": 4, "id": "acquired-state", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:40.124645Z", "iopub.status.busy": "2026-10-02T14:44:40.124438Z", "iopub.status.idle": "2026-10-02T14:44:40.192848Z", "shell.execute_reply": "2026-10-02T14:44:40.192201Z" }, "papermill": { "duration": 0.071376, "end_time": "2026-10-02T14:44:40.193415+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.122039+00:00", "status": "completed" }, "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "Coloring for 'arctan_yox' (class DynamicPartialsComp)\n", "\n", "Jacobian shape: (10, 20) (10.00% nonzero)\n", "FWD solves: 2 REV solves: 0\n", "Total colors vs. total size: 2 vs 20 (90.00% improvement)\n", "\n", "Sparsity computed using tolerance: 1e-20.\n", "Dense partial jacobian for DynamicPartialsComp 'arctan_yox' was computed 1 times.\n", "Time to compute sparsity: 0.0021 sec\n", "Time to compute coloring: 0.0012 sec\n", "Memory to compute coloring: 0.6953 MB\n", "\n", "Jacobian shape: (22, 21) (13.42% nonzero)\n", "FWD solves: 5 REV solves: 0\n", "Total colors vs. total size: 5 vs 21 (76.19% improvement)\n", "\n", "Sparsity computed using tolerance: 1e-25.\n", "Dense total jacobian for Problem 'problem' was computed 3 times.\n", "Time to compute sparsity: 0.0142 sec\n", "Time to compute coloring: 0.0013 sec\n", "Memory to compute coloring: 0.0391 MB\n", "Coloring created on: 2026-10-02 14:44:40\n" ] }, { "data": { "text/plain": [ "Problem: problem\n", "Driver: ScipyOptimizeDriver\n", " success : True\n", " iterations : 8\n", " runtime : 6.0351E-02 s\n", " model_evals : 8\n", " model_time : 1.5776E-03 s\n", " deriv_evals : 7\n", " deriv_time : 3.6812E-02 s\n", " exit_status : SUCCESS" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "p.run_driver()" ] }, { "cell_type": "code", "execution_count": 5, "id": "affiliated-newport", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:40.196920Z", "iopub.status.busy": "2026-10-02T14:44:40.196581Z", "iopub.status.idle": "2026-10-02T14:44:40.199380Z", "shell.execute_reply": "2026-10-02T14:44:40.198893Z" }, "papermill": { "duration": 0.005405, "end_time": "2026-10-02T14:44:40.200142+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.194737+00:00", "status": "completed" }, "tags": [ "remove-input", "remove-output", "active-ipynb" ] }, "outputs": [], "source": [ "np.testing.assert_allclose(p['circle.area'], np.pi)" ] }, { "cell_type": "markdown", "id": "surprised-helena", "metadata": { "papermill": { "duration": 0.001517, "end_time": "2026-10-02T14:44:40.311516+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.309999+00:00", "status": "completed" }, "tags": [] }, "source": [ "Let's see how many calls to compute we need to determine partials for arctan_yox. The partial derivatives are all diagonal, so we should be able to cover them using only 2 colors." ] }, { "cell_type": "code", "execution_count": 6, "id": "accessible-quick", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:40.315222Z", "iopub.status.busy": "2026-10-02T14:44:40.315000Z", "iopub.status.idle": "2026-10-02T14:44:40.319099Z", "shell.execute_reply": "2026-10-02T14:44:40.318431Z" }, "papermill": { "duration": 0.006692, "end_time": "2026-10-02T14:44:40.319563+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.312871+00:00", "status": "completed" }, "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2\n" ] } ], "source": [ "start_calls = arctan_yox.num_computes\n", "arctan_yox.run_linearize()\n", "print(arctan_yox.num_computes - start_calls)" ] }, { "cell_type": "code", "execution_count": 7, "id": "presidential-soldier", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:40.327169Z", "iopub.status.busy": "2026-10-02T14:44:40.326998Z", "iopub.status.idle": "2026-10-02T14:44:40.329801Z", "shell.execute_reply": "2026-10-02T14:44:40.328957Z" }, "papermill": { "duration": 0.005225, "end_time": "2026-10-02T14:44:40.330331+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.325106+00:00", "status": "completed" }, "tags": [ "remove-input", "remove-output", "active-ipynb" ] }, "outputs": [], "source": [ "assert arctan_yox.num_computes == start_calls + 2" ] }, { "cell_type": "markdown", "id": "short-petite", "metadata": { "papermill": { "duration": 0.016465, "end_time": "2026-10-02T14:44:40.348067+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.331602+00:00", "status": "completed" }, "tags": [] }, "source": [ "## Static Coloring\n", "\n", "Static partial coloring is activated by calling the `use_fixed_coloring` function on the corresponding component after calling `declare_coloring`.\n", "\n", "```{eval-rst}\n", ".. automethod:: openmdao.core.system.System.use_fixed_coloring\n", " :noindex:\n", "```\n", "Generally, no arg will be passed to `use_fixed_coloring`, and OpenMDAO will automatically determine the location and name of the appropriate coloring file, but it is possible to pass the name of a coloring file into `use_fixed_coloring`, and in that case the given coloring file will be used. Note that if a coloring filename is passed into `use_fixed_coloring`, it is assumed that the coloring in that file should *never* be regenerated, even if the user calls `openmdao total_coloring` or `openmdao partial_coloring` from the command line." ] }, { "cell_type": "markdown", "id": "c5049201", "metadata": { "papermill": { "duration": 0.001093, "end_time": "2026-10-02T14:44:40.350349+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.349256+00:00", "status": "completed" }, "tags": [] }, "source": [ "## Coloring Report\n", "\n", "Coloring files (`.pkl`) are written to `{prob.get_outputs_dir()}/coloring_files/`.\n", "OpenMDAO can also generate an HTML sparsity report that visualizes which Jacobian entries\n", "are nonzero and how they are grouped into colors." ] }, { "cell_type": "code", "execution_count": 8, "id": "f631955d", "metadata": { "execution": { "iopub.execute_input": "2026-10-02T14:44:40.432790Z", "iopub.status.busy": "2026-10-02T14:44:40.432507Z", "iopub.status.idle": "2026-10-02T14:44:40.933971Z", "shell.execute_reply": "2026-10-02T14:44:40.933093Z" }, "papermill": { "duration": 0.583645, "end_time": "2026-10-02T14:44:40.935162+00:00", "exception": false, "start_time": "2026-10-02T14:44:40.351517+00:00", "status": "completed" }, "tags": [] }, "outputs": [ { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from openmdao.utils.coloring import display_coloring\n", "from IPython.display import IFrame, display\n", "\n", "display_coloring(p.driver, output_file='total_coloring.html', show=False)\n", "display(IFrame(src='total_coloring.html', width='100%', height='800px'))" ] } ], "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": 4.968281, "end_time": "2026-10-02T14:44:41.653802+00:00", "environment_variables": {}, "exception": null, "input_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/openmdao_book/features/experimental/approx_coloring.ipynb", "output_path": "/home/runner/work/OpenMDAO/OpenMDAO/openmdao/docs/_executed_book/features/experimental/approx_coloring.ipynb", "parameters": {}, "start_time": "2026-10-02T14:44:36.685521+00:00", "version": "2.7.0" } }, "nbformat": 4, "nbformat_minor": 5 }