-
Notifications
You must be signed in to change notification settings - Fork 240
/
implicit_group.py
128 lines (107 loc) · 4.28 KB
/
implicit_group.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
"""Define a `Group` with two interconnected `ImplicitComponent`s for testing"""
from __future__ import division, print_function
from openmdao.api import Group, ImplicitComponent
from openmdao.api import LinearBlockGS, NonlinearBlockGS
class Comp(ImplicitComponent):
def __init__(self, use_varsets=True):
super(Comp, self).__init__()
self._use_var_sets = use_varsets
def setup(self):
if self._use_var_sets:
self.add_input('a', var_set=1)
self.add_input('b', var_set=0)
self.add_input('c', var_set=1)
self.add_input('d', var_set=2)
self.add_output('w', var_set=5)
self.add_output('x', var_set=1)
self.add_output('y', var_set=1)
self.add_output('z', var_set=5)
else:
self.add_input('a')
self.add_input('b')
self.add_input('c')
self.add_input('d')
self.add_output('w')
self.add_output('x')
self.add_output('y')
self.add_output('z')
def apply_nonlinear(self, inputs, outputs, residuals):
residuals['w'] = outputs['w'] + 2 * inputs['a']
residuals['x'] = outputs['x'] + 3 * inputs['b']
residuals['y'] = outputs['y'] + 4 * inputs['c']
residuals['z'] = outputs['z'] + 5 * inputs['d']
def apply_linear(self, inputs, outputs,
d_inputs, d_outputs, d_residuals, mode):
if mode == 'fwd':
if 'w' in d_outputs:
d_residuals['w'] += d_outputs['w']
if 'a' in d_inputs:
d_residuals['w'] += 2 * d_inputs['a']
if 'x' in d_outputs:
d_residuals['x'] += d_outputs['x']
if 'b' in d_inputs:
d_residuals['x'] += 3 * d_inputs['b']
if 'y' in d_outputs:
d_residuals['y'] += d_outputs['y']
if 'c' in d_inputs:
d_residuals['y'] += 4 * d_inputs['c']
if 'z' in d_outputs:
d_residuals['z'] += d_outputs['z']
if 'd' in d_inputs:
d_residuals['z'] += 5 * d_inputs['d']
else:
if 'w' in d_outputs:
d_outputs['w'] += d_residuals['w']
if 'a' in d_inputs:
d_inputs['a'] += 2 * d_residuals['w']
if 'x' in d_outputs:
d_outputs['x'] += d_residuals['x']
if 'b' in d_inputs:
d_inputs['b'] += 3 * d_residuals['x']
if 'y' in d_outputs:
d_outputs['y'] += d_residuals['y']
if 'c' in d_inputs:
d_inputs['c'] += 4 * d_residuals['y']
if 'z' in d_outputs:
d_outputs['z'] += d_residuals['z']
if 'd' in d_inputs:
d_inputs['d'] += 5 * d_residuals['z']
def solve_linear(self, d_outputs, d_residuals, mode):
if mode == 'fwd':
out_vec = d_outputs
in_vec = d_residuals
elif mode == 'rev':
in_vec = d_outputs
out_vec = d_residuals
for var in ['w', 'x', 'y', 'z']:
out_vec[var] = in_vec[var]
class TestImplicitGroup(Group):
"""
A `Group` with two interconnected <ImplicitComponent>s.
"""
def __init__(self, lnSolverClass=LinearBlockGS,
nlSolverClass=NonlinearBlockGS,
use_varsets=True):
super(TestImplicitGroup, self).__init__()
self.add_subsystem("C1", Comp(use_varsets))
self.add_subsystem("C2", Comp(use_varsets))
self.connect("C1.w", "C2.a")
self.connect("C1.x", "C2.b")
self.connect("C1.y", "C2.c")
self.connect("C1.z", "C2.d")
self.connect("C2.w", "C1.a")
self.connect("C2.x", "C1.b")
self.connect("C2.y", "C1.c")
self.connect("C2.z", "C1.d")
self.linear_solver = lnSolverClass()
self.nonlinear_solver = nlSolverClass()
if use_varsets:
self.expected_solution = [
[1./4., 1./5., 1./4., 1./5.],
[1./3., 1./6., 1./3., 1./6.]
]
else:
self.expected_solution = [
[1./3., 1./4., 1./5., 1./6.,
1./3., 1./4., 1./5., 1./6.]
]