Skip to content

Commit

Permalink
Merge pull request #3069 from naylor-b/transfers5
Browse files Browse the repository at this point in the history 
Refactor of reverse transfers and fixes for FD group that contains parallel group(s).
  • Loading branch information
@swryan
swryan committed Dec 5, 2023
2 parents 2545bfe + 6e16d7c commit 51f1873
Show file tree
Hide file tree
Showing 48 changed files with 3,497 additions and 818 deletions.
42 changes: 20 additions & 22 deletions openmdao/approximation_schemes/approximation_scheme.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
"""Base class used to define the interface for derivative approximation schemes."""
import time
from collections import defaultdict
from itertools import repeat
import numpy as np

from openmdao.core.constants import INT_DTYPE
from openmdao.vectors.vector import _full_slice
from openmdao.utils.array_utils import get_input_idx_split
import openmdao.utils.coloring as coloring_mod
from openmdao.utils.general_utils import _convert_auto_ivc_to_conn_name, LocalRangeIterable
Expand Down Expand Up @@ -134,8 +136,8 @@ def add_approximation(self, abs_key, system, kwargs):
raise NotImplementedError("add_approximation has not been implemented")

def _init_colored_approximations(self, system):
is_group = _is_group(system)
is_total = is_group and system.pathname == ''
is_total = system.pathname == ''
is_semi = _is_group(system) and not is_total
self._colored_approx_groups = []

# don't do anything if the coloring doesn't exist yet
Expand All @@ -155,17 +157,17 @@ def _init_colored_approximations(self, system):
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, _, cinds, _ in ordered_wrt_iter:
for abs_wrt, cstart, cend, vec, 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] = np.arange(cstart, cend, dtype=INT_DTYPE)
ccol2jcol[colored_start:colored_end] = range(cstart, cend)
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]
if cinds is None or cinds is _full_slice:
rng = range(slc.start, slc.stop)
else:
rng = np.arange(slc.start, slc.stop)[cinds]
ccol2vcol[colored_start:colored_end] = rng
colored_start = colored_end

Expand Down Expand Up @@ -198,7 +200,6 @@ def _init_colored_approximations(self, system):
inputs = system._inputs

from openmdao.core.implicitcomponent import ImplicitComponent
is_semi = is_group and not is_total
use_full_cols = is_semi or isinstance(system, ImplicitComponent)

for cols, nzrows in coloring.color_nonzero_iter('fwd'):
Expand Down Expand Up @@ -259,7 +260,7 @@ def _init_approximations(self, system):

if wrt in approx_wrt_idx:
if vec is None:
vec_idx = None
vec_idx = repeat(None, approx_wrt_idx[wrt].shaped_array().size)
else:
# local index into var
vec_idx = approx_wrt_idx[wrt].shaped_array(copy=True)
Expand All @@ -271,15 +272,9 @@ def _init_approximations(self, system):
in_idx = [list(in_idx)]
vec_idx = [vec_idx]
else:
if vec is None: # remote wrt
if wrt in abs2meta['input']:
vec_idx = range(abs2meta['input'][wrt]['global_size'])
else:
vec_idx = range(abs2meta['output'][wrt]['global_size'])
else:
vec_idx = LocalRangeIterable(system, wrt)
if directional:
vec_idx = [v for v in vec_idx if v is not None]
vec_idx = LocalRangeIterable(system, wrt)
if directional and vec is not None:
vec_idx = [v for v in vec_idx if v is not None]

# Directional derivatives for quick deriv checking.
# Place the indices in a list so that they are all stepped at the same time.
Expand Down Expand Up @@ -445,14 +440,16 @@ def _vec_ind_iter(self, vec_ind_list):
entry = [[None, None]]
ent0 = entry[0]
for vec, vec_idxs in vec_ind_list:
if vec_idxs is None:
continue
for vinds in vec_idxs:
ent0[0] = vec
ent0[1] = vinds
yield entry, vinds

def _uncolored_column_iter(self, system, approx_groups):
"""
Perform approximations and yields (column_index, column) for each jac column.
Perform approximations and yield (column_index, column) for each jac column.
Parameters
----------
Expand All @@ -477,9 +474,10 @@ def _uncolored_column_iter(self, system, approx_groups):
solution array corresponding to the jacobian column at the given column index
"""
total = system.pathname == ''
ordered_of_iter = list(system._jac_of_iter())
if total:
tot_result = np.zeros(ordered_of_iter[-1][2])
for _, _, end, _, _ in system._jac_of_iter():
pass
tot_result = np.zeros(end)

total_or_semi = total or _is_group(system)

Expand Down
13 changes: 1 addition & 12 deletions openmdao/core/component.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1006,17 +1006,6 @@ def _update_dist_src_indices(self, abs_in2out, all_abs2meta, all_abs2idx, all_si
if dist_in:
offset = np.sum(sizes_in[:iproc, i])
end = offset + sizes_in[iproc, i]
else:
if src.startswith('_auto_ivc.'):
nzs = np.nonzero(vout_sizes)[0]
if nzs.size == 1:
# special case where we have a 'distributed' auto_ivc output
# that has a nonzero value in only one proc, so we can treat
# it like a non-distributed output. This happens in cases
# where an auto_ivc output connects to a variable that is
# remote on at least one proc.
offset = 0
end = vout_sizes[nzs[0]]

# total sizes differ and output is distributed, so can't determine mapping
if offset is None:
Expand Down Expand Up @@ -1776,7 +1765,7 @@ def _get_dist_nz_dresids(self):
"""
nzresids = []
dresids = self._dresiduals.asarray()
for of, start, end, _full_slice, dist_sizes in self._jac_of_iter():
for of, start, end, _, dist_sizes in self._jac_of_iter():
if dist_sizes is not None:
if np.any(dresids[start:end]):
nzresids.append(of)
Expand Down

0 comments on commit 51f1873

Please sign in to comment.