group fd bug

gistfile1.txt

import numpy as np

import openmdao.api as om


class DistParab(om.ExplicitComponent):

    def initialize(self):

        self.options.declare('arr_size', types=int, default=10,
                             desc="Size of input and output vectors.")

    def setup(self):
        arr_size = self.options['arr_size']
        comm = self.comm
        rank = comm.rank

        self.add_input('x', val=np.ones(arr_size))
        self.add_input('y', val=np.ones(arr_size))

        self.add_output('f_xy', val=np.ones(arr_size))

        self.declare_partials('f_xy', ['x', 'y'])

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

        outputs['f_xy'] = x**2 + x * y + (y + 4.0)**2 - 3.0

    def compute_partials(self, inputs, partials):
        x = inputs['x']
        y = inputs['y']

        partials['f_xy', 'x'] = np.diag(2.0*x + y)
        partials['f_xy', 'y'] = np.diag(2.0*y + 8.0 + x)


class NonDistComp(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('arr_size', types=int, default=10,
                             desc="Size of input and output vectors.")

    def setup(self):
        arr_size = self.options['arr_size']

        self.add_input('f_xy', val=np.ones(arr_size))

        self.add_output('g', val=np.ones(arr_size))

        self.mat = np.array([3.5, -1, 5])

        row_col = np.arange(arr_size)
        self.declare_partials('g', ['f_xy'], rows=row_col, cols=row_col, val=self.mat.copy())

    def compute(self, inputs, outputs):
        x = inputs['f_xy']

        outputs['g'] = x * self.mat


size = 3
use_sub = True

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

ivc = om.IndepVarComp()
ivc.add_output('x', np.ones((size, )))

model.add_subsystem('p', ivc, promotes=['*'])
if use_sub:
    sub = model.add_subsystem('sub', om.Group(), promotes=['*'])
else:
    sub = model

sub.add_subsystem("parab", DistParab(arr_size=size), promotes=['*'])
sub.add_subsystem("ndp", NonDistComp(arr_size=size), promotes=['*'])

model.add_design_var('x', lower=-50.0, upper=50.0)
model.add_constraint('g', lower=0.0)

sub.approx_totals(method='fd')

prob.setup(force_alloc_complex=True)


prob.run_model()

J = prob.check_totals(method='fd')

of = ['sub.ndp.g']
J = prob.driver._compute_totals(of=of, wrt=['p.x'], return_format='dict')
print(J)
print('done')