def _init_colored_approximations(self, system): from openmdao.core.group import Group from openmdao.core.implicitcomponent import ImplicitComponent self._colored_approx_groups = [] self._j_colored = None self._j_data_sizes = None self._j_data_offsets = None # don't do anything if the coloring doesn't exist yet coloring = system._coloring_info['coloring'] if not isinstance(coloring, coloring_mod.Coloring): return outputs = system._outputs inputs = system._inputs abs2meta = system._var_allprocs_abs2meta prom2abs_out = system._var_allprocs_prom2abs_list['output'] prom2abs_in = system._var_allprocs_prom2abs_list['input'] approx_wrt_idx = system._owns_approx_wrt_idx out_slices = outputs.get_slice_dict() in_slices = inputs.get_slice_dict() is_total = isinstance(system, Group) system._update_wrt_matches() wrt_matches = system._coloring_info['wrt_matches'] data = None keys = set() for key, apprx in iteritems(self._exec_dict): if key[0] in wrt_matches: options = apprx[0][1] if 'coloring' in options: if data is None: # data is the same for all colored approxs so we only need the first data = self._get_approx_data(system, key) keys.update(a[0] for a in apprx) if is_total and system.pathname == '': # top level approx totals of_names = system._owns_approx_of full_wrts = system._var_allprocs_abs_names['output'] + \ system._var_allprocs_abs_names['input'] wrt_names = system._owns_approx_wrt else: of_names, wrt_names = system._get_partials_varlists() wrt_names = [ prom2abs_in[n][0] if n in prom2abs_in else prom2abs_out[n][0] for n in wrt_names ] full_wrts = wrt_names tmpJ = { '@nrows': coloring._shape[0], '@ncols': coloring._shape[1], '@out_slices': out_slices, '@approxs': keys, '@jac_slices': {}, } # FIXME: need to deal with mix of local/remote indices len_full_ofs = len(system._var_allprocs_abs_names['output']) full_idxs = [] approx_of_idx = system._owns_approx_of_idx jac_slices = tmpJ['@jac_slices'] for abs_of, roffset, rend, _ in system._jacobian_of_iter(): rslice = slice(roffset, rend) for abs_wrt, coffset, cend, _ in system._jacobian_wrt_iter( wrt_matches): jac_slices[(abs_of, abs_wrt)] = (rslice, slice(coffset, cend)) if is_total and (approx_of_idx or len_full_ofs > len(of_names)): slc = out_slices[abs_of] if abs_of in approx_of_idx: full_idxs.append( np.arange(slc.start, slc.stop)[approx_of_idx[abs_of]]) else: full_idxs.append(range(slc.start, slc.stop)) if full_idxs: tmpJ['@row_idx_map'] = np.hstack(full_idxs) if len(full_wrts) != len(wrt_matches) or approx_wrt_idx: if is_total and system.pathname == '': # top level approx totals full_wrt_sizes = [abs2meta[wrt]['size'] for wrt in wrt_names] else: _, full_wrt_sizes = system._get_partials_var_sizes() # need mapping from coloring jac columns (subset) to full jac columns col_map = sub2full_indices(full_wrts, wrt_matches, full_wrt_sizes, approx_wrt_idx) else: col_map = None # get groups of columns from the coloring and compute proper indices into # the inputs and outputs vectors. is_semi = is_total and system.pathname use_full_cols = isinstance(system, ImplicitComponent) or is_semi for cols, nzrows in coloring.color_nonzero_iter('fwd'): ccols = cols if col_map is None else col_map[cols] idx_info = get_input_idx_split(ccols, inputs, outputs, use_full_cols, is_total) self._colored_approx_groups.append( (data, cols, tmpJ, idx_info, nzrows))
def _init_colored_approximations(self, system): from openmdao.core.group import Group from openmdao.core.implicitcomponent import ImplicitComponent is_group = isinstance(system, Group) is_total = is_group and system.pathname == '' is_semi = is_group and not is_total use_full_cols = is_semi or isinstance(system, ImplicitComponent) self._colored_approx_groups = [] # don't do anything if the coloring doesn't exist yet coloring = system._coloring_info['coloring'] if not isinstance(coloring, coloring_mod.Coloring): return system._update_wrt_matches(system._coloring_info) wrt_matches = system._coloring_info['wrt_matches'] out_slices = system._outputs.get_slice_dict() if wrt_matches is not None: # this maps column indices into colored jac into indices into full jac ccol2jcol = np.empty(coloring._shape[1], dtype=INT_DTYPE) # colored col to out vec idx if is_total: ccol2vcol = np.empty(coloring._shape[1], dtype=INT_DTYPE) ordered_wrt_iter = list(system._jac_wrt_iter()) colored_start = colored_end = 0 for abs_wrt, cstart, cend, vec, cinds in ordered_wrt_iter: if wrt_matches is None or abs_wrt in wrt_matches: colored_end += cend - cstart ccol2jcol[colored_start:colored_end] = np.arange( cstart, cend, dtype=INT_DTYPE) if is_total and abs_wrt in out_slices: slc = out_slices[abs_wrt] rng = np.arange(slc.start, slc.stop) if cinds is not None: rng = rng[cinds] ccol2vcol[colored_start:colored_end] = rng colored_start = colored_end row_var_sizes = { v: sz for v, sz in zip(coloring._row_vars, coloring._row_var_sizes) } row_map = np.empty(coloring._shape[0], dtype=INT_DTYPE) abs2prom = system._var_allprocs_abs2prom['output'] if is_total: it = [(of, end - start) for of, start, end, _ in system._jac_of_iter()] else: it = [(n, arr.size) for n, arr in system._outputs._abs_item_iter()] start = end = colorstart = colorend = 0 for name, sz in it: end += sz prom = name if is_total else abs2prom[name] if prom in row_var_sizes: colorend += row_var_sizes[prom] row_map[colorstart:colorend] = np.arange(start, end, dtype=INT_DTYPE) colorstart = colorend start = end for wrt, meta in self._wrt_meta.items(): if wrt_matches is None or wrt in wrt_matches: # data is the same for all colored approxs so we only need the first data = self._get_approx_data(system, wrt, meta) break outputs = system._outputs inputs = system._inputs for cols, nzrows in coloring.color_nonzero_iter('fwd'): nzrows = [row_map[r] for r in nzrows] jaccols = cols if wrt_matches is None else ccol2jcol[cols] if is_total: vcols = ccol2vcol[cols] else: vcols = jaccols vec_ind_list = get_input_idx_split(vcols, inputs, outputs, use_full_cols, is_total) self._colored_approx_groups.append( (data, jaccols, vec_ind_list, nzrows))