/
build4test.py
133 lines (102 loc) · 4.02 KB
/
build4test.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
"""
Various functions to make it easier to build test models.
"""
from __future__ import print_function
import time
import numpy
from six.moves import range
from openmdao.core.group import Group
from openmdao.core.parallel_group import ParallelGroup
from openmdao.core.explicitcomponent import ExplicitComponent
from openmdao.core.indepvarcomp import IndepVarComp
from openmdao.test_suite.components.exec_comp_for_test import ExecComp4Test
from openmdao.utils.mpi import MPI
class DynComp(ExplicitComponent):
"""
A component with a settable number of params and outputs.
"""
def __init__(self, ninputs, noutputs,
nl_sleep=0.001, ln_sleep=0.001,
var_factory=float, vf_args=()):
super(DynComp, self).__init__()
self.ninputs = ninputs
self.noutputs = noutputs
self.var_factory = var_factory
self.vf_args = vf_args
self.nl_sleep = nl_sleep
self.ln_sleep = ln_sleep
def setup(self):
for i in range(self.ninputs):
self.add_input('i%d'%i, self.var_factory(*self.vf_args))
for i in range(self.noutputs):
self.add_output("o%d"%i, self.var_factory(*self.vf_args))
def compute(self, inputs, outputs):
time.sleep(self.nl_sleep)
def compute_partials(self, inputs, partials):
"""
Jacobian for Sellar discipline 1.
"""
time.sleep(self.ln_sleep)
def make_subtree(parent, nsubgroups, levels,
ncomps, ninputs, noutputs, nconns, var_factory=float):
"""Construct a system subtree under the given parent group."""
if levels <= 0:
return
if levels == 1: # add leaf nodes
create_dyncomps(parent, ncomps, ninputs, noutputs, nconns,
var_factory=var_factory)
else: # add more subgroup levels
for i in range(nsubgroups):
g = parent.add_subsystem("G%d"%i, Group())
make_subtree(g, nsubgroups, levels-1,
ncomps, ninputs, noutputs, nconns,
var_factory=var_factory)
def create_dyncomps(parent, ncomps, ninputs, noutputs, nconns,
var_factory=float):
"""Create a specified number of DynComps with a specified number
of variables (ninputs and noutputs), and add them to the given parent
and add the number of specified connections.
"""
for i in range(ncomps):
parent.add_subsystem("C%d" % i, DynComp(ninputs, noutputs, var_factory=var_factory))
if i > 0:
for j in range(nconns):
parent.connect("C%d.o%d" % (i-1,j), "C%d.i%d" % (i, j))
if __name__ == '__main__':
import sys
from openmdao.core.problem import Problem
from openmdao.devtools.debug import config_summary
vec_size = 1000
num_comps = 50
pts = 2
class SubGroup(Group):
def setup(self):
create_dyncomps(self, num_comps, 2, 2, 2,
var_factory=lambda: numpy.zeros(vec_size))
cname = "C%d"%(num_comps-1)
self.add_objective("%s.o0" % cname)
self.add_constraint("%s.o1" % cname, lower=0.0)
p = Problem()
g = p.model
if 'gmres' in sys.argv:
from openmdao.solvers.linear.scipy_iter_solver import ScipyKrylov
p.root.linear_solver = ScipyKrylov()
g.add_subsystem("P", IndepVarComp('x', numpy.ones(vec_size)))
g.add_design_var("P.x")
par = g.add_subsystem("par", ParallelGroup())
for pt in range(pts):
ptname = "G%d"%pt
ptg = par.add_subsystem(ptname, SubGroup())
#create_dyncomps(ptg, num_comps, 2, 2, 2,
#var_factory=lambda: numpy.zeros(vec_size))
g.connect("P.x", "par.%s.C0.i0" % ptname)
#cname = ptname + '.' + "C%d"%(num_comps-1)
#g.add_objective("par.%s.o0" % cname)
#g.add_constraint("par.%s.o1" % cname, lower=0.0)
p.setup()
p.final_setup()
p.run_model()
#
from openmdao.devtools.memory import max_mem_usage
print("mem:", max_mem_usage())
config_summary(p)