Source code for openmdao.jacobians.assembled_jacobian

"""Define the AssembledJacobian class."""
import sys
from collections import defaultdict, OrderedDict

import numpy as np

from openmdao.jacobians.jacobian import Jacobian
from openmdao.matrices.dense_matrix import DenseMatrix
from openmdao.matrices.coo_matrix import COOMatrix
from openmdao.matrices.csr_matrix import CSRMatrix
from openmdao.matrices.csc_matrix import CSCMatrix
from openmdao.utils.units import unit_conversion

_empty_dict = {}


[docs]class AssembledJacobian(Jacobian): """ Assemble a global <Jacobian>. Parameters ---------- matrix_class : type Class to use to create internal matrices. system : System Parent system to this jacobian. Attributes ---------- _view_ranges : dict Maps system pathnames to jacobian sub-view ranges _int_mtx : <Matrix> Global internal Jacobian. _ext_mtx : {str: <Matrix>, ...} External Jacobian for each viewing subsystem. _mask_caches : dict Contains masking arrays for when a subset of the variables are present in a vector, keyed by the input._names set. _matrix_class : type Class used to create Matrix objects. _subjac_iters : dict Mapping of system pathname to tuple of lists of absolute key tuples used to index into the jacobian. _in_ranges : dict Column ranges for inputs. _out_ranges : dict Row ranges for outputs. """
[docs] def __init__(self, matrix_class, system): """ Initialize all attributes. """ global Component # avoid circular imports from openmdao.core.component import Component super().__init__(system) self._view_ranges = {} self._int_mtx = None self._ext_mtx = {} self._mask_caches = {} self._matrix_class = matrix_class self._out_ranges = self._get_ranges(system, 'output') self._in_ranges = self._get_ranges(system, 'input') self._subjac_iters = defaultdict(lambda: None)
def _get_ranges(self, system, vtype): """ Return an ordered dict of ranges for each var of a particular type (input or output). Parameters ---------- system : System System owning this jacobian. vtype : str Type of variable, must be one of ('input', 'output'). Returns ------- OrderedDict Tuples of the form (start, end) keyed on variable name. """ ranges = OrderedDict() start = end = 0 for name, meta in system._var_abs2meta[vtype].items(): end += meta['size'] ranges[name] = (start, end) start = end return ranges def _initialize(self, system): """ Allocate the global matrices. Parameters ---------- system : System Parent system to this jacobian. """ # var_indices are the *global* indices for variables on this proc is_top = system.pathname == '' abs2meta_in = system._var_abs2meta['input'] all_meta = system._var_allprocs_abs2meta self._int_mtx = int_mtx = self._matrix_class(system.comm, True) ext_mtx = self._matrix_class(system.comm, False) iproc = system.comm.rank out_ranges = self._out_ranges in_ranges = self._in_ranges abs2prom_out = system._var_abs2prom['output'] conns = {} if isinstance(system, Component) else system._conn_global_abs_in2out abs_key2shape = self._abs_key2shape # create the matrix subjacs for abs_key, info in self._subjacs_info.items(): res_abs_name, wrt_abs_name = abs_key # because self._subjacs_info is shared among all 'related' assembled jacs, # we use out_ranges (and later in_ranges) to weed out keys outside of this jac if res_abs_name not in out_ranges: continue res_offset, res_end = out_ranges[res_abs_name] res_size = res_end - res_offset if wrt_abs_name in abs2prom_out: out_offset, out_end = out_ranges[wrt_abs_name] out_size = out_end - out_offset shape = (res_size, out_size) int_mtx._add_submat(abs_key, info, res_offset, out_offset, None, shape) elif wrt_abs_name in in_ranges: if wrt_abs_name in conns: # connected input out_abs_name = conns[wrt_abs_name] if out_abs_name not in out_ranges: continue meta_in = abs2meta_in[wrt_abs_name] all_out_meta = all_meta['output'][out_abs_name] # calculate unit conversion in_units = meta_in['units'] out_units = all_out_meta['units'] if in_units and out_units and in_units != out_units: factor, _ = unit_conversion(out_units, in_units) if factor == 1.0: factor = None else: factor = None out_offset, out_end = out_ranges[out_abs_name] out_size = out_end - out_offset shape = (res_size, out_size) src_indices = abs2meta_in[wrt_abs_name]['src_indices'] if src_indices is not None: # need to add an entry for d(output)/d(source) # instead of d(output)/d(input). int_mtx is a square matrix whose # rows and columns map to output/resid vars only. abs_key2 = (res_abs_name, out_abs_name) shape = abs_key2shape(abs_key2) int_mtx._add_submat(abs_key, info, res_offset, out_offset, src_indices, shape, factor) elif not is_top: # input is connected to something outside current system in_offset, in_end = in_ranges[wrt_abs_name] # don't use global offsets for ext_mtx res_offset, res_end = out_ranges[res_abs_name] res_size = res_end - res_offset shape = (res_size, in_end - in_offset) ext_mtx._add_submat(abs_key, info, res_offset, in_offset, None, shape) out_size = len(system._outputs) int_mtx._build(out_size, out_size, system) if ext_mtx._submats: ext_mtx._build(out_size, len(system._vectors['input']['linear'])) else: ext_mtx = None self._ext_mtx[system.pathname] = ext_mtx def _init_ranges(self, system): in_ranges = self._in_ranges out_ranges = self._out_ranges input_names = list(system._var_abs2meta['input']) if input_names: min_in_offset = in_ranges[input_names[0]][0] max_in_offset = in_ranges[input_names[-1]][1] else: min_in_offset = sys.maxsize max_in_offset = 0 output_names = list(system._var_abs2meta['output']) if output_names: min_res_offset = out_ranges[output_names[0]][0] max_res_offset = out_ranges[output_names[-1]][1] else: min_res_offset = sys.maxsize max_res_offset = 0 self._view_ranges[system.pathname] = (min_res_offset, max_res_offset, min_in_offset, max_in_offset) def _init_view(self, system): """ Determine the _ext_mtx for a sub-view of the assembled jacobian. Parameters ---------- system : <System> The system being solved using a sub-view of the jacobian. """ abs2meta = system._var_abs2meta['output'] ranges = self._view_ranges[system.pathname] ext_mtx = self._matrix_class(system.comm, False) conns = {} if isinstance(system, Component) else system._conn_global_abs_in2out iproc = system.comm.rank sizes = system._var_sizes['input'] abs2idx = system._var_allprocs_abs2idx in_offset = {n: np.sum(sizes[iproc, :abs2idx[n]]) for n in system._var_abs2meta['input'] if n not in conns} subjacs_info = self._subjacs_info sizes = system._var_sizes['output'] for s in system.system_iter(recurse=True, include_self=True, typ=Component): for res_abs_name, res_meta in s._var_abs2meta['output'].items(): res_offset = np.sum(sizes[iproc, :abs2idx[res_abs_name]]) res_size = res_meta['size'] for in_abs_name in s._var_abs2meta['input']: if in_abs_name not in conns: # unconnected input abs_key = (res_abs_name, in_abs_name) if abs_key not in subjacs_info: continue info = subjacs_info[abs_key] ext_mtx._add_submat(abs_key, info, res_offset - ranges[0], in_offset[in_abs_name] - ranges[2], None, info['shape']) if ext_mtx._submats: ext_mtx._build(len(system._vectors['output']['linear']), len(system._vectors['input']['linear'])) else: ext_mtx = None self._ext_mtx[system.pathname] = ext_mtx def _get_subjac_iters(self, system): # this determines the subjacs that get updated during _update() global _empty_dict subjac_iters = self._subjac_iters[system.pathname] if subjac_iters is None: int_mtx = self._int_mtx ext_mtx = self._ext_mtx[system.pathname] subjacs = system._subjacs_info sys_inputs = system._var_allprocs_abs2prom['input'] sys_outputs = system._var_allprocs_abs2prom['output'] if isinstance(system, Component): global_conns = _empty_dict else: global_conns = system._conn_global_abs_in2out output_names = set(system._var_abs2meta['output']) input_names = set(system._var_abs2meta['input']) rev_conns = defaultdict(list) for tgt, src in global_conns.items(): rev_conns[src].append(tgt) # This is the level where the AssembledJacobian is slotted. # The of and wrt are the inputs and outputs that it sees, if they are in the subjacs. # TODO - For top level FD, the subjacs might not contain all derivs. iters = [] iters_in_ext = [] for abs_key in subjacs: _, wrtname = abs_key if wrtname in sys_outputs: if wrtname in output_names: if abs_key in int_mtx._submats: iters.append(abs_key) else: # This happens when the src is an indepvarcomp that is # contained in the system. of, wrt = abs_key if wrt in rev_conns: for tgt in rev_conns[wrt]: if (of, tgt) in int_mtx._submats: iters.append(abs_key) break elif wrtname in sys_inputs: if wrtname in input_names: # wrt is an input if wrtname in global_conns: iters.append(abs_key) elif ext_mtx is not None: iters_in_ext.append(abs_key) elif ext_mtx is not None and wrtname in sys_inputs: iters_in_ext.append(abs_key) self._subjac_iters[system.pathname] = subjac_iters = (iters, iters_in_ext) return subjac_iters def _update(self, system): """ Read the user's sub-Jacobians and set into the global matrix. Parameters ---------- system : System System that is updating this jacobian. """ # _initialize has been delayed until the first _update call if self._int_mtx is None: self._initialize(system) self._init_ranges(system) if system.pathname: self._init_view(system) int_mtx = self._int_mtx ext_mtx = self._ext_mtx[system.pathname] subjacs = system._subjacs_info iters, iters_in_ext = self._get_subjac_iters(system) int_mtx._pre_update() if ext_mtx is not None: ext_mtx._pre_update() if self._randomize: for key in iters: int_mtx._update_submat(key, self._randomize_subjac(subjacs[key]['val'], key)) for key in iters_in_ext: ext_mtx._update_submat(key, self._randomize_subjac(subjacs[key]['val'], key)) else: for key in iters: int_mtx._update_submat(key, subjacs[key]['val']) for key in iters_in_ext: ext_mtx._update_submat(key, subjacs[key]['val']) int_mtx._post_update() if ext_mtx is not None: ext_mtx._post_update() if self._under_complex_step: # If we create a new _int_mtx while under complex step, we need to convert it to a # complex data type. self._int_mtx.set_complex_step_mode(True) def _apply(self, system, d_inputs, d_outputs, d_residuals, mode): """ Compute matrix-vector product. Parameters ---------- system : System System that is updating this jacobian. d_inputs : Vector inputs linear vector. d_outputs : Vector outputs linear vector. d_residuals : Vector residuals linear vector. mode : str 'fwd' or 'rev'. """ int_mtx = self._int_mtx ext_mtx = self._ext_mtx[system.pathname] if ext_mtx is None and not d_outputs._names: # avoid unnecessary unscaling return with system._unscaled_context(outputs=[d_outputs], residuals=[d_residuals]): do_mask = ext_mtx is not None and d_inputs._names if do_mask: # Masking try: mask = self._mask_caches[(d_inputs._names, mode)] except KeyError: mask = ext_mtx._create_mask_cache(d_inputs) self._mask_caches[(d_inputs._names, mode)] = mask dresids = d_residuals.asarray() if mode == 'fwd': if d_outputs._names: dresids += int_mtx._prod(d_outputs.asarray(), mode) if do_mask: dresids += ext_mtx._prod(d_inputs.asarray(), mode, mask=mask) else: # rev if d_outputs._names: d_outputs += int_mtx._prod(dresids, mode) if do_mask: d_inputs += ext_mtx._prod(dresids, mode, mask=mask)
[docs] def set_complex_step_mode(self, active): """ Turn on or off complex stepping mode. When turned on, the value in each subjac is cast as complex, and when turned off, they are returned to real values. Parameters ---------- active : bool Complex mode flag; set to True prior to commencing complex step. """ super().set_complex_step_mode(active) if self._int_mtx is not None: self._int_mtx.set_complex_step_mode(active) for mtx in self._ext_mtx.values(): if mtx: mtx.set_complex_step_mode(active)
[docs]class DenseJacobian(AssembledJacobian): """ Assemble dense global <Jacobian>. Parameters ---------- system : System Parent system to this jacobian. """
[docs] def __init__(self, system): """ Initialize all attributes. """ super().__init__(DenseMatrix, system=system)
[docs]class COOJacobian(AssembledJacobian): """ Assemble sparse global <Jacobian> in Coordinate list format. Parameters ---------- system : System Parent system to this jacobian. """
[docs] def __init__(self, system): """ Initialize all attributes. """ super().__init__(COOMatrix, system=system)
[docs]class CSRJacobian(AssembledJacobian): """ Assemble sparse global <Jacobian> in Compressed Row Storage format. Parameters ---------- system : System Parent system to this jacobian. """
[docs] def __init__(self, system): """ Initialize all attributes. """ super().__init__(CSRMatrix, system=system)
[docs]class CSCJacobian(AssembledJacobian): """ Assemble sparse global <Jacobian> in Compressed Col Storage format. Parameters ---------- system : System Parent system to this jacobian. """
[docs] def __init__(self, system): """ Initialize all attributes. """ super().__init__(CSCMatrix, system=system)