# coloring.py¶

Routines to compute coloring for use with simultaneous derivatives.

class openmdao.utils.coloring.Coloring(sparsity, row_vars=None, row_var_sizes=None, col_vars=None, col_var_sizes=None)[source]

Bases: object

Container for all information relevant to a coloring.

Parameters
sparsityndarray

Full jacobian sparsity matrix (dense bool form).

row_varslist of str or None

Names of variables corresponding to rows.

row_var_sizesndarray or None

Sizes of row variables.

col_varslist of str or None

Names of variables corresponding to columns.

col_var_sizesndarray or None

Sizes of column variables.

Attributes
_shapetuple of int (nrows, ncols)

Tuple describing the shape of the sparsity matrix.

_nzrowsndarray of int

Row indices of nonzero entries in the full jac sparsity matrix.

_nzcolsndarray of int

Column indices of nonzero entries in the full jac sparsity matrix.

_pct_nonzerofloat

If known, percentage of nonzero vs total array entries.

_fwdtuple (col_lists, row_maps) or None

Contains lists of grouped columns and nonzero rows for each column for forward coloring.

_revtuple (col_lists, row_maps) or None

Contains lists of grouped columns and nonzero rows for each column for reverse coloring.

_col_varslist of str or None

Names of variables corresponding to columns.

_col_var_sizesndarray or None

Sizes of column variables.

_row_varslist of str or None

Names of variables corresponding to rows.

_row_var_sizesndarray or None

Sizes of row variables.

Dictionary of metadata used to create the coloring.

_names_arrayndarray or None:

Names of total jacobian rows or columns.

_local_arrayndarray or None:

Indices of total jacobian rows or columns.

__init__(sparsity, row_vars=None, row_var_sizes=None, col_vars=None, col_var_sizes=None)[source]

Initialize data structures.

color_iter(direction)[source]

Given a direction, yield an iterator over column (or row) groups.

Parameters
directionstr

Derivative direction (‘fwd’ or ‘rev’).

Yields
list of int

Lists of column indices (in fwd mode) or row indices (in rev mode).

color_nonzero_iter(direction)[source]

Given a direction, yield an iterator over (columns, nz_rows) or (rows, nz_columns).

Parameters
directionstr

Indicates which coloring subdict (‘fwd’ or ‘rev’) to use.

Yields
(column or row groups, nonzero row or column lists)

Yields a list of columns/rows and their associated nonzero rows/columns for each color.

colored_jac_iter(compressed_j, direction, trans=None)[source]

Yield nonzero parts of columns (fwd) or rows (rev) of a colored jacobian.

Parameters
compressed_jndarray

The compressed jacobian.

directionstr

Derivative computation direction (‘fwd’ or ‘rev’).

transndarray

Index array to translate from compressed jac in function context to openmdao jac.

Yields
ndarray

Nonzero part of current jacobian column or row.

ndarray

Indices into full jacobian column or row where nonzero values should be placed.

int

Index into the full jacobian of the current column or row.

display(show=True, fname='coloring.png')[source]

Display a plot of the sparsity pattern, showing grouping by color.

Parameters
showbool

If True, show the plot. Otherwise, just save the plot in a file. Default is True.

fnamestr

Path to the location where the plot file should be saved.

display_txt()[source]

Print the structure of a boolean array with coloring info for each nonzero value.

Forward mode colored nonzeros are denoted by ‘f’, reverse mode nonzeros by ‘r’, overlapping nonzeros by ‘O’ and uncolored nonzeros by ‘x’. Zeros are denoted by ‘.’. Note that x’s and O’s should never appear unless there is a bug in the coloring algorithm.

If names and sizes of row and column vars are known, print the name of the row var alongside each row and print the names of the column vars, aligned with each column, at the bottom.

expand_jac(compressed_j, direction)[source]

Expand the given compressed jacobian into a full jacobian.

Parameters
compressed_jndarray

The compressed jacobian.

directionstr

Derivative computation direction (‘fwd’ or ‘rev’).

Returns
ndarray

The full jacobian.

get_declare_partials_calls()[source]

Return a string containing declare_partials() calls based on the subjac sparsity.

Returns
str

A string containing a declare_partials() call for each nonzero subjac. This string may be cut and pasted into a component’s setup() method.

get_dense_sparsity(dtype=<class 'bool'>)[source]

Return a dense array representing the full sparsity.

Parameters
dtypeobject

Data type of returned numpy array.

Returns
ndarray

Dense sparsity matrix.

get_row_col_map(direction)[source]

Return mapping of nonzero rows to each column (fwd) or nonzeros columns to each row (rev).

Parameters
directionstr

Indicator of forward mode (‘fwd’) or reverse mode (‘rev’).

Returns
list of lists of int

List where each entry contains list of nonzero rows/cols for the index corresponding to each column/row.

get_row_var_coloring(varname)[source]

Return the number of fwd and rev solves needed for a particular row variable.

Parameters
varnamestr

Name of the row variable.

Returns
int

Number of forward solves needed for the given variable.

int

Number of reverse solves needed for the given variable.

get_subjac_sparsity()[source]

Compute the sparsity structure of each subjacobian based on the full jac sparsity.

If row/col variables and sizes are not known, returns None.

Returns
dict or None

Mapping of (of, wrt) keys to their corresponding (nzrows, nzcols, shape).

Read the coloring object from the given pickle file.

Parameters
fnamestr

Name of file to read from.

Returns
Coloring

See docstring for Coloring class.

modes()[source]

Return a tuple containing the modes included in this coloring.

Returns
tuple

Tuple containing some subset of (‘fwd’, ‘rev’).

save(fname)[source]

Write the coloring object to the given stream, creating intermediate dirs if needed.

Parameters
fnamestr

File to save to.

summary()[source]

Print a summary of this coloring.

tangent_iter(direction, arr=None, trans=None)[source]

Given a direction, return input (fwd) or output (rev) tangent arrays.

Each array will contain multiple columns/rows of the identity matrix that share the same color.

Parameters
directionstr

Indicates which coloring subdict (‘fwd’ or ‘rev’) to use.

arrndarray or None

Storage for the current array value.

transndarray or None

Index translation array.

Yields
ndarray

tangent array for inputs (fwd) or outputs (rev)

tangent_matrix(direction, trans=None)[source]

Return a tangent or cotangent matrix for use with jax.

Parameters
directionstr

Derivative computation direction (‘fwd’ or ‘rev’).

transndarray or None

Index translation array.

Returns
ndarray

The tangent or cotangent matrix.

total_solves(fwd=True, rev=True)[source]

Return total number of solves required based on the given coloring info.

Parameters
fwdbool

If True, add fwd colors to total.

revbool

If True, add rev colors to total.

Returns
int

Total number of solves required to compute the jacobian.

openmdao.utils.coloring.MNCO_bidir(J)[source]

Compute bidirectional coloring using Minimum Nonzero Count Order (MNCO).

Based on the algorithm found in Coleman, T.F., Verma, A. (1998) The efficient Computation of Sparse Jacobian Matrices Using Automatic Differentiation. SIAM Journal on Scientific Computing, 19(4), 1210-1233.

Parameters
Jcoo_matrix

Jacobian sparsity matrix (boolean).

Returns
Coloring

See docstring for Coloring class.

openmdao.utils.coloring.compute_total_coloring(problem, mode=None, of=None, wrt=None, num_full_jacs=3, tol=1e-25, orders=None, setup=False, run_model=False, fname=None, use_abs_names=False)[source]

Compute simultaneous derivative colorings for the total jacobian of the given problem.

Parameters
problemProblem

The Problem being analyzed.

modestr or None

The direction for computing derivatives. If None, use problem._mode.

ofiter of str or None

Names of the ‘response’ variables.

wrtiter of str or None

Names of the ‘design’ variables.

num_full_jacsint

Number of times to repeat total jacobian computation.

tolfloat

Tolerance used to determine if an array entry is nonzero.

ordersint

Number of orders above and below the tolerance to check during the tolerance sweep.

setupbool

If True, run setup before calling compute_totals.

run_modelbool

If True, run run_model before calling compute_totals.

fnamefilename or None

File where output coloring info will be written. If None, no info will be written.

use_abs_namesbool

If True, use absolute naming for of and wrt variables unless they are aliases.

Returns
Coloring

See docstring for Coloring class.

openmdao.utils.coloring.dynamic_total_coloring(driver, run_model=True, fname=None)[source]

Compute simultaneous deriv coloring during runtime.

Parameters
driver<Driver>

The driver performing the optimization.

run_modelbool

If True, call run_model before computing coloring.

fnamestr or None

Name of file where coloring will be saved.

Returns
Coloring

The computed coloring.