Using src_indices with Promoted Variables

Inputs and outputs can be connected by promoting them both to the same name, but what if your output is an array and you only want to connect part of it to your input?

If you connect variables via promotion, you must set src_indices when you promote the input with the promotes method, Another argument, flat_src_indices is a boolean that determines whether the entries of the src_indices array are interpreted as indices into the flattened source or as tuples or lists of indices into the unflattened source. The default of flat_src_indices=False assumes indices map into an unflattened source.

Basic Example

Here is a simple example showing how to connect an independent array variable to two different components where each component gets part of the array.

import openmdao.api as om


class MyComp1(om.ExplicitComponent):
    """ multiplies input array by 2. """
    def setup(self):
        self.add_input('x', np.ones(3))
        self.add_output('y', 1.0)

    def compute(self, inputs, outputs):
        outputs['y'] = np.sum(inputs['x'])*2.0

class MyComp2(om.ExplicitComponent):
    """ multiplies input array by 4. """
    def setup(self):
        self.add_input('x', np.ones(2))
        self.add_output('y', 1.0)

    def compute(self, inputs, outputs):
        outputs['y'] = np.sum(inputs['x'])*4.0

class MyGroup(om.Group):
    def setup(self):
        self.add_subsystem('comp1', MyComp1())
        self.add_subsystem('comp2', MyComp2())

    def configure(self):
        # splits input via promotes using src_indices
        self.promotes('comp1', inputs=['x'], src_indices=[0, 1, 2])
        self.promotes('comp2', inputs=['x'], src_indices=[3, 4])

p = om.Problem()

p.model.set_input_defaults('x', np.ones(5))
p.model.add_subsystem('G1', MyGroup(), promotes_inputs=['x'])

p.setup()
inp = np.random.random(5)
p.set_val('x', inp)
p.run_model()

print(p.get_val('G1.comp1.x'))
print(p.get_val('G1.comp2.x'))
print(p.get_val('G1.comp1.y'))
print(p.get_val('G1.comp2.y'))
[0.06684484 0.11124793 0.51076008]
[0.09691343 0.99105836]
[1.37770571]
[4.35188714]

Using src_indices with 2D Arrays

In this example, the source array is shape (4,3) and the input array is shape (2,2)

class MyComp(om.ExplicitComponent):
    def setup(self):
        # We want to pull the following 4 values out of the source:
        # [(0,0), (3,1), (2,1), (1,1)].
        # Because our input is also non-flat we arrange the
        # source index tuples into an array having the same shape
        # as our input.  If we didn't set flat_src_indices to False,
        # we could specify src_indices as a 1D array of indices into
        # the flattened source.
        self.add_input('x', np.ones((2, 2)))
        self.add_output('y', 1.0)

    def compute(self, inputs, outputs):
        outputs['y'] = np.sum(inputs['x'])

p = om.Problem()

# by promoting the following output and inputs to 'x', they will
# be automatically connected
p.model.add_subsystem('indep',
                      om.IndepVarComp('x', np.arange(12).reshape((4, 3))),
                      promotes_outputs=['x'])
p.model.add_subsystem('C1', MyComp())
p.model.promotes('C1', inputs=['x'], src_indices=[[(0, 0), (3, 1)],
                                    [(2, 1), (1, 1)]], flat_src_indices=False)

p.setup()
p.run_model()

print(p.get_val('C1.x'))
print(p.get_val('C1.y'))
[[ 0. 10.]
 [ 7.  4.]]
[21.]

If the source array is shape (4,3), the input is scalar, and we want to connect it to the (3, 1) entry of the source, then the promotes call might look like the following if we use flat_src_indices:

    p.model.promotes('C1', inputs=['x'], src_indices=[10], shape=1, flat_src_indices=True)

If we instead use the default setting of flat_src_indices=False, we must specify the input shape since a scalar input value alone doesn’t necessarily indicate that the input variable is scalar. For example, in the case below, if we didn’t know the input shape, we wouldn’t know if it was scalar and connected to a 2-D source array, or if it was shape (1,2) and connected to a flat source array.

    p.model.promotes('C1', inputs=['x'], src_indices=np.array([[3,1]]), shape=1)

5: If the source array is flat and the input is shape (2,2), the promotes call might look like this:

    p.model.promotes('C1', inputs=['x'], src_indices=[[0, 10], [7, 4]], shape=(2,2))

Note

If the source array is flat, we allow the use of flat src_indices even without setting flat_src_indices=True.

Distributed component example

In the example, a distributed component promotes its input and receives certain entries of the source array based on its rank. Note that negative indices are supported.

%%px

import openmdao.api as om
import numpy as np

class MyComp(om.ExplicitComponent):
    def __init__(self, idxs, **kwargs):
        super().__init__(**kwargs)
        self.idxs = idxs

    def setup(self):
        self.add_input('x', np.ones(len(self.idxs)))
        self.add_output('y', 1.0)

    def compute(self, inputs, outputs):
        outputs['y'] = np.sum(inputs['x'])*2.0

p = om.Problem()

p.model.add_subsystem('indep', om.IndepVarComp('x', np.arange(5, dtype=float)),
                      promotes_outputs=['x'])

# decide what parts of the array we want based on our rank
if p.comm.rank == 0:
    idxs = [0, 1, 2]
else:
    # use [3, -1] here rather than [3, 4] just to show that we
    # can use negative indices.
    idxs = [3, -1]

p.model.add_subsystem('C1', MyComp(idxs))

p.model.promotes('C1', inputs=['x'], src_indices=idxs)

p.setup()
p.set_val('x', np.arange(5, dtype=float))
p.run_model()

# each rank holds the assigned portion of the input array
print(p.get_val('C1.x'))
[stdout:0] [0. 1. 2.]
[stdout:1] [3. 4.]
[stdout:2] [3. 4.]
[stdout:3] [3. 4.]
%%px

# the output in each rank is based on the local inputs
print(p.get_val('C1.y'))
[stdout:0] [6.]
[stdout:1] [14.]
[stdout:2] [14.]
[stdout:3] [14.]