Using src_indices with Promoted Variables

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 indices into the unflattened source. The default of flat_src_indices=False assumes indices map into an unflattened source. Note also that if src_indices are some form of mult-index, for example om.slicer[:, :] or ([0,1,2,3], [0,1,2,3]), then setting flat_src_indices=True is invalid and will raise an exception.

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 numpy as np
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.42769744 0.64733668 0.57657179]
[0.65695096 0.79511797]
[3.3032118]
[5.80827573]

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), and (1,1).  To do that, we use numpy style
        # indexing like this: ([[0, 3], [2, 1]], [[0, 1], [1, 1]]).
        # We've  split up our indexing into a row array and a column array, and we've
        # shaped those arrays to be 2x2, so the result of applying our src_indices to
        # the source array will be a 2x2 array.
        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, 3], [2, 1]], [[0, 1], [1, 1]]))

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 would just access the (3, 1) entry as expected.

    p.model.promotes('C1', inputs=['x'], src_indices=(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.

Note

This feature requires MPI, and may not be able to be run on Colab or Binder.

%%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)), distributed=True)
        self.add_output('y', 1.0, distributed=True)

    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:2] [3. 4.]
[stdout:1] [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:3] [14.]
[stdout:2] [14.]