-
Notifications
You must be signed in to change notification settings - Fork 240
/
mux_comp.py
136 lines (107 loc) · 4.43 KB
/
mux_comp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
"""Definition of the Mux Component."""
import numpy as np
from openmdao.core.explicitcomponent import ExplicitComponent
from openmdao.utils.array_utils import shape_to_len
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.
"""
def __init__(self, **kwargs):
"""
Instantiate MuxComp and populate private members.
"""
super().__init__(**kwargs)
self._vars = {}
self._input_names = {}
self._no_check_partials = True
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.')
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]))
self.declare_partials(of=name, wrt=in_name, rows=rs, cols=cs, val=1.0)
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)