-
Notifications
You must be signed in to change notification settings - Fork 240
/
complex_step.py
204 lines (161 loc) · 5.85 KB
/
complex_step.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
"""Complex Step derivative approximations."""
from __future__ import division, print_function
from six import iteritems, itervalues
from six.moves import range
from collections import defaultdict
import numpy as np
from openmdao.approximation_schemes.approximation_scheme import ApproximationScheme, \
_gather_jac_results, _get_wrt_subjacs
from openmdao.utils.general_utils import simple_warning
from openmdao.utils.array_utils import sub2full_indices
from openmdao.utils.coloring import Coloring
_full_slice = slice(None)
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,
}
def __init__(self):
"""
Initialize the ApproximationScheme.
"""
super(ComplexStep, self).__init__()
# Only used when nested under complex step.
self._fd = None
def add_approximation(self, abs_key, system, kwargs):
"""
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.
"""
options = self.DEFAULT_OPTIONS.copy()
options.update(kwargs)
key = (abs_key[1], options['step'], options['directional'])
self._exec_dict[key].append((abs_key, options))
self._reset() # force later regen of approx_groups
def _get_approx_data(self, system, data):
"""
Given approximation metadata, compute necessary delta for complex step.
Parameters
----------
system : System
System whose derivatives are being approximated.
data : tuple
Tuple of the form (wrt, delta, directional)
Returns
-------
float
Delta needed for complex step perturbation.
"""
_, delta, _ = data
delta *= 1j
return delta
def compute_approximations(self, system, jac, total=False):
"""
Execute the system to compute the approximate sub-Jacobians.
Parameters
----------
system : System
System on which the execution is run.
jac : dict-like
Approximations are stored in the given dict-like object.
total : bool
If True total derivatives are being approximated, else partials.
"""
if not self._exec_dict:
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
msg = "Nested complex step detected. Finite difference will be used for '%s'."
simple_warning(msg % system.pathname)
fd = self._fd = FiniteDifference()
empty = {}
for lst in itervalues(self._exec_dict):
for apprx in lst:
fd.add_approximation(apprx[0], system, empty)
self._fd.compute_approximations(system, jac, total=total)
return
# Turn on complex step.
system._set_complex_step_mode(True)
self._compute_approximations(system, jac, total, under_cs=True)
# Turn off complex step.
system._set_complex_step_mode(False)
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):
"""
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.
Returns
-------
Vector
Copy of the results from running the perturbed system.
"""
for vec, idxs in idx_info:
if vec is not None:
vec._data[idxs] += delta
if total:
system.run_solve_nonlinear()
results_vec = system._outputs
else:
system.run_apply_nonlinear()
results_vec = system._residuals
result_array[:] = results_vec._data
for vec, idxs in idx_info:
if vec is not None:
vec._data[idxs] -= delta
return result_array