Skip to content

Commit

Permalink
Merge d6a498c into 2f6d165
Browse files Browse the repository at this point in the history
  • Loading branch information
@naylor-b
naylor-b committed Oct 2, 2019
2 parents 2f6d165 + d6a498c commit 63525bb
Show file tree
Hide file tree
Showing 7 changed files with 225 additions and 91 deletions.
89 changes: 49 additions & 40 deletions openmdao/approximation_schemes/approximation_scheme.py
View file
Expand Up @@ -41,13 +41,20 @@ def __init__(self):
Initialize the ApproximationScheme.
"""
self._approx_groups = None
self._colored_approx_groups = []
self._colored_approx_groups = None
self._j_colored = None
self._j_data_sizes = None
self._j_data_offsets = None
self._approx_groups_cached_under_cs = False
self._exec_dict = defaultdict(list)

def _reset(self):
"""
Get rid of any existing approx groups.
"""
self._colored_approx_groups = None
self._approx_groups = None

def _get_approx_groups(self, system, under_cs=False):
"""
Retrieve data structure that contains all the approximations.
Expand Down Expand Up @@ -304,36 +311,37 @@ def _compute_approximations(self, system, jac, total, under_cs):
do_rows_cols = self._j_colored is None

# do colored solves first
for data, col_idxs, tmpJ, idx_info, nz_rows in colored_approx_groups:
colored_shape = (tmpJ['@nrows'], tmpJ['@ncols'])

if fd_count % num_par_fd == system._par_fd_id:
# run the finite difference
result = self._run_point(system, idx_info, data, results_array, total)
if par_fd_w_serial_model or not is_parallel:
rowmap = tmpJ['@row_idx_map'] if '@row_idx_map' in tmpJ else None
if rowmap is not None:
result = result[rowmap]
result = self._transform_result(result)

if nz_rows is None: # uncolored column
if do_rows_cols:
nrows = tmpJ['@nrows']
jrows.extend(range(nrows))
jcols.extend(col_idxs * nrows)
jdata.extend(result)
else:
for i, col in enumerate(col_idxs):
if colored_approx_groups is not None:
for data, col_idxs, tmpJ, idx_info, nz_rows in colored_approx_groups:
colored_shape = (tmpJ['@nrows'], tmpJ['@ncols'])

if fd_count % num_par_fd == system._par_fd_id:
# run the finite difference
result = self._run_point(system, idx_info, data, results_array, total)
if par_fd_w_serial_model or not is_parallel:
rowmap = tmpJ['@row_idx_map'] if '@row_idx_map' in tmpJ else None
if rowmap is not None:
result = result[rowmap]
result = self._transform_result(result)

if nz_rows is None: # uncolored column
if do_rows_cols:
jrows.extend(nz_rows[i])
jcols.extend([col] * len(nz_rows[i]))
jdata.extend(result[nz_rows[i]])
else: # parallel model (some vars are remote)
raise NotImplementedError("simul approx coloring with parallel FD/CS is "
"only supported currently when using "
"a serial model, i.e., when "
"num_par_fd == number of MPI procs.")
fd_count += 1
nrows = tmpJ['@nrows']
jrows.extend(range(nrows))
jcols.extend(col_idxs * nrows)
jdata.extend(result)
else:
for i, col in enumerate(col_idxs):
if do_rows_cols:
jrows.extend(nz_rows[i])
jcols.extend([col] * len(nz_rows[i]))
jdata.extend(result[nz_rows[i]])
else: # parallel model (some vars are remote)
raise NotImplementedError("simul approx coloring with parallel FD/CS is "
"only supported currently when using "
"a serial model, i.e., when "
"num_par_fd == number of MPI procs.")
fd_count += 1

# now do uncolored solves
for wrt, data, col_idxs, tmpJ, idx_info, nz_rows in approx_groups:
Expand Down Expand Up @@ -392,16 +400,17 @@ def _compute_approximations(self, system, jac, total, under_cs):
elif is_parallel: # uncolored with parallel systems
results = _gather_jac_results(mycomm, results)

for _, _, tmpJ, _, _ in colored_approx_groups:
# TODO: coloring when using parallel FD and/or FD with remote comps
for key in tmpJ['@approxs']:
slc = tmpJ['@jac_slices'][key]
if uses_voi_indices:
jac._override_checks = True
jac[key] = _from_dense(jacobian, key, Jcolored[slc])
jac._override_checks = False
else:
jac[key] = _from_dense(jacobian, key, Jcolored[slc])
if colored_approx_groups is not None:
for _, _, tmpJ, _, _ in colored_approx_groups:
# TODO: coloring when using parallel FD and/or FD with remote comps
for key in tmpJ['@approxs']:
slc = tmpJ['@jac_slices'][key]
if uses_voi_indices:
jac._override_checks = True
jac[key] = _from_dense(jacobian, key, Jcolored[slc])
jac._override_checks = False
else:
jac[key] = _from_dense(jacobian, key, Jcolored[slc])

Jcolored = None # clean up memory

Expand Down
2 changes: 1 addition & 1 deletion openmdao/approximation_schemes/complex_step.py
View file
Expand Up @@ -66,7 +66,7 @@ def add_approximation(self, abs_key, system, kwargs):

key = (abs_key[1], options['step'], options['directional'])
self._exec_dict[key].append((abs_key, options))
self._approx_groups = None # force later regen of approx_groups
self._reset() # force later regen of approx_groups

def _get_approx_data(self, system, data):
"""
Expand Down
2 changes: 1 addition & 1 deletion openmdao/approximation_schemes/finite_difference.py
View file
Expand Up @@ -127,7 +127,7 @@ def add_approximation(self, abs_key, system, kwargs):
key = (abs_key[1], options['form'], options['order'],
options['step'], options['step_calc'], options['directional'])
self._exec_dict[key].append((abs_key, options))
self._approx_groups = None # force later regen of approx_groups
self._reset() # force later regen of approx_groups

def _get_approx_data(self, system, data):
"""
Expand Down
5 changes: 1 addition & 4 deletions openmdao/core/group.py
View file
Expand Up @@ -2079,10 +2079,7 @@ def _setup_approx_partials(self):
approx = self._get_approx_scheme(method)
# reset the approx if necessary
approx._exec_dict = defaultdict(list)
approx._approx_groups = None
if wrt_matches:
# if coloring is active, force regen of coloring approxs
approx._colored_approx_groups = None
approx._reset()

approx_keys = self._get_approx_subjac_keys()
for key in approx_keys:
Expand Down
65 changes: 42 additions & 23 deletions openmdao/core/system.py
View file
Expand Up @@ -61,14 +61,14 @@
}

_DEFAULT_COLORING_META = {
'wrt_patterns': ('*',),
'method': 'fd',
'wrt_matches': None,
'per_instance': False,
'coloring': None, # this will contain the actual Coloring object
'dynamic': False, # True if dynamic coloring is being used
'static': None, # either _STD_COLORING_FNAME, a filename, or a Coloring object
# if use_fixed_coloring was called
'wrt_patterns': ('*',), # patterns used to match wrt variables
'method': 'fd', # finite differencing method ('fd' or 'cs')
'wrt_matches': None, # where matched wrt names are stored
'per_instance': True, # assume each instance can have a different coloring
'coloring': None, # this will contain the actual Coloring object
'dynamic': False, # True if dynamic coloring is being used
'static': None, # either _STD_COLORING_FNAME, a filename, or a Coloring object
# if use_fixed_coloring was called
}

_DEFAULT_COLORING_META.update(_DEF_COMP_SPARSITY_ARGS)
Expand Down Expand Up @@ -944,8 +944,7 @@ def use_fixed_coloring(self, coloring=_STD_COLORING_FNAME, recurse=True):
if isinstance(coloring, Coloring):
approx = self._get_approx_scheme(coloring._meta['method'])
# force regen of approx groups on next call to compute_approximations
approx._colored_approx_groups = None
approx._approx_groups = None
approx._reset()
return

if recurse:
Expand Down Expand Up @@ -1101,6 +1100,27 @@ def _compute_approx_coloring(self, recurse=False, **overrides):

approx_scheme = self._get_approx_scheme(self._coloring_info['method'])

if self._coloring_info['coloring'] is None and self._coloring_info['static'] is None:
self._coloring_info['dynamic'] = True

coloring_fname = self.get_approx_coloring_fname()

# if we find a previously computed class coloring for our class, just use that
# instead of regenerating a coloring.
if not info['per_instance'] and coloring_fname in coloring_mod._CLASS_COLORINGS:
info['coloring'] = coloring = coloring_mod._CLASS_COLORINGS[coloring_fname]
if coloring is None:
print("\nClass coloring for class '{}' wasn't good enough, "
"so skipping for '{}'".format(type(self).__name__, self.pathname))
info['static'] = info['dynamic'] = None
else:
print("\n{} using class coloring for class '{}'".format(self.pathname,
type(self).__name__))
info.update(coloring._meta)
# force regen of approx groups during next compute_approximations
approx_scheme._reset()
return [coloring]

from openmdao.core.group import Group
is_total = isinstance(self, Group)

Expand All @@ -1117,9 +1137,6 @@ def _compute_approx_coloring(self, recurse=False, **overrides):

starting_resids = self._residuals._data.copy()

if self._coloring_info['coloring'] is None and self._coloring_info['static'] is None:
self._coloring_info['dynamic'] = True

# for groups, this does some setup of approximations
self._setup_approx_coloring()

Expand All @@ -1141,14 +1158,14 @@ def _compute_approx_coloring(self, recurse=False, **overrides):
self._apply_nonlinear()

for scheme in self._approx_schemes.values():
scheme._approx_groups = None # force a re-initialization of approx
scheme._reset() # force a re-initialization of approx

self.run_linearize()
self._jacobian._save_sparsity(self)

sparsity_time = time.time() - sparsity_start_time

self._update_wrt_matches(self._coloring_info)
self._update_wrt_matches(info)

ordered_of_info = list(self._jacobian_of_iter())
ordered_wrt_info = list(self._jacobian_wrt_iter(info['wrt_matches']))
Expand Down Expand Up @@ -1176,6 +1193,8 @@ def _compute_approx_coloring(self, recurse=False, **overrides):
info['coloring'] = info['static'] = info['dynamic'] = None
simple_warning("%s: Coloring was deactivated. Improvement of %.1f%% was less than min "
"allowed (%.1f%%)." % (self.msginfo, pct, info['min_improve_pct']))
if not info['per_instance']:
coloring_mod._CLASS_COLORINGS[coloring_fname] = None
return [None]

coloring._row_vars = [t[0] for t in ordered_of_info]
Expand All @@ -1191,11 +1210,10 @@ def _compute_approx_coloring(self, recurse=False, **overrides):

approx = self._get_approx_scheme(coloring._meta['method'])
# force regen of approx groups during next compute_approximations
approx._colored_approx_groups = None
approx._approx_groups = None
approx._reset()

if info['show_sparsity'] or info['show_summary']:
print("\nApprox coloring for '%s' (class %s)\n" % (self.pathname, type(self).__name__))
print("\nApprox coloring for '%s' (class %s)" % (self.pathname, type(self).__name__))

if info['show_sparsity']:
coloring.display_txt()
Expand All @@ -1204,6 +1222,10 @@ def _compute_approx_coloring(self, recurse=False, **overrides):

self._save_coloring(coloring)

if not info['per_instance']:
# save the class coloring for other instances of this class to use
coloring_mod._CLASS_COLORINGS[coloring_fname] = coloring

# restore original inputs/outputs
self._inputs._data[:] = starting_inputs
self._outputs._data[:] = starting_outputs
Expand Down Expand Up @@ -1273,9 +1295,7 @@ def get_approx_coloring_fname(self):
# total coloring
return os.path.join(directory, 'total_coloring.pkl')

per_instance = self._coloring_info.get('per_instance')

if per_instance:
if self._coloring_info.get('per_instance'):
# base the name on the instance pathname
fname = 'coloring_' + self.pathname.replace('.', '_') + '.pkl'
else:
Expand Down Expand Up @@ -1331,8 +1351,7 @@ def _get_static_coloring(self):
info.update(info['coloring']._meta)
approx = self._get_approx_scheme(info['method'])
# force regen of approx groups during next compute_approximations
approx._colored_approx_groups = None
approx._approx_groups = None
approx._reset()
elif isinstance(static, coloring_mod.Coloring):
info['coloring'] = coloring = static

Expand Down

0 comments on commit 63525bb

Please sign in to comment.