Search code examples
openmdao

Issue using Demux Comp with 3 dimentional arrray (OpenMDAO)


I'm trying to demux an array with 3 axis as seen in the following sample code

import openmdao.api as om

p = om.Problem()

ivc = p.model.add_subsystem('idv', om.IndepVarComp(), promotes=['*'])
ivc.add_output(name='f0', shape=(2,3), units='m', val=[[111,112,113], [121,122,123]])
ivc.add_output(name='f1', shape=(2,3), units='m', val=[[211,212,213], [221,222,223]])
ivc.add_output(name='f2', shape=(2,3), units='m', val=[[311,312,313], [321,322,323]])
ivc.add_output(name='f3', shape=(2,3), units='m', val=[[411,412,413], [421,422,423]])

mux_comp = p.model.add_subsystem(name='mux', subsys=om.MuxComp(vec_size=4))
mux_comp.add_var('r', shape=(2,3), axis=0, units='m')

demux_comp = p.model.add_subsystem(name='demux', subsys=om.DemuxComp(vec_size=4))
demux_comp.add_var('g', shape=(4,2,3), axis=0, units='m')

p.model.connect('f0', 'mux.r_0')
p.model.connect('f1', 'mux.r_1')
p.model.connect('f2', 'mux.r_2')
p.model.connect('f3', 'mux.r_3')

p.model.connect('mux.r', 'demux.g')

p.setup()
p.run_model()

print(p['mux.r'])
print(p['mux.r'].shape)

print(p['demux.g_0'])
print(p['demux.g_1'])
print(p['demux.g_2'])
print(p['demux.g_3'])
print(p['demux.g_0'].shape)

When this runs, I get the following error

RuntimeError: DemuxComp (demux): d(g_0)/d(g): declare_partials has been called with rows and cols, which should be arrays of equal length, but rows is length 6 while cols is length 2.

As the error only pertained to the partials, I took a look into the demux_comp.py code in the OpenMDAO library and modified the declare_partials line from

self.declare_partials(of=out_name, wrt=var, rows=rs, cols=cs, val=1.0)

to

self.declare_partials(of=out_name, wrt=var, val=1.0)

This allowed the code to run successfully and output the proper demuxed variables. Will this have any adverse effects on the rest of my code and optimizations?


Solution

  • It looks like you've uncovered a bug. We'll get on fixing this.

    While your fix will allow it to run, the partials will be incorrect. In the mean time you'd be better off replacing the declare partials call with instructions to use finite difference or complex step.

    self.declare_partials(of=out_name, wrt=var, method='cs')
    

    For complex-step, or method='fd' for finite difference.