openmdao

OpenMDAO 3.x: relevance graph for ParallelGroup approximated totals

For the below problem, I would like to improve the efficiency of the gradient computation when 1) approximating totals with FD and when using matrix free forward mode. The ParallelGroup components all get executed when perturbing/seeding individual indices of the vector input indeps.x_vec, which adds unnecessary overhead.

If I remove the indepvar vector and instead declare each 'par.case%02d.x' % i as design variables, the matrix free derivative approach correctly only calls compute_jacvec_product once for each component, but this is not the case when approximating totals.

So the questions are:

  1. Is there a way to get the vector input fan-out model to work correctly for forward matrix-free derivatives?
  2. Can relevance reduction be applied when computing approximated totals?

For the actual code we have, the MultComp is a group with a nonlinear solver with two disciplines with large I/O vectors being exchanged.

import numpy as np
import openmdao.api as om


class MultComp(om.ExplicitComponent):
    """
    """

    def __init__(self, delay=1.0, size=3, mult=2.0):
        super().__init__()
        self.delay = delay
        self.size = size
        self.mult = mult

    def setup(self):
        self.add_input('x', val=self.mult)
        self.add_output('y', val=0.0)

        self.declare_partials(of='*', wrt='*')

    def compute(self, inputs, outputs):

        outputs['y'] = inputs['x'] * self.mult
        print('hello from comp', self.mult, inputs['x'], outputs['y'])

    def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):

        d_outputs['y'] += d_inputs['x'] * self.mult
        print('hello from jacvec comp', self.mult, inputs['x'])

    # def compute_partials(self, inputs, partials):

    #     partials[('y', 'x')] = self.mult
    #     print('hello from partials comp', self.mult, inputs['x'])

p = om.Problem(group_by_pre_opt_post=True)

indeps = p.model.add_subsystem('indeps', om.IndepVarComp())
indeps.add_output('x_vec', np.linspace(1., 4., 4))

par = p.model.add_subsystem('par', om.ParallelGroup())

mux = p.model.add_subsystem('mux_comp', om.MuxComp(vec_size=4))
mux.add_var("y_vec", shape=(1,))

for i in range(4):
    par.add_subsystem('case%02d' % i, MultComp(mult=float(i+1)))
    p.model.connect('indeps.x_vec', 'par.case%02d.x' % i, src_indices=[i])
    p.model.connect('par.case%02d.y' % i, 'mux_comp.y_vec_%i' % i)

p.model.add_subsystem('sum_comp', om.ExecComp('ysum=sum(y_vec)',
                                              y_vec=np.zeros(4)))
p.model.connect('mux_comp.y_vec', 'sum_comp.y_vec')
p.model.add_design_var('indeps.x_vec', upper=20.)
p.model.add_objective('sum_comp.ysum', scaler=-1.)

p.driver = om.ScipyOptimizeDriver()
p.driver.options["optimizer"] = "SLSQP"
p.model.par.approx_totals(method='fd', step=1e-4, form='central')

p.setup(mode='fwd')

p.run_model()

totals = p.compute_totals()
Solution

  • The main issue here is that currently, OpenMDAO computes relevance based on whole variables rather than on individual entries of variables that are arrays. This makes OpenMDAO treat all of the components in the ParallelGroup as being relevant to the indeps.x_vec design variable, so if any part of indeps.x_vec changes, OpenMDAO runs all of those components.

    So in answer to your question 1, unless we add knowledge of individual array entries to OpenMDAO's relevance (which we're not opposed to as long as we can figure out some efficient way to do it), it won't work the way you'd like without splitting up the single array design variable into multiple separate variables.

    In answer to question 2, relevance does work for approx totals if they're done at the top level. In your example you do it down at the ParallelGroup level and OpenMDAO currently doesn't do approx relevance at that level.