Source code for openmdao.components.cross_product_comp

"""Definition of the Cross Product Component."""

import numpy as np

from openmdao.core.explicitcomponent import ExplicitComponent


[docs]class CrossProductComp(ExplicitComponent): """ Compute a vectorized cross product. math:: c = np.cross(a, b) where a is of shape (vec_size, 3) b is of shape (vec_size, 3) c is of shape (vec_size, 3) if vec_size > 1 and where a is of shape (3,) b is of shape (3,) c is of shape (3,) otherwise. Parameters ---------- **kwargs : dict of keyword arguments Keyword arguments that will be mapped into the Component options. Attributes ---------- _products : list Cache the data provided during `add_product` so everything can be saved until setup is called. """
[docs] def __init__(self, **kwargs): """ Initialize the Cross Product component. """ super().__init__(**kwargs) self._products = [] opt = self.options self.add_product(c_name=opt['c_name'], a_name=opt['a_name'], b_name=opt['b_name'], c_units=opt['c_units'], a_units=opt['a_units'], b_units=opt['b_units'], vec_size=opt['vec_size']) self._no_check_partials = True
[docs] def initialize(self): """ Declare options. """ self.options.declare('vec_size', types=int, default=1, desc='The number of points at which the cross product is computed') self.options.declare('a_name', types=str, default='a', desc='The variable name for vector a.') self.options.declare('b_name', types=str, default='b', desc='The variable name for vector b.') self.options.declare('c_name', types=str, default='c', desc='The variable name for vector c.') self.options.declare('a_units', types=str, default=None, allow_none=True, desc='The units for vector a.') self.options.declare('b_units', types=str, default=None, allow_none=True, desc='The units for vector b.') self.options.declare('c_units', types=str, default=None, allow_none=True, desc='The units for vector c.') self._k = np.array([[0, 0, 0, -1, 0, 1], [0, 1, 0, 0, -1, 0], [-1, 0, 1, 0, 0, 0]], dtype=float) self._minus_k = -self._k
[docs] def add_product(self, c_name, a_name='a', b_name='b', c_units=None, a_units=None, b_units=None, vec_size=1): """ Add a new output product to the cross product component. Parameters ---------- c_name : str The name of the vector product output. a_name : str The name of the first vector input. b_name : str The name of the second vector input. c_units : str or None The units of the output. a_units : str or None The units of input a. b_units : str or None The units of input b. vec_size : int The number of points at which the dot vector product should be computed simultaneously. The shape of the output is (vec_size,). """ self._products.append({ 'a_name': a_name, 'b_name': b_name, 'c_name': c_name, 'a_units': a_units, 'b_units': b_units, 'c_units': c_units, 'vec_size': vec_size, }) # add inputs and outputs for all products if self._static_mode: var_rel2meta = self._static_var_rel2meta var_rel_names = self._static_var_rel_names else: var_rel2meta = self._var_rel2meta var_rel_names = self._var_rel_names shape = (vec_size, 3) if vec_size > 1 else (3,) if c_name not in var_rel2meta: self.add_output(name=c_name, val=np.ones(shape=shape), units=c_units) elif c_name in var_rel_names['input']: raise NameError(f"{self.msginfo}: '{c_name}' specified as an output, " "but it has already been defined as an input.") else: raise NameError(f"{self.msginfo}: Multiple definition of output '{c_name}'.") if a_name not in var_rel2meta: self.add_input(name=a_name, shape=shape, units=a_units) elif a_name in var_rel_names['output']: raise NameError(f"{self.msginfo}: '{a_name}' specified as an input, " "but it has already been defined as an output.") else: meta = var_rel2meta[a_name] if a_units != meta['units']: raise ValueError(f"{self.msginfo}: Conflicting units '{a_units}' specified " f"for input '{a_name}', which has already been defined " f"with units '{meta['units']}'.") meta_shape = meta['shape'] if shape != meta_shape: raise ValueError(f"{self.msginfo}: Conflicting vec_size={vec_size} specified " f"for input '{a_name}', which has already been defined with " f"vec_size={meta_shape[0] if len(meta_shape) > 1 else 1}.") if b_name not in var_rel2meta: self.add_input(name=b_name, shape=shape, units=b_units) elif b_name in var_rel_names['output']: raise NameError(f"{self.msginfo}: '{b_name}' specified as an input, " "but it has already been defined as an output.") else: meta = var_rel2meta[b_name] if b_units != meta['units']: raise ValueError(f"{self.msginfo}: Conflicting units '{b_units}' specified " f"for input '{b_name}', which has already been defined " f"with units '{meta['units']}'.") meta_shape = meta['shape'] if shape != meta_shape: raise ValueError(f"{self.msginfo}: Conflicting vec_size={vec_size} specified " f"for input '{b_name}', which has already been defined with " f"vec_size={meta_shape[0] if len(meta_shape) > 1 else 1}.") row_idxs = np.repeat(np.arange(vec_size * 3, dtype=int), 2) col_idxs = np.empty((0,), dtype=int) M = np.array([1, 2, 0, 2, 0, 1], dtype=int) for i in range(vec_size): col_idxs = np.concatenate((col_idxs, M + i * 3)) self.declare_partials(of=c_name, wrt=a_name, rows=row_idxs, cols=col_idxs) self.declare_partials(of=c_name, wrt=b_name, rows=row_idxs, cols=col_idxs)
[docs] def compute(self, inputs, outputs): """ Compute the cross product of inputs `a` and `b` using np.cross. Parameters ---------- inputs : Vector Unscaled, dimensional input variables read via inputs[key]. outputs : Vector Unscaled, dimensional output variables read via outputs[key]. """ for product in self._products: a = inputs[product['a_name']] b = inputs[product['b_name']] outputs[product['c_name']] = np.cross(a, b)
[docs] def compute_partials(self, inputs, partials): """ Compute the sparse partials for the cross product w.r.t. the inputs. Parameters ---------- inputs : Vector Unscaled, dimensional input variables read via inputs[key]. partials : Jacobian Sub-jac components written to partials[output_name, input_name]. """ for product in self._products: a = inputs[product['a_name']] b = inputs[product['b_name']] # Use the following for sparse partials partials[product['c_name'], product['a_name']] = \ np.einsum('...j,ji->...i', b, self._minus_k).ravel() partials[product['c_name'], product['b_name']] = \ np.einsum('...j,ji->...i', a, self._k).ravel()