A Guide to Using Complex Step to Compute Derivatives

A Guide to Using Complex Step to Compute Derivatives#

The intent of this guide is to summarize in detail how to use complex step to compute derivatives. It is assumed that you’re already familiar with OpenMDAO usage in general.

Setting Up Complex Step#

When complex step is used somewhere in an OpenMDAO model, OpenMDAO must allocate sufficient memory to contain a complex version of the nonlinear vector, and in many cases also a complex version of the linear vector. OpenMDAO can figure this out automatically most of the time. For example, if any component in your model calls either declare_partials or declare_coloring with method='cs', complex vectors will be allocated automatically.

The main situation where complex vectors are needed but are not allocated automatically is when you call either check_totals or check_partials with method='cs' and nothing in your model natively uses complex step. In that case, you must tell OpenMDAO that complex vectors are required by passing a force_alloc_complex=True argument when calling setup on your Problem. The force_alloc_complex flag will force OpenMDAO to allocate complex nonlinear vectors regardless of what it detects in the model.

Note that while ExecComp components use complex step to compute derivatives by default, they do not require that the OpenMDAO nonlinear vectors are complex because they perform their own internal complex step operation. However, if you declare your own partials on an ExecComp using declare_partials or declare_coloring with method='cs', then that component will use the framework level complex step routines and will be treated as any other component with partials declared in that manner.

How to Tell if a Component is Running Under Complex Step#

A component can tell when it’s running under complex step by checking the value of its under_complex_step attribute. A similar flag, under_finite_difference can be used to tell if a component is running under finite difference. Here’s an example of a component that checks its complex step status in its compute method:

import openmdao.api as om

class MyCheckComp(om.ExplicitComponent):
    def setup(self):
        self.add_input('a', 1.0)
        self.add_output('x', 0.0)
        # because we set method='cs' here, OpenMDAO automatically knows to allocate
        # complex nonlinear vectors
        self.declare_partials(of='*', wrt='*', method='cs')

    def compute(self, inputs, outputs):
        a = inputs['a']
        if self.under_complex_step:
            print('under complex step')
        else:
            print('not under complex step')
        outputs['x'] = a * 2.

p = om.Problem()
p.model.add_subsystem('comp', MyCheckComp())
# don't need to set force_alloc_complex=True here since we call declare_partials with
# method='cs' in our model.
p.setup()

# during run_model, our component's compute will *not* be running under complex step
p.run_model()

# during compute_partials, our component's compute *will* be running under complex step
J = p.compute_totals(of=['comp.x'], wrt=['comp.a'])
print(J['comp.x', 'comp.a'])

not under complex step
under complex step
[[2.]]

Complex Stepping through Solvers#

See Complex Step Guidelines for important issues to consider when your model has nonlinear solvers under a group that is performing complex step.

Using Complex Step to Compute a Partial Jacobian Coloring#

OpenMDAO can compute a coloring of the partial jacobian matrix for a component that uses complex step or finite difference to compute its derivatives. For a detailed explanation of this, see Simultaneous Coloring of Approximated Derivatives. Assuming your component is complex safe, generally using complex step is more accurate than finite difference and should be preferred when computing a coloring of the partial jacobian.

The cs_safe Module#

The openmdao.utils.cs_safe module contains complex-safe versions of a few common functions, namely, abs, norm, and arctan2. The numpy versions of these functions are not complex-safe and so must be replaced with the safe versions if you intend to use such functions in your component under complex step. Note that the ExecComp component, which uses complex step by default, automatically uses the complex safe versions of these functions if they are referenced in one of its expressions.

The following example shows how to make a component that uses abs safe to use under complex step:

import openmdao.api as om
from openmdao.utils.cs_safe import abs as cs_abs

class MyComp(om.ExplicitComponent):
    def setup(self):
        self.add_input('a', 1.0)
        self.add_input('b', -2.0)
        self.add_output('x', 0.0)
        # because we set method='cs' here, OpenMDAO automatically knows to allocate
        # complex nonlinear vectors
        self.declare_partials(of='*', wrt='*', method='cs')

    def compute(self, inputs, outputs):
        a, b = inputs.values()
        # normal abs isn't complex safe, so use cs_abs
        outputs['x'] = a * cs_abs(b) * 2.

p = om.Problem()
p.model.add_subsystem('comp', MyComp())
# don't need to set force_alloc_complex=True here since we call declare_partials with
# method='cs' in our model.
p.setup()
p.run_model()

J = p.compute_totals(of=['comp.x'], wrt=['comp.a', 'comp.b'])
print(J['comp.x', 'comp.a'])
print(J['comp.x', 'comp.b'])

[[4.]]
[[-2.]]

Model Level Complex Step Mode for Debugging#

Sometimes when debugging it may be useful to run a model or part of a model using complex inputs. This can be done by calling set_complex_step_mode(True) on the Problem instance. You can then call

prob['some_var'] = a_complex_val

or

prob.set_val('some_var', a_complex_val)

then run the model by calling

prob.run_model()

The complex values will carry through the model as it runs. Note that this only works if all of the components in the model are complex-safe. After the model run has completed, the outputs can then be inspected using one of the following:

x = prob['some_output']

or

x = prob.get_val('some_output')

or

prob.model.list_outputs()

Here’s a short example:

p = om.Problem()
p.model.add_subsystem('comp', MyComp())
p.setup()

p['comp.a'] = 1.5
p['comp.b'] = -3.

p.run_model()

# output x should be a float here
print('float', p['comp.x'])

# now we're setting the problem to use complex step
p.set_complex_step_mode(True)

p['comp.a'] = 1.5+2j
p['comp.b'] = -3.-7j

p.run_model()

# output x should be complex here
print('complex', p['comp.x'])

float [9.]
complex [-19.+33.j]