Simultaneous Coloring of Approximated Derivatives

In OpenMDAO, partial derivatives for components and semi-total derivatives for groups can be approximated using either finite difference or complex step. Sometimes the partial or semi-total 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. 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.

Dynamic Coloring

Setting up a problem to use dynamic coloring of approximated derivatives requires a call to the declare_coloring function.

System.declare_coloring(wrt=('*'), method='fd', form=None, step=None, per_instance=True, num_full_jacs=3, tol=1e-25, orders=None, perturb_size=1e-09, min_improve_pct=5.0, show_summary=True, show_sparsity=False)[source]

Set options for deriv coloring of a set of wrt vars matching the given pattern(s).

Parameters
wrtstr or list of str

The name or names of the variables that derivatives are taken with respect to. This can contain input names, output names, or glob patterns.

methodstr

Method used to compute derivative: “fd” for finite difference, “cs” for complex step.

formstr

Finite difference form, can be “forward”, “central”, or “backward”. Leave undeclared to keep unchanged from previous or default value.

stepfloat

Step size for finite difference. Leave undeclared to keep unchanged from previous or default value.

per_instancebool

If True, a separate coloring will be generated for each instance of a given class. Otherwise, only one coloring for a given class will be generated and all instances of that class will use it.

num_full_jacsint

Number of times to repeat partial jacobian computation when computing sparsity.

tolfloat

Tolerance used to determine if an array entry is nonzero during sparsity determination.

ordersint

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

perturb_sizefloat

Size of input/output perturbation during generation of sparsity.

min_improve_pctfloat

If coloring does not improve (decrease) the number of solves more than the given percentage, coloring will not be used.

show_summarybool

If True, display summary information after generating coloring.

show_sparsitybool

If True, display sparsity with coloring info after generating coloring.

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.

comp = prob.model.add_subsystem('comp', MyComp())
comp.declare_coloring('x*', method='cs', num_full_jacs=2, min_improve_pct=10.)

Note that in the call to declare_coloring, we also set num_full_jacs to 2. This means that the first 2 times that a partial jacobian is computed for ‘comp’, it’s values will be computed without coloring and stored. Just prior to the 3rd time, the coloring will 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, then comp will not use coloring at all.

Semi-total derivative coloring can be performed in a similar way, except that declare_coloring would be called on a Group instead of a Component. num_full_jacs can also be passed as an arg to declare_coloring on the Group.

The purpose of declare_coloring is to provide all of the necessary information to allow OpenMDAO to generate a coloring, either dynamically or manually using openmdao partial_coloring.

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 or group being colored. If False, the name will be based on the full module pathname of the class of the given component or group.

declare_coloring should generally be called in the setup function of the component or group.

Here’s a modified version of our total coloring example, where we replace one of our components with one 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.

import numpy as np
import openmdao.api as om

class DynamicPartialsComp(om.ExplicitComponent):
    def __init__(self, size):
        super().__init__()
        self.size = size
        self.num_computes = 0

    def setup(self):
        self.add_input('y', np.ones(self.size))
        self.add_input('x', np.ones(self.size))
        self.add_output('g', np.ones(self.size))

        self.declare_partials('*', '*', method='cs')

        # turn on dynamic partial coloring
        self.declare_coloring(wrt='*', method='cs', perturb_size=1e-5, num_full_jacs=2, tol=1e-20,
                              orders=20, show_summary=True, show_sparsity=True)

    def compute(self, inputs, outputs):
        outputs['g'] = np.arctan(inputs['y'] / inputs['x'])
        self.num_computes += 1


SIZE = 10

p = om.Problem()
model = p.model

########################################################################
# DynamicPartialsComp is set up to do dynamic partial coloring
arctan_yox = model.add_subsystem('arctan_yox', DynamicPartialsComp(SIZE),
                                 promotes_inputs=['x', 'y'])
########################################################################

model.add_subsystem('circle', om.ExecComp('area=pi*r**2'),
                    promotes_inputs=['r'])

model.add_subsystem('r_con', om.ExecComp('g=x**2 + y**2 - r', has_diag_partials=True,
                                         g=np.ones(SIZE), x=np.ones(SIZE), y=np.ones(SIZE)),
                    promotes_inputs=['x', 'y', 'r'])

thetas = np.linspace(0, np.pi/4, SIZE)
model.add_subsystem('theta_con', om.ExecComp('g = x - theta', has_diag_partials=True,
                                               g=np.ones(SIZE), x=np.ones(SIZE),
                                               theta=thetas))
model.add_subsystem('delta_theta_con', om.ExecComp('g = even - odd', has_diag_partials=True,
                                                     g=np.ones(SIZE//2), even=np.ones(SIZE//2),
                                                     odd=np.ones(SIZE//2)))

model.add_subsystem('l_conx', om.ExecComp('g=x-1', has_diag_partials=True,
                                          g=np.ones(SIZE), x=np.ones(SIZE)),
                    promotes_inputs=['x'])

IND = np.arange(SIZE, dtype=int)
ODD_IND = IND[1::2]  # all odd indices
EVEN_IND = IND[0::2]  # all even indices

model.connect('arctan_yox.g', 'theta_con.x')
model.connect('arctan_yox.g', 'delta_theta_con.even', src_indices=EVEN_IND)
model.connect('arctan_yox.g', 'delta_theta_con.odd', src_indices=ODD_IND)

p.driver = om.ScipyOptimizeDriver()
p.driver.options['optimizer'] = 'SLSQP'
p.driver.options['disp'] = False

#####################################
# set up dynamic total coloring here
p.driver.declare_coloring(show_summary=True, show_sparsity=True)
#####################################

model.add_design_var('x')
model.add_design_var('y')
model.add_design_var('r', lower=.5, upper=10)

# nonlinear constraints
model.add_constraint('r_con.g', equals=0)

model.add_constraint('theta_con.g', lower=-1e-5, upper=1e-5, indices=EVEN_IND)
model.add_constraint('delta_theta_con.g', lower=-1e-5, upper=1e-5)

# this constrains x[0] to be 1 (see definition of l_conx)
model.add_constraint('l_conx.g', equals=0, linear=False, indices=[0,])

# linear constraint
model.add_constraint('y', equals=0, indices=[0,], linear=True)

model.add_objective('circle.area', ref=-1)

p.setup(mode='fwd')

# the following were randomly generated using np.random.random(10)*2-1 to randomly
# disperse them within a unit circle centered at the origin.
p.set_val('x', np.array([ 0.55994437, -0.95923447,  0.21798656, -0.02158783,  0.62183717,
                          0.04007379,  0.46044942, -0.10129622,  0.27720413, -0.37107886]))
p.set_val('y', np.array([ 0.52577864,  0.30894559,  0.8420792 ,  0.35039912, -0.67290778,
                          -0.86236787, -0.97500023,  0.47739414,  0.51174103,  0.10052582]))
p.set_val('r', .7)

# 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.
p.run_driver()
Approx coloring for 'arctan_yox' (class DynamicPartialsComp)
f.........f......... 0  g
.f.........f........ 1  g
..f.........f....... 2  g
...f.........f...... 3  g
....f.........f..... 4  g
.....f.........f.... 5  g
......f.........f... 6  g
.......f.........f.. 7  g
........f.........f. 8  g
.........f.........f 9  g
|y
          |x

Jacobian shape: (10, 20)  (10.00% nonzero)
FWD solves: 2   REV solves: 0
Total colors vs. total size: 2 vs 20  (90.0% improvement)

Sparsity computed using tolerance: 1e-40
Most common number of nonzero entries (20 of 200) repeated 39 times out of 40 tolerances tested.

Time to compute sparsity: 0.002229 sec.
Time to compute coloring: 0.000459 sec.
Full total jacobian was computed 3 times, taking 0.023547 seconds.
Total jacobian shape: (22, 21) 

....................f 0  circle.area
f.........f.........f 1  r_con.g
.f.........f........f 2  r_con.g
..f.........f.......f 3  r_con.g
...f.........f......f 4  r_con.g
....f.........f.....f 5  r_con.g
.....f.........f....f 6  r_con.g
......f.........f...f 7  r_con.g
.......f.........f..f 8  r_con.g
........f.........f.f 9  r_con.g
.........f.........ff 10  r_con.g
f.........f.......... 11  theta_con.g
..f.........f........ 12  theta_con.g
....f.........f...... 13  theta_con.g
......f.........f.... 14  theta_con.g
........f.........f.. 15  theta_con.g
ff........ff......... 16  delta_theta_con.g
..ff........ff....... 17  delta_theta_con.g
....ff........ff..... 18  delta_theta_con.g
......ff........ff... 19  delta_theta_con.g
........ff........ff. 20  delta_theta_con.g
f.................... 21  l_conx.g
|_auto_ivc.v1
          |_auto_ivc.v0
                    |_auto_ivc.v2

Jacobian shape: (22, 21)  (13.42% nonzero)
FWD solves: 5   REV solves: 0
Total colors vs. total size: 5 vs 21  (76.2% improvement)

Sparsity computed using tolerance: 1e-25
Time to compute sparsity: 0.023547 sec.
Time to compute coloring: 0.000589 sec.
print(p['circle.area'])
[3.14159265]
# 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.
start_calls = arctan_yox.num_computes
arctan_yox.run_linearize()
print(arctan_yox.num_computes - start_calls)
2

Static Coloring

Static partial or semi-total derivative coloring is activated by calling the use_fixed_coloring function on the corresponding component or group, after calling declare_coloring.

System.use_fixed_coloring(coloring=<object object>, recurse=True)[source]

Use a precomputed coloring for this System.

Parameters
coloringstr

A coloring filename. If no arg is passed, filename will be determined automatically.

recursebool

If True, set fixed coloring in all subsystems that declare a coloring. Ignored if a specific coloring is passed in.

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.