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
complex vectors will be allocated automatically.
The main situation where complex vectors are
needed but are not allocated automatically is when you call either
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
setup on your
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
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
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
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#
openmdao.utils.cs_safe module contains complex-safe versions of a few common functions, namely,
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'])
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
then run the model by calling
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']
x = prob.get_val('some_output')
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]