Splitting up connections between groups - openmdao

Question

I would like to know the best way to split up the connection command. I have two groups that I want to be modular, an inner group and an outer group. I want the inner group to be a kind of black box where I can switch out or change the inner group without changing all the connections for the outer group. I just want the outer group to have to know the inputs and outputs of the inner group. For an example:
import numpy as np
from openmdao.api import Group, Problem, Component, IndepVarComp, ExecComp
class C(Component):
def __init__(self, n):
super(C, self).__init__()
self.add_param('array_element', shape=1)
self.add_output('new_element', shape=1)
def solve_nonlinear(self, params, unknowns, resids):
unknowns['new_element'] = params['array_element']*2.0
class MUX(Component):
def __init__(self, n):
super(MUX, self).__init__()
for i in range(n):
self.add_param('new_element' + str(i), shape=1)
self.add_output('new_array', shape=n)
self.n = n
def solve_nonlinear(self, params, unknowns, resids):
new_array = np.zeros(n)
for i in range(n):
new_array[i] = params['new_element' + str(i)]
unknowns['new_array'] = new_array
class GroupInner(Group):
def __init__(self, n):
super(GroupInner, self).__init__()
for i in range(n):
self.add('c'+str(i), C(n))
self.connect('array', 'c'+str(i) + '.array_element', src_indices=[i])
self.connect('c'+str(i)+'.new_element', 'new_element'+str(i))
self.add('mux', MUX(n), promotes=['*'])
class GroupOuter(Group):
def __init__(self, n):
super(GroupOuter, self).__init__()
self.add('array', IndepVarComp('array', np.zeros(n)), promotes=['*'])
self.add('inner', GroupInner(n), promotes=['new_array'])
for i in range(n):
# self.connect('array', 'inner.c'+str(i) + '.array_element', src_indices=[i])
self.connect('array', 'inner.array', src_indices=[i])
n = 3
p = Problem()
p.root = GroupOuter(n)
p.setup(check=False)
p['array'] = np.ones(n)
p.run()
print p['new_array']
When I run the code I get the error that:
NameError: Source 'array' cannot be connected to target 'c0.array_element': 'array' does not exist.
To try to solve this I made 'array' an IndepVarComp in the GroupInner group. However, when I do this I get the error:
NameError: Source 'array' cannot be connected to target 'inner.array': Target must be a parameter but 'inner.array' is an unknown.
I know that if I just make the full connection: self.connect('array', 'inner.c'+str(i) + '.array_element', src_indices=[i]) then it will work. But like I said I want GroupInner to be kind of a black box where I don't know what groups or components are in it. I also can't just promote all because the array_elements are different. Is it possible to do this or do you have to do the entire connection in one command?

Answer 1

I have two answers to your question. First I'll get the problem working as you specified it. Second, I'll suggest a modification that I think might be more efficient for some applications of this model structure.
First, the main issue with the problem as you specified it was the following line
self.connect('array', 'c'+str(i) + '.array_element', src_indices=[i])
There simply isn't a output or state named array anywhere inside the Inner group, so that connection isn't going to work. You did create a variable called 'array' in the Outer group, but you can't issue a connection to it from inside the Inner definition because its not available in that scope. To make it work the way you've specified, the simplest way would be to do the following:
class GroupInner(Group):
def __init__(self, n):
super(GroupInner, self).__init__()
for i in range(n):
self.add('c'+str(i), C(n))
self.connect('c%d.new_element'%i, 'new_element'+str(i))
self.add('mux', MUX(n), promotes=['*'])
class GroupOuter(Group):
def __init__(self, n):
super(GroupOuter, self).__init__()
self.add('array', IndepVarComp('array', np.zeros(n)), promotes=['*'])
self.add('inner', GroupInner(n), promotes=['new_array'])
for i in range(n):
self.connect('array', 'inner.c%d.array_element'%i, src_indices=[i])
Here is an alternate approach that will reduce the number of variables and components in your model, which will help reduce setup times if n grows large is to use an actual distributed component, and partition the array using the MPI comm. This has some nice properties, besides the setup cost savings, because it will also allow you to scale your calculations more flexibility and improves efficiency when you run in serial. This solution works well if your model would have really had multiple c instances that were all doing the same thing and the process can be simply vectorized via numpy.
import numpy as np
from openmdao.api import Group, Problem, Component, IndepVarComp
from openmdao.util.array_util import evenly_distrib_idxs
from openmdao.core.mpi_wrap import MPI
if MPI:
# if you called this script with 'mpirun', then use the petsc data passing
from openmdao.core.petsc_impl import PetscImpl as impl
else:
# if you didn't use `mpirun`, then use the numpy data passing
from openmdao.api import BasicImpl as impl
class C(Component):
def __init__(self, n):
super(C, self).__init__()
self.add_param('in_array', shape=n)
self.add_output('new_array', shape=n)
self.n = n
def get_req_procs(self):
"""
min/max number of procs that this component can use
"""
return (1,self.n)
#NOTE: This needs to be setup_distrib_idx for <= version 1.5.0
def setup_distrib(self):
comm = self.comm
rank = comm.rank
# NOTE: evenly_distrib_idxs is a helper function to split the array
# up as evenly as possible
sizes, offsets = evenly_distrib_idxs(comm.size, self.n)
local_size, local_offset = sizes[rank], offsets[rank]
self.local_size = int(local_size)
start = local_offset
end = local_offset + local_size
self.set_var_indices('in_array', val=np.zeros(local_size, float),
src_indices=np.arange(start, end, dtype=int))
self.set_var_indices('new_array', val=np.zeros(local_size, float),
src_indices=np.arange(start, end, dtype=int))
def solve_nonlinear(self, params, unknowns, resids):
unknowns['new_array'] = params['in_array']*2.0
print "computing new_array: ", unknowns['new_array']
class GroupInner(Group):
def __init__(self, n):
super(GroupInner, self).__init__()
self.add('c', C(n), promotes=['new_array', 'in_array'])
class GroupOuter(Group):
def __init__(self, n):
super(GroupOuter, self).__init__()
self.add('array', IndepVarComp('array', np.zeros(n)), promotes=['*'])
self.add('inner', GroupInner(n), promotes=['new_array',])
self.connect('array', 'inner.in_array')
n = 3
p = Problem(impl=impl)
p.root = GroupOuter(n)
p.setup(check=False)
p['array'] = np.ones(n)
p.run()
print p['new_array']