def setUp(self):
self.nn = 10
self.p = om.Problem()
ivc = om.IndepVarComp()
ivc.add_output(name='a', shape=(self.nn,))
ivc.add_output(name='b', shape=(self.nn,))
self.p.model.add_subsystem(name='ivc',
subsys=ivc,
promotes_outputs=['a', 'b'])
demux_comp = self.p.model.add_subsystem(name='demux_comp',
subsys=om.DemuxComp(vec_size=self.nn))
demux_comp.add_var('a', shape=(self.nn,))
demux_comp.add_var('b', shape=(self.nn,))
self.p.model.connect('a', 'demux_comp.a')
self.p.model.connect('b', 'demux_comp.b')
self.p.setup(force_alloc_complex=True)
self.p['a'] = np.random.rand(self.nn)
self.p['b'] = np.random.rand(self.nn)
self.p.run_model()
def test_axis_with_wrong_size(self):
nn = 10
p = om.Problem()
ivc = om.IndepVarComp()
ivc.add_output(name='a', shape=(nn, 7))
ivc.add_output(name='b', shape=(3, nn))
p.model.add_subsystem(name='ivc',
subsys=ivc,
promotes_outputs=['a', 'b'])
demux_comp = p.model.add_subsystem(name='demux_comp',
subsys=om.DemuxComp(vec_size=nn))
demux_comp.add_var('a', shape=(nn, 7), axis=1)
demux_comp.add_var('b', shape=(3, nn), axis=1)
p.model.connect('a', 'demux_comp.a')
p.model.connect('b', 'demux_comp.b')
with self.assertRaises(RuntimeError) as ctx:
p.setup()
self.assertEqual(
str(ctx.exception),
"Variable 'a' cannot be demuxed along axis 1. Axis size is "
"7 but vec_size is 10.")
def test(self):
"""
An example demonstrating a trivial use case of DemuxComp
"""
import numpy as np
import openmdao.api as om
from openmdao.utils.assert_utils import assert_rel_error
# The number of elements to be demuxed
n = 3
# The size of each element to be demuxed
m = 100
p = om.Problem()
ivc = om.IndepVarComp()
ivc.add_output(name='pos_ecef', shape=(m, 3), units='km')
p.model.add_subsystem(name='ivc',
subsys=ivc,
promotes_outputs=['pos_ecef'])
mux_comp = p.model.add_subsystem(name='demux',
subsys=om.DemuxComp(vec_size=n))
mux_comp.add_var('pos', shape=(m, n), axis=1, units='km')
p.model.add_subsystem(name='longitude_comp',
subsys=om.ExecComp('long = atan(y/x)',
x={
'value': np.ones(m),
'units': 'km'
},
y={
'value': np.ones(m),
'units': 'km'
},
long={
'value': np.ones(m),
'units': 'rad'
}))
p.model.connect('demux.pos_0', 'longitude_comp.x')
p.model.connect('demux.pos_1', 'longitude_comp.y')
p.model.connect('pos_ecef', 'demux.pos')
p.setup()
p['pos_ecef'][:, 0] = 6378 * np.cos(np.linspace(0, 2 * np.pi, m))
p['pos_ecef'][:, 1] = 6378 * np.sin(np.linspace(0, 2 * np.pi, m))
p['pos_ecef'][:, 2] = 0.0
p.run_model()
expected = np.arctan(p['pos_ecef'][:, 1] / p['pos_ecef'][:, 0])
assert_rel_error(self, p.get_val('longitude_comp.long'), expected)
def test(self):
"""
An example demonstrating a trivial use case of DemuxComp
"""
import numpy as np
import openmdao.api as om
# The number of elements to be demuxed
n = 3
# The size of each element to be demuxed
m = 100
p = om.Problem()
demux_comp = p.model.add_subsystem(name='demux',
subsys=om.DemuxComp(vec_size=n),
promotes_inputs=['pos_ecef'])
demux_comp.add_var('pos_ecef', shape=(m, n), axis=1, units='km')
p.model.add_subsystem(name='longitude_comp',
subsys=om.ExecComp('long = atan(y/x)',
x={
'value': np.ones(m),
'units': 'km'
},
y={
'value': np.ones(m),
'units': 'km'
},
long={
'value': np.ones(m),
'units': 'rad'
}))
p.model.connect('demux.pos_ecef_0', 'longitude_comp.x')
p.model.connect('demux.pos_ecef_1', 'longitude_comp.y')
p.setup()
p.set_val('pos_ecef',
6378 * np.cos(np.linspace(0, 2 * np.pi, m)),
indices=om.slicer[:, 0])
p.set_val('pos_ecef',
6378 * np.sin(np.linspace(0, 2 * np.pi, m)),
indices=om.slicer[:, 1])
p.set_val('pos_ecef', 0.0, indices=om.slicer[:, 2])
p.run_model()
expected = np.arctan(
p.get_val('pos_ecef', indices=om.slicer[:, 1]) /
p.get_val('pos_ecef', indices=om.slicer[:, 0]))
assert_near_equal(p.get_val('longitude_comp.long'), expected)
def test_invalid_axis(self):
nn = 10
p = om.Problem()
ivc = om.IndepVarComp()
ivc.add_output(name='a', shape=(nn, ))
p.model.add_subsystem(name='ivc', subsys=ivc, promotes_outputs=['a'])
demux_comp = p.model.add_subsystem(name='demux_comp',
subsys=om.DemuxComp(vec_size=nn))
with self.assertRaises(RuntimeError) as ctx:
demux_comp.add_var('a', shape=(nn, ), axis=1)
self.assertEqual(
str(ctx.exception),
"DemuxComp (demux_comp): Invalid axis (1) for variable 'a' of shape (10,)"
)
