Source code for openmdao.jacobians.jacobian

"""Define the base Jacobian class."""
import weakref
import sys

import numpy as np
from numpy.random import rand

from collections import OrderedDict, defaultdict
from scipy.sparse import issparse, coo_matrix

from openmdao.core.constants import INT_DTYPE
from openmdao.utils.name_maps import key2abs_key, rel_name2abs_name
from openmdao.utils.array_utils import sparse_subinds
from openmdao.matrices.matrix import sparse_types
from openmdao.vectors.vector import _full_slice

SUBJAC_META_DEFAULTS = {
    'rows': None,
    'cols': None,
    'val': None,
    'dependent': False,
}


[docs]class Jacobian(object): """ Base Jacobian class. This class provides a dictionary interface for sub-Jacobians and performs matrix-vector products when apply_linear is called. Parameters ---------- system : System Parent system to this jacobian. Attributes ---------- _system : <System> Pointer to the system that is currently operating on this Jacobian. _subjacs_info : dict Dictionary of the sub-Jacobian metadata keyed by absolute names. _under_complex_step : bool When True, this Jacobian is under complex step, using a complex jacobian. _abs_keys : defaultdict A cache dict for key to absolute key. _randomize : bool If True, sparsity is being computed for simultaneous derivative coloring. _col_var_info : dict Maps column name to start, end, and slice/indices into the result array. _colnames : list List of column var names. _col2name_ind : ndarray Array that maps jac col index to index of column name. """
[docs] def __init__(self, system): """ Initialize all attributes. """ self._system = weakref.ref(system) self._subjacs_info = system._subjacs_info self._under_complex_step = False self._abs_keys = defaultdict(bool) self._randomize = False self._col_var_info = None self._colnames = None self._col2name_ind = None
def _get_abs_key(self, key): abskey = self._abs_keys[key] if not abskey: self._abs_keys[key] = abskey = key2abs_key(self._system(), key) return abskey def _abs_key2shape(self, abs_key): """ Return shape of sub-jacobian for variables making up the key tuple. Parameters ---------- abs_key : (str, str) Absolute name pair of sub-Jacobian. Returns ------- out_size : int local size of the output variable. in_size : int local size of the input variable. """ abs2meta = self._system()._var_allprocs_abs2meta of, wrt = abs_key if wrt in abs2meta['input']: sz = abs2meta['input'][wrt]['size'] else: sz = abs2meta['output'][wrt]['size'] return (abs2meta['output'][of]['size'], sz)
[docs] def __contains__(self, key): """ Return whether there is a subjac for the given promoted or relative name pair. Parameters ---------- key : (str, str) Promoted or relative name pair of sub-Jacobian. Returns ------- bool return whether sub-Jacobian has been defined. """ return self._get_abs_key(key) in self._subjacs_info
[docs] def __getitem__(self, key): """ Get sub-Jacobian. Parameters ---------- key : (str, str) Promoted or relative name pair of sub-Jacobian. Returns ------- ndarray or spmatrix or list[3] sub-Jacobian as an array, sparse mtx, or AIJ/IJ list or tuple. """ abs_key = self._get_abs_key(key) if abs_key in self._subjacs_info: return self._subjacs_info[abs_key]['val'] else: msg = '{}: Variable name pair ("{}", "{}") not found.' raise KeyError(msg.format(self.msginfo, key[0], key[1]))
[docs] def __setitem__(self, key, subjac): """ Set sub-Jacobian. Parameters ---------- key : (str, str) Promoted or relative name pair of sub-Jacobian. subjac : int or float or ndarray or sparse matrix sub-Jacobian as a scalar, vector, array, or AIJ list or tuple. """ abs_key = self._get_abs_key(key) if abs_key is not None: # You can only set declared subjacobians. if abs_key not in self._subjacs_info: msg = '{}: Variable name pair ("{}", "{}") must first be declared.' raise KeyError(msg.format(self.msginfo, key[0], key[1])) subjacs_info = self._subjacs_info[abs_key] if issparse(subjac): subjacs_info['val'] = subjac else: # np.promote_types will choose the smallest dtype that can contain both arguments subjac = np.atleast_1d(subjac) safe_dtype = np.promote_types(subjac.dtype, float) subjac = subjac.astype(safe_dtype, copy=False) rows = subjacs_info['rows'] if rows is None: # Dense subjac subjac = np.atleast_2d(subjac) if subjac.shape != (1, 1): shape = self._abs_key2shape(abs_key) subjac = subjac.reshape(shape) else: # Sparse subjac if subjac.shape != (1,) and subjac.shape != rows.shape: msg = '{}: Sub-jacobian for key {} has the wrong shape ({}), expected ({}).' raise ValueError(msg.format(self.msginfo, abs_key, subjac.shape, rows.shape)) subjacs_info['val'][:] = subjac else: msg = '{}: Variable name pair ("{}", "{}") not found.' raise KeyError(msg.format(self.msginfo, key[0], key[1]))
[docs] def __iter__(self): """ Yield next name pair of sub-Jacobian. Yields ------ str """ yield from self._subjacs_info.keys()
[docs] def keys(self): """ Yield next name pair of sub-Jacobian. Yields ------ str """ yield from self._subjacs_info.keys()
[docs] def items(self): """ Yield name pair and value of sub-Jacobian. Yields ------ str """ for key, meta in self._subjacs_info.items(): yield key, meta['val']
@property def msginfo(self): """ Return info to prepend to messages. Returns ------- str Info to prepend to messages. """ if self._system() is None: return type(self).__name__ return '{} in {}'.format(type(self).__name__, self._system().msginfo) 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. """ pass 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'. """ pass def _randomize_subjac(self, subjac, key): """ Return a subjac that is the given subjac filled with random values. Parameters ---------- subjac : ndarray or csc_matrix Sub-jacobian to be randomized. key : tuple (of, wrt) Key for subjac within the jacobian. Returns ------- ndarray or csc_matrix Randomized version of the subjac. """ if isinstance(subjac, sparse_types): # sparse sparse = subjac.copy() sparse.data = rand(sparse.data.size) sparse.data += 1.0 return sparse # if a subsystem has computed a dynamic partial or semi-total coloring, # we use that sparsity information to set the sparsity of the randomized # subjac. Otherwise all subjacs that didn't have sparsity declared by the # user will appear completely dense, which will lead to a total jacobian that # is more dense than it should be, causing any total coloring that we compute # to be overly conservative. subjac_info = self._subjacs_info[key] if 'sparsity' in subjac_info: assert subjac_info['rows'] is None rows, cols, shape = subjac_info['sparsity'] r = np.zeros(shape) val = rand(len(rows)) val += 1.0 r[rows, cols] = val else: r = rand(*subjac.shape) r += 1.0 return r
[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. """ for meta in self._subjacs_info.values(): if active: meta['val'] = meta['val'].astype(np.complex) else: meta['val'] = meta['val'].real self._under_complex_step = active
def _setup_index_maps(self, system): self._col_var_info = col_var_info = { t[0]: t for t in system._jac_wrt_iter() } self._colnames = list(col_var_info) # map var id to varname ncols = np.sum(end - start for _, start, end, _, _, _ in col_var_info.values()) self._col2name_ind = np.empty(ncols, dtype=INT_DTYPE) # jac col to var id start = end = 0 for i, (wrt, _start, _end, _, _, _) in enumerate(col_var_info.values()): end += _end - _start self._col2name_ind[start:end] = i start = end # for total derivs, we can have sub-indices making some subjacs smaller if system.pathname == '': for key, meta in system._subjacs_info.items(): nrows, ncols = meta['shape'] if key[0] in system._owns_approx_of_idx: ridxs = system._owns_approx_of_idx[key[0]] if len(ridxs) == nrows: ridxs = _full_slice # value was already changed else: ridxs = ridxs.shaped_array() else: ridxs = _full_slice if key[1] in system._owns_approx_wrt_idx: cidxs = system._owns_approx_wrt_idx[key[1]] if len(cidxs) == ncols: cidxs = _full_slice # value was already changed else: cidxs = cidxs.shaped_array() else: cidxs = _full_slice if ridxs is not _full_slice or cidxs is not _full_slice: # replace our local subjac with a smaller one but don't # change the subjac belonging to the system (which has values # shared with subsystems) if self._subjacs_info is system._subjacs_info: self._subjacs_info = system._subjacs_info.copy() meta = self._subjacs_info[key] = meta.copy() val = meta['val'] if ridxs is not _full_slice: nrows = len(ridxs) if cidxs is not _full_slice: ncols = len(cidxs) if meta['rows'] is None: # dense val = val[ridxs, :] val = val[:, cidxs] meta['val'] = val else: # sparse sprows = meta['rows'] spcols = meta['cols'] if ridxs is not _full_slice: sprows, mask = sparse_subinds(sprows, ridxs) spcols = spcols[mask] val = val[mask] if cidxs is not _full_slice: spcols, mask = sparse_subinds(spcols, cidxs) sprows = sprows[mask] val = val[mask] meta['rows'] = sprows meta['cols'] = spcols meta['val'] = val meta['shape'] = (nrows, ncols)
[docs] def set_col(self, system, icol, column): """ Set a column of the jacobian. The column is assumed to be the same size as a column of the jacobian. This also assumes that the column does not attempt to set any nonzero values that are outside of specified sparsity patterns for any of the subjacs. Parameters ---------- system : System The system that owns this jacobian. icol : int Column index. column : ndarray Column value. """ if self._colnames is None: self._setup_index_maps(system) wrt = self._colnames[self._col2name_ind[icol]] _, offset, _, _, _, _ = self._col_var_info[wrt] loc_idx = icol - offset # local col index into subjacs for of, start, end, _, _ in system._jac_of_iter(): key = (of, wrt) if key in self._subjacs_info: subjac = self._subjacs_info[key] if subjac['cols'] is None: # dense subjac['val'][:, loc_idx] = column[start:end] else: # our COO format match_inds = np.nonzero(subjac['cols'] == loc_idx)[0] if match_inds.size > 0: subjac['val'][match_inds] = column[start:end][subjac['rows'][match_inds]]
def _restore_approx_sparsity(self): """ Revert all subjacs back to the way they were as declared by the user. """ self._subjacs_info = self._system()._subjacs_info self._colnames = None # force recompute of internal index maps on next set_col