Simultaneous Coloring of Approximated Derivatives

Simultaneous Coloring of Approximated Derivatives#

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. 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, plot 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 nonzero values will be computed 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.

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

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

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.

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.

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=1, tol=1e-20,
                              show_summary=True, show_sparsity=True)

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

import openmdao.api as om

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()

Coloring for 'arctan_yox' (class DynamicPartialsComp)

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

Sparsity computed using tolerance: 1e-20
Time to compute sparsity:   0.0008 sec
Time to compute coloring:   0.0019 sec
Memory to compute coloring:   0.2500 MB

Full total jacobian for problem 'problem' was computed 3 times, taking 0.22742042999993828 seconds.
Total jacobian shape: (22, 21) 


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

Sparsity computed using tolerance: 1e-25
Time to compute sparsity:   0.2274 sec
Time to compute coloring:   0.0024 sec
Memory to compute coloring:   0.2500 MB
Coloring created on: 2024-04-24 17:50:12

False

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)

Static Coloring#

Static partial coloring is activated by calling the use_fixed_coloring function on the corresponding component 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.

Simultaneous Coloring of Approximated Derivatives

Contents

Simultaneous Coloring of Approximated Derivatives#

Dynamic Coloring#

Static Coloring#