parallel_list_outputs.py

import openmdao.api as om
import mpi4py.MPI as MPI

class TestComp(om.ExplicitComponent):

    def initialize(self):
        self.options['distributed'] = False

    def setup(self):
        self.add_input('x', shape=1)
        self.add_output('y', shape=1)
        self.declare_partials('y', 'x')
    
    def compute(self, inputs, outputs):
        outputs['y'] = inputs['x'] ** 2

    def compute_partials(self, inputs, J):
        J['y', 'x'] = 2 * inputs[x]


class TestComp2(om.ExplicitComponent):

    def initialize(self):
        self.options['distributed'] = False

    def setup(self):
        self.add_input('x', shape=1)
        self.add_output('y', shape=1)
        self.declare_partials('y', 'x')
    
    def compute(self, inputs, outputs):
        print(self.comm.rank)
        outputs['y'] = inputs['x'] ** 2

    def compute_partials(self, inputs, J):
        J['y', 'x'] = 2 * inputs[x]

prob = om.Problem()
model = prob.model

model.add_subsystem('p1', om.IndepVarComp('x', 1.0))
model.add_subsystem('p2', om.IndepVarComp('x', 1.0))

parallel = model.add_subsystem('parallel', om.ParallelGroup())
parallel.add_subsystem('c1', TestComp())
parallel.add_subsystem('c2', TestComp2())

model.add_subsystem('c3', om.ExecComp(['y=3.0*x1+7.0*x2']))

model.connect("parallel.c1.y", "c3.x1")
model.connect("parallel.c2.y", "c3.x2")

model.connect("p1.x", "parallel.c1.x")
model.connect("p2.x", "parallel.c2.x")

prob.setup(check=False, mode='fwd')
prob.run_model()

# fails if n procs > 3
prob.model.list_outputs()
# also will fail if n procs > 3
# prob.model.list_inputs