Merge 87b4ada into fa88683

openmdao/approximation_schemes/approximation_scheme.py

@@ -6,10 +6,12 @@
 
 from openmdao.core.constants import INT_DTYPE
 from openmdao.vectors.vector import _full_slice
-from openmdao.utils.array_utils import get_input_idx_split
+from openmdao.utils.array_utils import get_input_idx_split, ValueRepeater
 import openmdao.utils.coloring as coloring_mod
 from openmdao.utils.general_utils import _convert_auto_ivc_to_conn_name, LocalRangeIterable
 from openmdao.utils.mpi import check_mpi_env
+from openmdao.utils.rangemapper import RangeMapper
+
 
 use_mpi = check_mpi_env()
 if use_mpi is False:
@@ -40,9 +42,6 @@ class ApproximationScheme(object):
         A dict that maps wrt name to its fd/cs metadata.
     _progress_out : None or file-like object
         Attribute to output the progress of check_totals
-    _during_sparsity_comp : bool
-        If True, we're doing a sparsity computation and uncolored approxs need to be restricted
-        to only colored columns.
     _jac_scatter : tuple
         Data needed to scatter values from results array to a total jacobian column.
     _totals_directions : dict
@@ -62,7 +61,6 @@ def __init__(self):
         self._approx_groups_cached_under_cs = False
         self._wrt_meta = {}
         self._progress_out = None
-        self._during_sparsity_comp = False
         self._jac_scatter = None
         self._totals_directions = {}
         self._totals_directional_mode = None
@@ -84,7 +82,6 @@ def _reset(self):
         """
         self._colored_approx_groups = None
         self._approx_groups = None
-        self._during_sparsity_comp = False
 
     def _get_approx_groups(self, system, under_cs=False):
         """
@@ -139,54 +136,58 @@ def _init_colored_approximations(self, system):
         is_total = system.pathname == ''
         is_semi = _is_group(system) and not is_total
         self._colored_approx_groups = []
+        wrt_ranges = []
 
         # don't do anything if the coloring doesn't exist yet
-        coloring = system._coloring_info['coloring']
+        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']
+        wrt_matches = system._coloring_info._update_wrt_matches(system)
         out_slices = system._outputs.get_slice_dict()
 
+        # this maps column indices into colored jac into indices into full jac
         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)
+        # colored col to out vec idx
+        if is_total:
+            ccol2outvec = np.empty(coloring._shape[1], dtype=INT_DTYPE)
 
+        if is_total or wrt_matches is not None:
             colored_start = colored_end = 0
-            for abs_wrt, cstart, cend, vec, cinds, _ in system._jac_wrt_iter():
+            for abs_wrt, cstart, cend, _, cinds, _ in system._jac_wrt_iter():
                 if wrt_matches is None or abs_wrt in wrt_matches:
                     colored_end += cend - cstart
-                    ccol2jcol[colored_start:colored_end] = range(cstart, cend)
+                    if wrt_matches is not None:
+                        ccol2jcol[colored_start:colored_end] = range(cstart, cend)
                     if is_total and abs_wrt in out_slices:
                         slc = out_slices[abs_wrt]
-                        if cinds is None or cinds is _full_slice:
-                            rng = range(slc.start, slc.stop)
-                        else:
+                        if cinds is not None:
                             rng = np.arange(slc.start, slc.stop)[cinds]
-                        ccol2vcol[colored_start:colored_end] = rng
+                        else:
+                            rng = range(slc.start, slc.stop)
+                        wrt_ranges.append((abs_wrt, slc.stop - slc.start))
+                        ccol2outvec[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()]
+            it = ((of, end - start) for of, start, end, _, _ in system._jac_of_iter())
+            rangemapper = RangeMapper.create(wrt_ranges)
         else:
-            it = [(n, arr.size) for n, arr in system._outputs._abs_item_iter()]
+            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)
+                row_map[colorstart:colorend] = range(start, end)
                 colorstart = colorend
             start = end
 
@@ -206,11 +207,13 @@ def _init_colored_approximations(self, system):
             nzrows = [row_map[r] for r in nzrows]
             jaccols = cols if wrt_matches is None else ccol2jcol[cols]
             if is_total:
-                vcols = ccol2vcol[cols]
+                vcols = ccol2outvec[cols]
+                seed_vars = rangemapper.inds2keys(cols)
             else:
                 vcols = jaccols
+                seed_vars = None
             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))
+            self._colored_approx_groups.append((data, jaccols, vec_ind_list, nzrows, seed_vars))
 
     def _init_approximations(self, system):
         """
@@ -222,19 +225,17 @@ def _init_approximations(self, system):
             The system having its derivs approximated.
         """
         total = system.pathname == ''
-        abs2meta = system._var_allprocs_abs2meta
 
         in_slices = system._inputs.get_slice_dict()
         out_slices = system._outputs.get_slice_dict()
 
-        approx_wrt_idx = system._owns_approx_wrt_idx
         coloring = system._get_static_coloring()
 
         self._approx_groups = []
         self._nruns_uncolored = 0
 
-        if self._during_sparsity_comp:
-            wrt_matches = system._coloring_info['wrt_matches']
+        if system._during_sparsity:
+            wrt_matches = system._coloring_info.wrt_matches
         else:
             wrt_matches = None
 
@@ -243,7 +244,8 @@ def _init_approximations(self, system):
             in_inds_directional = []
             vec_inds_directional = defaultdict(list)
 
-        for wrt, start, end, vec, _, _ in system._jac_wrt_iter(wrt_matches):
+        # wrt here is an absolute name (source if total)
+        for wrt, start, end, vec, sinds, _ in system._jac_wrt_iter(wrt_matches):
             if wrt in self._wrt_meta:
                 meta = self._wrt_meta[wrt]
                 if coloring is not None and 'coloring' in meta:
@@ -258,12 +260,12 @@ def _init_approximations(self, system):
 
                 in_idx = range(start, end)
 
-                if wrt in approx_wrt_idx:
+                if total and sinds is not _full_slice:
                     if vec is None:
-                        vec_idx = repeat(None, approx_wrt_idx[wrt].shaped_array().size)
+                        vec_idx = ValueRepeater(None, sinds.size)
                     else:
                         # local index into var
-                        vec_idx = approx_wrt_idx[wrt].shaped_array(copy=True)
+                        vec_idx = sinds.copy()
                         # convert into index into input or output vector
                         vec_idx += slices[wrt].start
                         # Directional derivatives for quick deriv checking.
@@ -297,8 +299,8 @@ def _init_approximations(self, system):
                         direction = self._totals_directions['fwd'][start:end]
                     else:
                         direction = meta['vector']
-                    self._approx_groups.append((wrt, data, in_idx, [(vec, vec_idx)], directional,
-                                                direction))
+                    self._approx_groups.append(((wrt,) if directional else wrt, data, in_idx,
+                                                [(vec, vec_idx)], directional, direction))
 
         if total:
             if self._totals_directional_mode == 'rev':
@@ -347,8 +349,8 @@ def _colored_column_iter(self, system, colored_approx_groups):
         ndarray
             solution array corresponding to the jacobian column at the given column index
         """
-        total_or_semi = _is_group(system)
         total = system.pathname == ''
+        total_or_semi = total or _is_group(system)
 
         if total:
             tot_result = np.zeros(sum([end - start for _, start, end, _, _
@@ -371,12 +373,18 @@ def _colored_column_iter(self, system, colored_approx_groups):
         nruns = len(colored_approx_groups)
         tosend = None
 
-        for data, jcols, vec_ind_list, nzrows in colored_approx_groups:
+        for data, jcols, vec_ind_list, nzrows, seed_vars, in colored_approx_groups:
             mult = self._get_multiplier(data)
 
             if fd_count % num_par_fd == system._par_fd_id:
                 # run the finite difference
-                result = self._run_point(system, vec_ind_list, data, results_array, total_or_semi)
+                if total:
+                    with system._relevant.seeds_active(fwd_seeds=seed_vars):
+                        result = self._run_point(system, vec_ind_list, data, results_array,
+                                                 total_or_semi)
+                else:
+                    result = self._run_point(system, vec_ind_list, data, results_array,
+                                             total_or_semi)
 
                 if par_fd_w_serial_model or not use_parallel_fd:
                     result = self._transform_result(result)
@@ -411,7 +419,7 @@ def _colored_column_iter(self, system, colored_approx_groups):
 
                 i, res = tup
 
-                _, jcols, _, nzrows = colored_approx_groups[i]
+                _, jcols, _, nzrows, _ = colored_approx_groups[i]
 
                 for i, col in enumerate(jcols):
                     scratch[:] = 0.0
@@ -511,9 +519,16 @@ def _uncolored_column_iter(self, system, approx_groups):
 
                 if fd_count % num_par_fd == system._par_fd_id:
                     # run the finite difference
-                    result = self._run_point(system, vec_ind_info,
-                                             app_data, results_array, total_or_semi,
-                                             jcol_idxs)
+                    if total:
+                        seeds = wrt if directional else (wrt,)
+                        with system._relevant.seeds_active(fwd_seeds=seeds):
+                            result = self._run_point(system, vec_ind_info,
+                                                     app_data, results_array, total_or_semi,
+                                                     jcol_idxs)
+                    else:
+                        result = self._run_point(system, vec_ind_info,
+                                                 app_data, results_array, total_or_semi,
+                                                 jcol_idxs)
 
                     result = self._transform_result(result)
 
@@ -580,15 +595,14 @@ def compute_approximations(self, system, jac=None):
         if not self._wrt_meta:
             return
 
-        if jac is None:
+        if system._tot_jac is not None:
+            jac = system._tot_jac
+        elif jac is None:
             jac = system._jacobian
 
         for ic, col in self.compute_approx_col_iter(system,
                                                     under_cs=system._outputs._under_complex_step):
-            if system._tot_jac is None:
-                jac.set_col(system, ic, col)
-            else:
-                system._tot_jac.set_col(ic, col)
+            jac.set_col(system, ic, col)
 
     def _compute_approx_col_iter(self, system, under_cs):
         # This will either generate new approx groups or use cached ones

openmdao/approximation_schemes/complex_step.py

@@ -136,7 +136,11 @@ def compute_approx_col_iter(self, system, under_cs=False):
         system._set_complex_step_mode(True)
 
         try:
-            yield from self._compute_approx_col_iter(system, under_cs=True)
+            for tup in self._compute_approx_col_iter(system, under_cs=True):
+                yield tup
+                # this was needed after adding relevance to the NL solve in order to clean
+                # out old results left over in the output array from a previous solve.
+                system._outputs.set_val(saved_outputs)
         finally:
             # Turn off complex step.
             system._set_complex_step_mode(False)

openmdao/components/cross_product_comp.py

@@ -74,6 +74,8 @@ def initialize(self):
                             [0, 1, 0, 0, -1, 0],
                             [-1, 0, 1, 0, 0, 0]], dtype=float)
 
+        self._minus_k = -self._k
+
     def add_product(self, c_name, a_name='a', b_name='b',
                     c_units=None, a_units=None, b_units=None,
                     vec_size=1):
@@ -169,11 +171,8 @@ def add_product(self, c_name, a_name='a', b_name='b',
         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, val=0)
-
-        self.declare_partials(of=c_name, wrt=b_name,
-                              rows=row_idxs, cols=col_idxs, val=0)
+        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)
 
     def compute(self, inputs, outputs):
         """
@@ -208,6 +207,6 @@ def compute_partials(self, inputs, partials):
 
             # Use the following for sparse partials
             partials[product['c_name'], product['a_name']] = \
-                np.einsum('...j,ji->...i', b, self._k * -1).ravel()
+                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()