# Source code for openmdao.approximation_schemes.complex_step

"""Complex Step derivative approximations."""

from collections import defaultdict

import numpy as np

from openmdao.utils.om_warnings import issue_warning, DerivativesWarning
from openmdao.approximation_schemes.approximation_scheme import ApproximationScheme

[docs]class ComplexStep(ApproximationScheme):
r"""
Approximation scheme using complex step to calculate derivatives.

For example, using  a step size of 'h' will approximate the derivative in
the following way:

.. math::

f'(x) = \Im{\frac{f(x+ih)}{h}}.

Attributes
----------
_fd : <FiniteDifference>
When nested complex step is detected, we switch to Finite Difference.
"""

DEFAULT_OPTIONS = {
'step': 1e-40,
'directional': False,
}

[docs]    def __init__(self):
"""
Initialize the ApproximationScheme.
"""
super().__init__()

# Only used when nested under complex step.
self._fd = None

[docs]    def add_approximation(self, abs_key, system, kwargs, vector=None):
"""
Use this approximation scheme to approximate the derivative d(of)/d(wrt).

Parameters
----------
abs_key : tuple(str,str)
Absolute name pairing of (of, wrt) for the derivative.
system : System
Containing System.
kwargs : dict
Additional keyword arguments, to be interpreted by sub-classes.
vector : ndarray or None
Direction for difference when using directional derivatives.
"""
options = self.DEFAULT_OPTIONS.copy()
options.update(kwargs)
options['vector'] = vector

wrt = abs_key[1]
if wrt in self._wrt_meta:
self._wrt_meta[wrt].update(options)
else:
self._wrt_meta[wrt] = options
self._reset()  # force later regen of approx_groups

def _get_approx_data(self, system, wrt, meta):
"""
Given approximation metadata, compute necessary delta for complex step.

Parameters
----------
system : System
System whose derivatives are being approximated.
wrt : str
Name of wrt variable.
meta : dict

Returns
-------
float
Delta needed for complex step perturbation.
"""
step = meta['step']
step *= 1j
return step

[docs]    def compute_approx_col_iter(self, system, under_cs=False):
"""
Execute the system to compute the approximate sub-Jacobians.

Parameters
----------
system : System
System on which the execution is run.
under_cs : bool
True if we're currently under complex step at a higher level.

Yields
------
int
column index
ndarray
solution array corresponding to the jacobian column at the given column index
"""
if not self._wrt_meta:
return

if system.under_complex_step:

# If we are nested under another complex step, then warn and swap to FD.
if not self._fd:
from openmdao.approximation_schemes.finite_difference import FiniteDifference

issue_warning("Nested complex step detected. Finite difference will be used.",
prefix=system.pathname, category=DerivativesWarning)

fd = self._fd = FiniteDifference()
empty = {}
for wrt in self._wrt_meta:

yield from self._fd.compute_approx_col_iter(system)
return

saved_inputs = system._inputs._get_data().copy()
system._inputs._data.imag[:] = 0.0
saved_outputs = system._outputs.asarray(copy=True)
system._outputs._data.imag[:] = 0.0
saved_resids = system._residuals.asarray(copy=True)
system._residuals._data.imag[:] = 0.0

# Turn on complex step.
system._set_complex_step_mode(True)

try:
yield from self._compute_approx_col_iter(system, under_cs=True)
finally:
# Turn off complex step.
system._set_complex_step_mode(False)

system._inputs.set_val(saved_inputs)
system._outputs.set_val(saved_outputs)
system._residuals.set_val(saved_resids)

def _get_multiplier(self, delta):
"""
Return a multiplier to be applied to the jacobian.

Parameters
----------
delta : complex
Complex number used to compute the multiplier.

Returns
-------
float
multiplier to apply to the jacobian.
"""
return (1.0 / delta * 1j).real

def _transform_result(self, array):
"""
Return the imaginary part of the given array.

Parameters
----------
array : ndarray of complex
Result array after doing a complex step.

Returns
-------
ndarray
Imaginary part of the result array.
"""
return array.imag

def _run_point(self, system, idx_info, delta, result_array, total, idx_start=0):
"""
Perturb the system inputs with a complex step, run, and return the results.

Parameters
----------
system : System
The system having its derivs approximated.
idx_info : tuple of (Vector, ndarray of int)
Tuple of wrt indices and corresponding data vector to perturb.
delta : complex
Perturbation amount.
result_array : ndarray
An array used to store the results.
total : bool
If True total derivatives are being approximated, else partials.
idx_start : int
Vector index of the first element of this wrt variable.

Returns
-------
ndarray
Copy of the outputs or residuals array after running the perturbed system.
"""
for vec, idxs in idx_info:
if vec is not None and idxs is not None:

if total:
system.run_solve_nonlinear()
result_array[:] = system._outputs.asarray()
else:
system.run_apply_nonlinear()
result_array[:] = system._residuals.asarray()

for vec, idxs in idx_info:
if vec is not None and idxs is not None:
vec.isub(delta, idxs)

return result_array

[docs]    def apply_directional(self, data, direction):
"""
Apply stepsize to direction and embed into approximation data.

Parameters
----------
data : float
Step size for complex step.
direction : ndarray
Vector containing derivative direction.

Returns
-------
ndarray
New step direction.
"""
return data * direction