Source code for openmdao.components.mux_comp
"""Definition of the Mux Component."""
import numpy as np
from openmdao.core.explicitcomponent import ExplicitComponent
from openmdao.utils.array_utils import shape_to_len
[docs]class MuxComp(ExplicitComponent):
"""
Mux one or more inputs along a given axis.
Parameters
----------
**kwargs : dict
Arguments to be passed to the component initialization method.
Attributes
----------
_vars : dict
Container mapping name of variables to be muxed with additional data.
_input_names : dict
Container mapping name of variables to be muxed with associated inputs.
"""
[docs] def __init__(self, **kwargs):
"""
Instantiate MuxComp and populate private members.
"""
super().__init__(**kwargs)
self._vars = {}
self._input_names = {}
self._no_check_partials = True
[docs] def initialize(self):
"""
Declare options.
"""
self.options.declare('vec_size', types=int, default=2,
desc='The number of elements to be combined into an output.')
[docs] def add_var(self, name, val=1.0, shape=None, units=None, desc='', axis=0):
"""
Add an output variable to be muxed, and all associated input variables.
Parameters
----------
name : str
Name of the variable in this component's namespace.
val : float or list or tuple or ndarray or Iterable
The initial value of the variable being added in user-defined units.
Default is 1.0.
shape : int or tuple or list or None
Shape of the input variables to be muxed, only required if val is not an array.
Default is None.
units : str or None
Units in which this input variable will be provided to the component
during execution. Default is None, which means it is unitless.
desc : str
Description of the variable.
axis : int
The axis along which the elements will be stacked. Note that N-dimensional inputs
cannot be stacked along an axis greater than N.
"""
self._vars[name] = {'val': val, 'shape': shape, 'units': units, 'desc': desc, 'axis': axis}
opts = self.options
vec_size = opts['vec_size']
options = self._vars[name]
kwgs = dict(options)
in_shape = np.asarray(options['val']).shape \
if options['shape'] is None else options['shape']
in_size = shape_to_len(in_shape)
out_shape = list(in_shape)
out_shape.insert(options['axis'], vec_size)
kwgs.pop('shape')
ax = kwgs.pop('axis')
in_dimension = len(in_shape)
if ax > in_dimension:
raise ValueError('{3}: Cannot mux a {0}D inputs for {2} along axis greater '
'than {0} ({1})'.format(in_dimension, ax, name, self.msginfo))
self.add_output(name=name,
val=options['val'],
shape=out_shape,
units=options['units'],
desc=options['desc'])
self._input_names[name] = []
for i in range(vec_size):
in_name = '{0}_{1}'.format(name, i)
self._input_names[name].append(in_name)
self.add_input(name=in_name, shape=in_shape, **kwgs)
in_templates = [np.zeros(in_shape, dtype=int) for _ in range(vec_size)]
rs = []
cs = []
for j in range(in_size):
in_templates[i].flat[:] = 0
in_templates[i].flat[j] = 1
temp_out = np.stack(in_templates, axis=ax)
cs.append(j)
rs.append(int(np.nonzero(temp_out.ravel())[0].item()))
self.declare_partials(of=name, wrt=in_name, rows=rs, cols=cs, val=1.0)
[docs] def compute(self, inputs, outputs):
"""
Mux the inputs into the appropriate outputs.
Parameters
----------
inputs : Vector
Unscaled, dimensional input variables read via inputs[key].
outputs : Vector
Unscaled, dimensional output variables read via outputs[key].
"""
opts = self.options
vec_size = opts['vec_size']
for var in self._vars:
ax = self._vars[var]['axis']
invar = self._input_names[var]
vals = [inputs[invar[i]] for i in range(vec_size)]
outputs[var][...] = np.stack(vals, axis=ax)