-
Notifications
You must be signed in to change notification settings - Fork 240
/
dictionary_jacobian.py
170 lines (147 loc) · 6.63 KB
/
dictionary_jacobian.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
"""Define the DictionaryJacobian class."""
import numpy as np
from openmdao.jacobians.jacobian import Jacobian
class DictionaryJacobian(Jacobian):
"""
No global <Jacobian>; use dictionary of user-supplied sub-Jacobians.
Attributes
----------
_iter_keys : list of (vname, vname) tuples
List of tuples of variable names that match subjacs in the this Jacobian.
"""
def __init__(self, system, **kwargs):
"""
Initialize all attributes.
Parameters
----------
system : System
Parent system to this jacobian.
**kwargs : dict
options dictionary.
"""
super().__init__(system, **kwargs)
self._iter_keys = {}
def _iter_abs_keys(self, system, vec_name):
"""
Iterate over subjacs keyed by absolute names.
This includes only subjacs that have been set and are part of the current system.
Parameters
----------
system : System
System that is updating this jacobian.
vec_name : str
The name of the current RHS vector.
Returns
-------
list
List of keys matching this jacobian for the current system.
"""
entry = (system.pathname, vec_name)
if entry not in self._iter_keys:
subjacs = self._subjacs_info
keys = []
for res_name in system._var_relevant_names[vec_name]['output']:
for type_ in ('output', 'input'):
for name in system._var_relevant_names[vec_name][type_]:
key = (res_name, name)
if key in subjacs:
keys.append(key)
self._iter_keys[entry] = keys
return self._iter_keys[entry]
def _apply(self, system, d_inputs, d_outputs, d_residuals, mode):
"""
Compute matrix-vector product.
Parameters
----------
system : System
System that is updating this jacobian.
d_inputs : Vector
inputs linear vector.
d_outputs : Vector
outputs linear vector.
d_residuals : Vector
residuals linear vector.
mode : str
'fwd' or 'rev'.
"""
# avoid circular import
from openmdao.core.explicitcomponent import ExplicitComponent
fwd = mode == 'fwd'
d_res_names = d_residuals._names
d_out_names = d_outputs._names
d_inp_names = d_inputs._names
if not d_out_names and not d_inp_names:
return
rflat = d_residuals._abs_get_val
oflat = d_outputs._abs_get_val
iflat = d_inputs._abs_get_val
ncol = d_residuals._ncol
subjacs_info = self._subjacs_info
is_explicit = isinstance(system, ExplicitComponent)
with system._unscaled_context(outputs=[d_outputs], residuals=[d_residuals]):
for abs_key in self._iter_abs_keys(system, d_residuals._name):
res_name, other_name = abs_key
if res_name in d_res_names:
if other_name in d_out_names:
# skip the matvec mult completely for identity subjacs
if is_explicit and res_name is other_name:
if fwd:
val = rflat(res_name)
val -= oflat(other_name)
else:
val = oflat(other_name)
val -= rflat(res_name)
continue
if fwd:
left_vec = rflat(res_name)
right_vec = oflat(other_name)
else:
left_vec = oflat(other_name)
right_vec = rflat(res_name)
elif other_name in d_inp_names:
if fwd:
left_vec = rflat(res_name)
right_vec = iflat(other_name)
else:
left_vec = iflat(other_name)
right_vec = rflat(res_name)
else:
continue
subjac_info = subjacs_info[abs_key]
if self._randomize:
subjac = self._randomize_subjac(subjac_info['value'], abs_key)
else:
subjac = subjac_info['value']
rows = subjac_info['rows']
if rows is not None: # our homegrown COO format
linds, rinds = rows, subjac_info['cols']
if not fwd:
linds, rinds = rinds, linds
if self._under_complex_step:
# bincount only works with float, so split into parts
if ncol > 1:
for i in range(ncol):
prod = right_vec[:, i][rinds] * subjac
left_vec[:, i].real += np.bincount(linds, prod.real,
minlength=left_vec.shape[0])
left_vec[:, i].imag += np.bincount(linds, prod.imag,
minlength=left_vec.shape[0])
else:
prod = right_vec[rinds] * subjac
left_vec[:].real += np.bincount(linds, prod.real,
minlength=left_vec.size)
left_vec[:].imag += np.bincount(linds, prod.imag,
minlength=left_vec.size)
else:
if ncol > 1:
for i in range(ncol):
left_vec[:, i] += np.bincount(linds,
right_vec[:, i][rinds] * subjac,
minlength=left_vec.shape[0])
else:
left_vec[:] += np.bincount(linds, right_vec[rinds] * subjac,
minlength=left_vec.size)
else:
if not fwd:
subjac = subjac.transpose()
left_vec += subjac.dot(right_vec)