def _initialize(self, system):
"""
Allocate the global matrices.
Parameters
----------
system : System
Parent system to this jacobian.
"""
# var_indices are the *global* indices for variables on this proc
is_top = system.pathname == ''
abs2meta = system._var_abs2meta
all_meta = system._var_allprocs_abs2meta
self._int_mtx = int_mtx = self._matrix_class(system.comm)
ext_mtx = self._matrix_class(system.comm)
iproc = system.comm.rank
abs2idx = system._var_allprocs_abs2idx['nonlinear']
out_sizes = system._var_sizes['nonlinear']['output']
in_sizes = system._var_sizes['nonlinear']['input']
out_ranges = self._out_ranges
in_ranges = self._in_ranges
abs2prom_out = system._var_abs2prom['output']
conns = {} if isinstance(system,
Component) else system._conn_global_abs_in2out
keymap = self._keymap
abs_key2shape = self._abs_key2shape
# create the matrix subjacs
for abs_key, info in iteritems(self._subjacs_info):
res_abs_name, wrt_abs_name = abs_key
# because self._subjacs_info is shared among all 'related' assembled jacs,
# we use out_ranges (and later in_ranges) to weed out keys outside of this jac
if res_abs_name not in out_ranges:
continue
res_offset, _ = out_ranges[res_abs_name]
if wrt_abs_name in abs2prom_out:
out_offset, _ = out_ranges[wrt_abs_name]
int_mtx._add_submat(abs_key, info, res_offset, out_offset,
None, info['shape'])
keymap[abs_key] = abs_key
elif wrt_abs_name in in_ranges:
if wrt_abs_name in conns: # connected input
out_abs_name = conns[wrt_abs_name]
meta_in = abs2meta[wrt_abs_name]
all_out_meta = all_meta[out_abs_name]
# calculate unit conversion
in_units = meta_in['units']
out_units = all_out_meta['units']
if in_units and out_units and in_units != out_units:
factor, _ = get_conversion(out_units, in_units)
if factor == 1.0:
factor = None
else:
factor = None
out_offset, _ = out_ranges[out_abs_name]
src_indices = abs2meta[wrt_abs_name]['src_indices']
# need to add an entry for d(output)/d(source)
# instead of d(output)/d(input)
abs_key2 = (res_abs_name, out_abs_name)
keymap[abs_key] = abs_key2
if src_indices is None:
shape = info['shape']
else:
shape = abs_key2shape(abs_key2)
if len(src_indices.shape) > 1:
src_indices = _flatten_src_indices(
src_indices, meta_in['shape'],
all_out_meta['global_shape'],
all_out_meta['global_size'])
int_mtx._add_submat(abs_key, info, res_offset, out_offset,
src_indices, shape, factor)
elif not is_top: # input is connected to something outside current system
ext_mtx._add_submat(abs_key, info, res_offset,
in_ranges[wrt_abs_name][0], None,
info['shape'])
iproc = system.comm.rank
out_size = np.sum(out_sizes[iproc, :])
int_mtx._build(out_size, out_size, in_ranges, out_ranges)
if ext_mtx._submats:
in_size = np.sum(in_sizes[iproc, :])
ext_mtx._build(out_size, in_size, in_ranges, out_ranges)
else:
ext_mtx = None
self._ext_mtx[system.pathname] = ext_mtx
def _setup_transfers(group):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
"""
rev = group._mode != 'fwd'
for subsys in group._subgroups_myproc:
subsys._setup_transfers()
abs2meta = group._var_abs2meta
allprocs_abs2meta = group._var_allprocs_abs2meta
myproc = group.comm.rank
transfers = group._transfers = {}
vectors = group._vectors
offsets = group._get_var_offsets()
vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names
mypathlen = len(group.pathname + '.' if group.pathname else '')
sub_inds = group._subsystems_inds
for vec_name in vec_names:
relvars, _ = group._relevant[vec_name]['@all']
# Initialize empty lists for the transfer indices
nsub_allprocs = len(group._subsystems_allprocs)
xfer_in = []
xfer_out = []
fwd_xfer_in = [[] for s in group._subsystems_allprocs]
fwd_xfer_out = [[] for s in group._subsystems_allprocs]
if rev:
rev_xfer_in = [[] for s in group._subsystems_allprocs]
rev_xfer_out = [[] for s in group._subsystems_allprocs]
allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
sizes_in = group._var_sizes[vec_name]['input']
sizes_out = group._var_sizes[vec_name]['output']
offsets_in = offsets[vec_name]['input']
offsets_out = offsets[vec_name]['output']
# Loop through all connections owned by this system
for abs_in, abs_out in group._conn_abs_in2out.items():
if abs_out not in relvars['output'] or abs_in not in relvars[
'input']:
continue
# Only continue if the input exists on this processor
if abs_in in abs2meta:
# Get meta
meta_in = abs2meta[abs_in]
meta_out = allprocs_abs2meta[abs_out]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is None:
owner = group._owning_rank[abs_out]
# if the input is larger than the output on a single proc, we have
# to just loop over the procs in the same way we do when src_indices
# is defined.
if meta_in['size'] > sizes_out[owner, idx_out]:
src_indices = np.arange(meta_in['size'],
dtype=INT_DTYPE)
elif src_indices.ndim == 1:
if _is_slice(src_indices):
indices = _slice_indices(src_indices,
meta_out['global_size'],
meta_out['global_shape'])
src_indices = convert_neg(indices,
meta_out['global_size'])
else:
src_indices = _flatten_src_indices(
src_indices, meta_in['shape'],
meta_out['global_shape'], meta_out['global_size'])
# 1. Compute the output indices
# NOTE: src_indices are relative to a single, possibly distributed variable,
# while the output_inds that we compute are relative to the full distributed
# array that contains all local variables from each rank stacked in rank order.
if src_indices is None:
if meta_out['distributed']:
# input in this case is non-distributed (else src_indices would be
# defined by now). The input size must match the full
# distributed size of the output.
for rank, sz in enumerate(sizes_out[:, idx_out]):
if sz > 0:
out_offset = offsets_out[rank, idx_out]
break
output_inds = np.arange(out_offset,
out_offset +
meta_out['global_size'],
dtype=INT_DTYPE)
else:
rank = myproc if abs_out in abs2meta else owner
offset = offsets_out[rank, idx_out]
output_inds = np.arange(offset,
offset + meta_in['size'],
dtype=INT_DTYPE)
else:
output_inds = np.zeros(src_indices.size, INT_DTYPE)
start = end = 0
for iproc in range(group.comm.size):
end += sizes_out[iproc, idx_out]
if start == end:
continue
# The part of src on iproc
on_iproc = np.logical_and(start <= src_indices,
src_indices < end)
if np.any(on_iproc):
# This converts from iproc-then-ivar to ivar-then-iproc ordering
# Subtract off part of previous procs
# Then add all variables on previous procs
# Then all previous variables on this proc
# - np.sum(out_sizes[:iproc, idx_out])
# + np.sum(out_sizes[:iproc, :])
# + np.sum(out_sizes[iproc, :idx_out])
# + inds
offset = offsets_out[iproc, idx_out] - start
output_inds[
on_iproc] = src_indices[on_iproc] + offset
start = end
# 2. Compute the input indices
input_inds = np.arange(offsets_in[myproc, idx_in],
offsets_in[myproc, idx_in] +
sizes_in[myproc, idx_in],
dtype=INT_DTYPE)
# Now the indices are ready - input_inds, output_inds
sub_in = abs_in[mypathlen:].split('.', 1)[0]
isub = sub_inds[sub_in]
fwd_xfer_in[isub].append(input_inds)
fwd_xfer_out[isub].append(output_inds)
if rev:
sub_out = abs_out[mypathlen:].split('.', 1)[0]
isub = sub_inds[sub_out]
rev_xfer_in[isub].append(input_inds)
rev_xfer_out[isub].append(output_inds)
transfers[vec_name] = {}
for isub in range(nsub_allprocs):
fwd_xfer_in[isub] = _merge(fwd_xfer_in[isub])
fwd_xfer_out[isub] = _merge(fwd_xfer_out[isub])
if rev:
rev_xfer_in[isub] = _merge(rev_xfer_in[isub])
rev_xfer_out[isub] = _merge(rev_xfer_out[isub])
xfer_in = np.concatenate(fwd_xfer_in)
xfer_out = np.concatenate(fwd_xfer_out)
out_vec = vectors['output'][vec_name]
xfer_all = PETScTransfer(vectors['input'][vec_name], out_vec,
xfer_in, xfer_out, group.comm)
transfers[vec_name]['fwd', None] = xfer_all
if rev:
transfers[vec_name]['rev', None] = xfer_all
for isub in range(nsub_allprocs):
transfers[vec_name]['fwd', isub] = PETScTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm)
if rev:
transfers[vec_name]['rev', isub] = PETScTransfer(
vectors['input'][vec_name],
vectors['output'][vec_name], rev_xfer_in[isub],
rev_xfer_out[isub], group.comm)
if group._use_derivatives:
transfers['nonlinear'] = transfers['linear']
def _setup_transfers(group, recurse=True):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
recurse : bool
Whether to call this method in subsystems.
"""
group._transfers = {}
iproc = group.comm.rank
rev = group._mode == 'rev' or group._mode == 'auto'
def merge(indices_list):
if len(indices_list) > 0:
return np.concatenate(indices_list)
else:
return _empty_idx_array
if recurse:
for subsys in group._subgroups_myproc:
subsys._setup_transfers(recurse)
# Pre-compute map from abs_names to the index of the containing subsystem
abs2isub = {}
for subsys, isub in zip(group._subsystems_myproc, group._subsystems_myproc_inds):
for type_ in ['input', 'output']:
for abs_name in subsys._var_allprocs_abs_names[type_]:
abs2isub[abs_name] = isub
abs2meta = group._var_abs2meta
allprocs_abs2meta = group._var_allprocs_abs2meta
transfers = group._transfers
vectors = group._vectors
offsets = _global2local_offsets(group._get_var_offsets())
vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names
for vec_name in vec_names:
relvars, _ = group._relevant[vec_name]['@all']
relvars_in = relvars['input']
relvars_out = relvars['output']
# Initialize empty lists for the transfer indices
nsub_allprocs = len(group._subsystems_allprocs)
xfer_in = []
xfer_out = []
fwd_xfer_in = [[] for i in range(nsub_allprocs)]
fwd_xfer_out = [[] for i in range(nsub_allprocs)]
if rev:
rev_xfer_in = [[] for i in range(nsub_allprocs)]
rev_xfer_out = [[] for i in range(nsub_allprocs)]
allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
sizes_in = group._var_sizes[vec_name]['input']
sizes_out = group._var_sizes[vec_name]['output']
offsets_in = offsets[vec_name]['input']
offsets_out = offsets[vec_name]['output']
# Loop through all explicit / implicit connections owned by this system
for abs_in, abs_out in iteritems(group._conn_abs_in2out):
if abs_out not in relvars_out or abs_in not in relvars_in:
continue
# Only continue if the input exists on this processor
if abs_in in abs2meta and abs_in in relvars['input']:
# Get meta
meta_in = abs2meta[abs_in]
meta_out = allprocs_abs2meta[abs_out]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is None:
pass
elif src_indices.ndim == 1:
src_indices = convert_neg(src_indices, meta_out['global_size'])
else:
src_indices = _flatten_src_indices(src_indices, meta_in['shape'],
meta_out['global_shape'],
meta_out['global_size'])
# 1. Compute the output indices
offset = offsets_out[iproc, idx_out]
if src_indices is None:
output_inds = np.arange(offset, offset + meta_in['size'], dtype=INT_DTYPE)
else:
output_inds = src_indices + offset
# 2. Compute the input indices
input_inds = np.arange(offsets_in[iproc, idx_in],
offsets_in[iproc, idx_in] +
sizes_in[iproc, idx_in], dtype=INT_DTYPE)
# Now the indices are ready - input_inds, output_inds
xfer_in.append(input_inds)
xfer_out.append(output_inds)
isub = abs2isub[abs_in]
fwd_xfer_in[isub].append(input_inds)
fwd_xfer_out[isub].append(output_inds)
if rev and abs_out in abs2isub:
isub = abs2isub[abs_out]
rev_xfer_in[isub].append(input_inds)
rev_xfer_out[isub].append(output_inds)
xfer_in = merge(xfer_in)
xfer_out = merge(xfer_out)
for isub in range(nsub_allprocs):
fwd_xfer_in[isub] = merge(fwd_xfer_in[isub])
fwd_xfer_out[isub] = merge(fwd_xfer_out[isub])
if rev:
rev_xfer_in[isub] = merge(rev_xfer_in[isub])
rev_xfer_out[isub] = merge(rev_xfer_out[isub])
out_vec = vectors['output'][vec_name]
transfers[vec_name] = {}
xfer_all = DefaultTransfer(vectors['input'][vec_name], out_vec,
xfer_in, xfer_out, group.comm)
transfers[vec_name]['fwd', None] = xfer_all
if rev:
transfers[vec_name]['rev', None] = xfer_all
for isub in range(nsub_allprocs):
transfers[vec_name]['fwd', isub] = DefaultTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm)
if rev:
transfers[vec_name]['rev', isub] = DefaultTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
rev_xfer_in[isub], rev_xfer_out[isub], group.comm)
if group._use_derivatives:
transfers['nonlinear'] = transfers['linear']
def _setup_transfers(group, recurse=True):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
recurse : bool
Whether to call this method in subsystems.
"""
group._transfers = {}
iproc = group.comm.rank
rev = group._mode == 'rev' or group._mode == 'auto'
def merge(indices_list):
if len(indices_list) > 0:
return np.concatenate(indices_list)
else:
return _empty_idx_array
if recurse:
for subsys in group._subgroups_myproc:
subsys._setup_transfers(recurse)
# Pre-compute map from abs_names to the index of the containing subsystem
abs2isub = {}
for subsys, isub in zip(group._subsystems_myproc,
group._subsystems_myproc_inds):
for type_ in ['input', 'output']:
for abs_name in subsys._var_allprocs_abs_names[type_]:
abs2isub[abs_name] = isub
abs2meta = group._var_abs2meta
allprocs_abs2meta = group._var_allprocs_abs2meta
transfers = group._transfers
vectors = group._vectors
offsets = _global2local_offsets(group._get_var_offsets())
vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names
for vec_name in vec_names:
relvars, _ = group._relevant[vec_name]['@all']
relvars_in = relvars['input']
relvars_out = relvars['output']
# Initialize empty lists for the transfer indices
nsub_allprocs = len(group._subsystems_allprocs)
xfer_in = []
xfer_out = []
fwd_xfer_in = [[] for i in range(nsub_allprocs)]
fwd_xfer_out = [[] for i in range(nsub_allprocs)]
if rev:
rev_xfer_in = [[] for i in range(nsub_allprocs)]
rev_xfer_out = [[] for i in range(nsub_allprocs)]
allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
sizes_in = group._var_sizes[vec_name]['input']
sizes_out = group._var_sizes[vec_name]['output']
offsets_in = offsets[vec_name]['input']
offsets_out = offsets[vec_name]['output']
# Loop through all explicit / implicit connections owned by this system
for abs_in, abs_out in iteritems(group._conn_abs_in2out):
if abs_out not in relvars_out or abs_in not in relvars_in:
continue
# Only continue if the input exists on this processor
if abs_in in abs2meta and abs_in in relvars['input']:
# Get meta
meta_in = abs2meta[abs_in]
meta_out = allprocs_abs2meta[abs_out]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is None:
pass
elif src_indices.ndim == 1:
src_indices = convert_neg(src_indices,
meta_out['global_size'])
else:
src_indices = _flatten_src_indices(
src_indices, meta_in['shape'],
meta_out['global_shape'], meta_out['global_size'])
# 1. Compute the output indices
offset = offsets_out[iproc, idx_out]
if src_indices is None:
output_inds = np.arange(offset,
offset + meta_in['size'],
dtype=INT_DTYPE)
else:
output_inds = src_indices + offset
# 2. Compute the input indices
input_inds = np.arange(offsets_in[iproc, idx_in],
offsets_in[iproc, idx_in] +
sizes_in[iproc, idx_in],
dtype=INT_DTYPE)
# Now the indices are ready - input_inds, output_inds
xfer_in.append(input_inds)
xfer_out.append(output_inds)
isub = abs2isub[abs_in]
fwd_xfer_in[isub].append(input_inds)
fwd_xfer_out[isub].append(output_inds)
if rev and abs_out in abs2isub:
isub = abs2isub[abs_out]
rev_xfer_in[isub].append(input_inds)
rev_xfer_out[isub].append(output_inds)
xfer_in = merge(xfer_in)
xfer_out = merge(xfer_out)
for isub in range(nsub_allprocs):
fwd_xfer_in[isub] = merge(fwd_xfer_in[isub])
fwd_xfer_out[isub] = merge(fwd_xfer_out[isub])
if rev:
rev_xfer_in[isub] = merge(rev_xfer_in[isub])
rev_xfer_out[isub] = merge(rev_xfer_out[isub])
out_vec = vectors['output'][vec_name]
transfers[vec_name] = {}
xfer_all = DefaultTransfer(vectors['input'][vec_name], out_vec,
xfer_in, xfer_out, group.comm)
transfers[vec_name]['fwd', None] = xfer_all
if rev:
transfers[vec_name]['rev', None] = xfer_all
for isub in range(nsub_allprocs):
transfers[vec_name]['fwd', isub] = DefaultTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm)
if rev:
transfers[vec_name]['rev', isub] = DefaultTransfer(
vectors['input'][vec_name],
vectors['output'][vec_name], rev_xfer_in[isub],
rev_xfer_out[isub], group.comm)
if group._use_derivatives:
transfers['nonlinear'] = transfers['linear']
def _setup_transfers(group, recurse=True):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
recurse : bool
Whether to call this method in subsystems.
"""
rev = group._mode == 'rev' or group._mode == 'auto'
def merge(indices_list):
if len(indices_list) > 0:
return np.concatenate(indices_list)
else:
return _empty_idx_array
if recurse:
for subsys in group._subgroups_myproc:
subsys._setup_transfers(recurse)
# Pre-compute map from abs_names to the index of the containing subsystem
abs2isub = {}
for subsys, isub in zip(group._subsystems_myproc,
group._subsystems_myproc_inds):
for type_ in ['input', 'output']:
for abs_name in subsys._var_allprocs_abs_names[type_]:
abs2isub[abs_name] = isub
abs2meta = group._var_abs2meta
allprocs_abs2meta = group._var_allprocs_abs2meta
myproc = group.comm.rank
transfers = group._transfers = {}
vectors = group._vectors
offsets = group._get_var_offsets()
vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names
for vec_name in vec_names:
relvars, _ = group._relevant[vec_name]['@all']
# Initialize empty lists for the transfer indices
nsub_allprocs = len(group._subsystems_allprocs)
xfer_in = []
xfer_out = []
fwd_xfer_in = [[] for i in range(nsub_allprocs)]
fwd_xfer_out = [[] for i in range(nsub_allprocs)]
if rev:
rev_xfer_in = [[] for i in range(nsub_allprocs)]
rev_xfer_out = [[] for i in range(nsub_allprocs)]
allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
sizes_in = group._var_sizes[vec_name]['input']
sizes_out = group._var_sizes[vec_name]['output']
offsets_in = offsets[vec_name]['input']
offsets_out = offsets[vec_name]['output']
# Loop through all connections owned by this system
for abs_in, abs_out in iteritems(group._conn_abs_in2out):
if abs_out not in relvars['output']:
continue
# Only continue if the input exists on this processor
if abs_in in abs2meta and abs_in in relvars['input']:
# Get meta
meta_in = abs2meta[abs_in]
meta_out = allprocs_abs2meta[abs_out]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is None:
owner = group._owning_rank[abs_out]
# if the input is larger than the output on a single proc, we have
# to just loop over the procs in the same way we do when src_indices
# is defined.
if meta_in['size'] > sizes_out[owner, idx_out]:
src_indices = np.arange(meta_in['size'],
dtype=INT_DTYPE)
elif src_indices.ndim == 1:
src_indices = convert_neg(src_indices,
meta_out['global_size'])
else:
src_indices = _flatten_src_indices(
src_indices, meta_in['shape'],
meta_out['global_shape'], meta_out['global_size'])
# 1. Compute the output indices
if src_indices is None:
start = 0 if owner == 0 else np.sum(sizes_out[:owner,
idx_out])
offset = offsets_out[owner, idx_out] - start
output_inds = np.arange(offset,
offset + meta_in['size'],
dtype=INT_DTYPE)
else:
output_inds = np.zeros(src_indices.size, INT_DTYPE)
start = end = 0
for iproc in range(group.comm.size):
end += sizes_out[iproc, idx_out]
# The part of src on iproc
on_iproc = np.logical_and(start <= src_indices,
src_indices < end)
if np.any(on_iproc):
# This converts from iproc-then-ivar to ivar-then-iproc ordering
# Subtract off part of previous procs
# Then add all variables on previous procs
# Then all previous variables on this proc
# - np.sum(out_sizes[:iproc, idx_out])
# + np.sum(out_sizes[:iproc, :])
# + np.sum(out_sizes[iproc, :idx_out])
# + inds
offset = offsets_out[iproc, idx_out] - start
output_inds[
on_iproc] = src_indices[on_iproc] + offset
start = end
# 2. Compute the input indices
input_inds = np.arange(offsets_in[myproc, idx_in],
offsets_in[myproc, idx_in] +
sizes_in[myproc, idx_in],
dtype=INT_DTYPE)
# Now the indices are ready - input_inds, output_inds
xfer_in.append(input_inds)
xfer_out.append(output_inds)
isub = abs2isub[abs_in]
fwd_xfer_in[isub].append(input_inds)
fwd_xfer_out[isub].append(output_inds)
if rev and abs_out in abs2isub:
isub = abs2isub[abs_out]
rev_xfer_in[isub].append(input_inds)
rev_xfer_out[isub].append(output_inds)
xfer_in = merge(xfer_in)
xfer_out = merge(xfer_out)
for isub in range(nsub_allprocs):
fwd_xfer_in[isub] = merge(fwd_xfer_in[isub])
fwd_xfer_out[isub] = merge(fwd_xfer_out[isub])
if rev:
rev_xfer_in[isub] = merge(rev_xfer_in[isub])
rev_xfer_out[isub] = merge(rev_xfer_out[isub])
out_vec = vectors['output'][vec_name]
transfers[vec_name] = {}
xfer_all = PETScTransfer(vectors['input'][vec_name], out_vec,
xfer_in, xfer_out, group.comm)
transfers[vec_name]['fwd', None] = xfer_all
if rev:
transfers[vec_name]['rev', None] = xfer_all
for isub in range(nsub_allprocs):
transfers[vec_name]['fwd', isub] = PETScTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm)
if rev:
transfers[vec_name]['rev', isub] = PETScTransfer(
vectors['input'][vec_name],
vectors['output'][vec_name], rev_xfer_in[isub],
rev_xfer_out[isub], group.comm)
if group._use_derivatives:
transfers['nonlinear'] = transfers['linear']
def _setup_transfers(group, recurse=True):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
recurse : bool
Whether to call this method in subsystems.
"""
iproc = group.comm.rank
rev = group._mode == 'rev' or group._mode == 'auto'
if recurse:
for subsys in group._subgroups_myproc:
subsys._setup_transfers(recurse)
abs2meta = group._var_abs2meta
allprocs_abs2meta = group._var_allprocs_abs2meta
group._transfers = transfers = {}
vectors = group._vectors
offsets = _global2local_offsets(group._get_var_offsets())
vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names
mypathlen = len(group.pathname + '.' if group.pathname else '')
sub_inds = group._subsystems_inds
for vec_name in vec_names:
relvars, _ = group._relevant[vec_name]['@all']
relvars_in = relvars['input']
relvars_out = relvars['output']
# Initialize empty lists for the transfer indices
nsub_allprocs = len(group._subsystems_allprocs)
xfer_in = []
xfer_out = []
fwd_xfer_in = [[] for s in group._subsystems_allprocs]
fwd_xfer_out = [[] for s in group._subsystems_allprocs]
if rev:
rev_xfer_in = [[] for s in group._subsystems_allprocs]
rev_xfer_out = [[] for s in group._subsystems_allprocs]
allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
sizes_in = group._var_sizes[vec_name]['input']
sizes_out = group._var_sizes[vec_name]['output']
offsets_in = offsets[vec_name]['input']
offsets_out = offsets[vec_name]['output']
# Loop through all connections owned by this group
for abs_in, abs_out in group._conn_abs_in2out.items():
if abs_out not in relvars_out or abs_in not in relvars_in:
continue
# Only continue if the input exists on this processor
if abs_in in abs2meta:
# Get meta
meta_in = abs2meta[abs_in]
meta_out = allprocs_abs2meta[abs_out]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is None:
pass
elif src_indices.ndim == 1:
src_indices = convert_neg(src_indices,
meta_out['global_size'])
else:
src_indices = _flatten_src_indices(
src_indices, meta_in['shape'],
meta_out['global_shape'], meta_out['global_size'])
# 1. Compute the output indices
offset = offsets_out[iproc, idx_out]
if src_indices is None:
output_inds = np.arange(offset,
offset + meta_in['size'],
dtype=INT_DTYPE)
else:
output_inds = src_indices + offset
# 2. Compute the input indices
input_inds = np.arange(offsets_in[iproc, idx_in],
offsets_in[iproc, idx_in] +
sizes_in[iproc, idx_in],
dtype=INT_DTYPE)
# Now the indices are ready - input_inds, output_inds
sub_in = abs_in[mypathlen:].split('.', 1)[0]
isub = sub_inds[sub_in]
fwd_xfer_in[isub].append(input_inds)
fwd_xfer_out[isub].append(output_inds)
if rev and abs_out in abs2meta:
sub_out = abs_out[mypathlen:].split('.', 1)[0]
isub = sub_inds[sub_out]
rev_xfer_in[isub].append(input_inds)
rev_xfer_out[isub].append(output_inds)
tot_size = 0
for isub in range(nsub_allprocs):
fwd_xfer_in[isub] = _merge(fwd_xfer_in[isub])
fwd_xfer_out[isub] = _merge(fwd_xfer_out[isub])
tot_size += fwd_xfer_in[isub].size
if rev:
rev_xfer_in[isub] = _merge(rev_xfer_in[isub])
rev_xfer_out[isub] = _merge(rev_xfer_out[isub])
transfers[vec_name] = {}
if tot_size > 0:
xfer_in = np.concatenate(fwd_xfer_in)
xfer_out = np.concatenate(fwd_xfer_out)
out_vec = vectors['output'][vec_name]
xfer_all = DefaultTransfer(vectors['input'][vec_name], out_vec,
xfer_in, xfer_out, group.comm)
else:
xfer_all = None
transfers[vec_name]['fwd', None] = xfer_all
if rev:
transfers[vec_name]['rev', None] = xfer_all
for isub in range(nsub_allprocs):
if fwd_xfer_in[isub].size > 0:
transfers[vec_name]['fwd', isub] = DefaultTransfer(
vectors['input'][vec_name],
vectors['output'][vec_name], fwd_xfer_in[isub],
fwd_xfer_out[isub], group.comm)
else:
transfers[vec_name]['fwd', isub] = None
if rev:
if rev_xfer_out[isub].size > 0:
transfers[vec_name]['rev', isub] = DefaultTransfer(
vectors['input'][vec_name],
vectors['output'][vec_name], rev_xfer_in[isub],
rev_xfer_out[isub], group.comm)
else:
transfers[vec_name]['rev', isub] = None
if group._use_derivatives:
transfers['nonlinear'] = transfers['linear']
def _initialize(self, system):
"""
Allocate the global matrices.
Parameters
----------
system : System
Parent system to this jacobian.
"""
# var_indices are the *global* indices for variables on this proc
is_top = system.pathname == ''
abs2meta = system._var_abs2meta
all_meta = system._var_allprocs_abs2meta
self._int_mtx = int_mtx = self._matrix_class(system.comm, True)
ext_mtx = self._matrix_class(system.comm, False)
iproc = system.comm.rank
abs2idx = system._var_allprocs_abs2idx['nonlinear']
in_sizes = system._var_sizes['nonlinear']['input']
out_ranges = self._out_ranges
# for non-MPI case, global_out_ranges is out_ranges
global_out_ranges = self._nodup_out_ranges
in_ranges = self._in_ranges
abs2prom_out = system._var_abs2prom['output']
owns = system._owning_rank
conns = {} if isinstance(system, Component) else system._conn_global_abs_in2out
abs_key2shape = self._abs_key2shape
check_owns = system._use_owned_sizes()
# create the matrix subjacs
for abs_key, info in iteritems(self._subjacs_info):
res_abs_name, wrt_abs_name = abs_key
# because self._subjacs_info is shared among all 'related' assembled jacs,
# we use global_out_ranges (and later in_ranges) to weed out keys outside of this jac
if res_abs_name not in global_out_ranges:
continue
res_offset, res_end = global_out_ranges[res_abs_name]
res_size = res_end - res_offset
if wrt_abs_name in abs2prom_out:
if check_owns and (owns[wrt_abs_name] != iproc and
not all_meta[wrt_abs_name]['distributed']):
continue
out_offset, out_end = global_out_ranges[wrt_abs_name]
out_size = out_end - out_offset
shape = (res_size, out_size)
int_mtx._add_submat(abs_key, info, res_offset, out_offset, None, shape)
elif wrt_abs_name in in_ranges:
if wrt_abs_name in conns: # connected input
out_abs_name = conns[wrt_abs_name]
if out_abs_name not in global_out_ranges:
continue
if check_owns and (owns[wrt_abs_name] != iproc and
not (all_meta[wrt_abs_name]['distributed'] or
all_meta[out_abs_name]['distributed'])):
continue
meta_in = abs2meta[wrt_abs_name]
all_out_meta = all_meta[out_abs_name]
# calculate unit conversion
in_units = meta_in['units']
out_units = all_out_meta['units']
if in_units and out_units and in_units != out_units:
factor, _ = get_conversion(out_units, in_units)
if factor == 1.0:
factor = None
else:
factor = None
out_offset, out_end = global_out_ranges[out_abs_name]
out_size = out_end - out_offset
shape = (res_size, out_size)
src_indices = abs2meta[wrt_abs_name]['src_indices']
if src_indices is not None:
# need to add an entry for d(output)/d(source)
# instead of d(output)/d(input). int_mtx is a square matrix whose
# rows and columns map to output/resid vars only.
abs_key2 = (res_abs_name, out_abs_name)
shape = abs_key2shape(abs_key2)
if len(src_indices.shape) > 1:
src_indices = _flatten_src_indices(src_indices, meta_in['shape'],
all_out_meta['global_shape'],
all_out_meta['global_size'])
int_mtx._add_submat(abs_key, info, res_offset, out_offset,
src_indices, shape, factor)
elif not is_top: # input is connected to something outside current system
in_offset, in_end = in_ranges[wrt_abs_name]
# don't use global offsets for ext_mtx
res_offset, res_end = out_ranges[res_abs_name]
res_size = res_end - res_offset
shape = (res_size, in_end - in_offset)
ext_mtx._add_submat(abs_key, info, res_offset, in_offset, None, shape)
out_size = system._outputs._data.size
if check_owns:
total_non_dup_size = np.sum(system._owned_sizes)
self._int_mtx._build(total_non_dup_size, total_non_dup_size, system)
else:
int_mtx._build(out_size, out_size, system)
if ext_mtx._submats:
in_size = np.sum(in_sizes[iproc, :])
ext_mtx._build(out_size, in_size)
else:
ext_mtx = None
self._ext_mtx[system.pathname] = ext_mtx
def _initialize(self, system):
"""
Allocate the global matrices.
Parameters
----------
system : System
Parent system to this jacobian.
"""
# var_indices are the *global* indices for variables on this proc
is_top = system.pathname == ''
abs2meta = system._var_abs2meta
all_meta = system._var_allprocs_abs2meta
self._int_mtx = int_mtx = self._matrix_class(system.comm)
ext_mtx = self._matrix_class(system.comm)
iproc = system.comm.rank
abs2idx = system._var_allprocs_abs2idx['nonlinear']
out_sizes = system._var_sizes['nonlinear']['output']
in_sizes = system._var_sizes['nonlinear']['input']
out_ranges = self._out_ranges
in_ranges = self._in_ranges
abs2prom_out = system._var_abs2prom['output']
conns = {} if isinstance(system, Component) else system._conn_global_abs_in2out
keymap = self._keymap
abs_key2shape = self._abs_key2shape
# create the matrix subjacs
for abs_key, info in iteritems(self._subjacs_info):
res_abs_name, wrt_abs_name = abs_key
# because self._subjacs_info is shared among all 'related' assembled jacs,
# we use out_ranges (and later in_ranges) to weed out keys outside of this jac
if res_abs_name not in out_ranges:
continue
res_offset, _ = out_ranges[res_abs_name]
if wrt_abs_name in abs2prom_out:
out_offset, _ = out_ranges[wrt_abs_name]
int_mtx._add_submat(abs_key, info, res_offset, out_offset, None, info['shape'])
keymap[abs_key] = abs_key
elif wrt_abs_name in in_ranges:
if wrt_abs_name in conns: # connected input
out_abs_name = conns[wrt_abs_name]
meta_in = abs2meta[wrt_abs_name]
all_out_meta = all_meta[out_abs_name]
# calculate unit conversion
in_units = meta_in['units']
out_units = all_out_meta['units']
if in_units and out_units and in_units != out_units:
factor, _ = get_conversion(out_units, in_units)
if factor == 1.0:
factor = None
else:
factor = None
out_offset, _ = out_ranges[out_abs_name]
src_indices = abs2meta[wrt_abs_name]['src_indices']
# need to add an entry for d(output)/d(source)
# instead of d(output)/d(input)
abs_key2 = (res_abs_name, out_abs_name)
keymap[abs_key] = abs_key2
if src_indices is None:
shape = info['shape']
else:
shape = abs_key2shape(abs_key2)
if len(src_indices.shape) > 1:
src_indices = _flatten_src_indices(src_indices, meta_in['shape'],
all_out_meta['global_shape'],
all_out_meta['global_size'])
int_mtx._add_submat(abs_key, info, res_offset, out_offset,
src_indices, shape, factor)
elif not is_top: # input is connected to something outside current system
ext_mtx._add_submat(abs_key, info, res_offset,
in_ranges[wrt_abs_name][0], None, info['shape'])
iproc = system.comm.rank
out_size = np.sum(out_sizes[iproc, :])
int_mtx._build(out_size, out_size, in_ranges, out_ranges)
if ext_mtx._submats:
in_size = np.sum(in_sizes[iproc, :])
ext_mtx._build(out_size, in_size, in_ranges, out_ranges)
else:
ext_mtx = None
self._ext_mtx[system.pathname] = ext_mtx
def _setup_transfers(group, recurse=True):
"""
Compute all transfers that are owned by our parent group.
Parameters
----------
group : <Group>
Parent group.
recurse : bool
Whether to call this method in subsystems.
"""
group._transfers = {}
rev = group._mode == 'rev' or group._mode == 'auto'
def merge(indices_list):
if len(indices_list) > 0:
return np.concatenate(indices_list)
else:
return _empty_idx_array
if recurse:
for subsys in group._subgroups_myproc:
subsys._setup_transfers(recurse)
# Pre-compute map from abs_names to the index of the containing subsystem
abs2isub = {}
for subsys, isub in zip(group._subsystems_myproc, group._subsystems_myproc_inds):
for type_ in ['input', 'output']:
for abs_name in subsys._var_allprocs_abs_names[type_]:
abs2isub[abs_name] = isub
abs2meta = group._var_abs2meta
allprocs_abs2meta = group._var_allprocs_abs2meta
myproc = group.comm.rank
transfers = group._transfers
vectors = group._vectors
offsets = group._get_var_offsets()
vec_names = group._lin_rel_vec_name_list if group._use_derivatives else group._vec_names
for vec_name in vec_names:
relvars, _ = group._relevant[vec_name]['@all']
# Initialize empty lists for the transfer indices
nsub_allprocs = len(group._subsystems_allprocs)
xfer_in = []
xfer_out = []
fwd_xfer_in = [[] for i in range(nsub_allprocs)]
fwd_xfer_out = [[] for i in range(nsub_allprocs)]
if rev:
rev_xfer_in = [[] for i in range(nsub_allprocs)]
rev_xfer_out = [[] for i in range(nsub_allprocs)]
allprocs_abs2idx = group._var_allprocs_abs2idx[vec_name]
sizes_in = group._var_sizes[vec_name]['input']
sizes_out = group._var_sizes[vec_name]['output']
offsets_in = offsets[vec_name]['input']
offsets_out = offsets[vec_name]['output']
# Loop through all connections owned by this system
for abs_in, abs_out in iteritems(group._conn_abs_in2out):
if abs_out not in relvars['output']:
continue
# Only continue if the input exists on this processor
if abs_in in abs2meta and abs_in in relvars['input']:
# Get meta
meta_in = abs2meta[abs_in]
meta_out = allprocs_abs2meta[abs_out]
idx_in = allprocs_abs2idx[abs_in]
idx_out = allprocs_abs2idx[abs_out]
# Read in and process src_indices
src_indices = meta_in['src_indices']
if src_indices is None:
owner = group._owning_rank[abs_out]
# if the input is larger than the output on a single proc, we have
# to just loop over the procs in the same way we do when src_indices
# is defined.
if meta_in['size'] > sizes_out[owner, idx_out]:
src_indices = np.arange(meta_in['size'], dtype=INT_DTYPE)
elif src_indices.ndim == 1:
src_indices = convert_neg(src_indices, meta_out['global_size'])
else:
src_indices = _flatten_src_indices(src_indices, meta_in['shape'],
meta_out['global_shape'],
meta_out['global_size'])
# 1. Compute the output indices
if src_indices is None:
start = 0 if owner == 0 else np.sum(sizes_out[:owner, idx_out])
offset = offsets_out[owner, idx_out] - start
output_inds = np.arange(offset, offset + meta_in['size'], dtype=INT_DTYPE)
else:
output_inds = np.zeros(src_indices.size, INT_DTYPE)
start = end = 0
for iproc in range(group.comm.size):
end += sizes_out[iproc, idx_out]
# The part of src on iproc
on_iproc = np.logical_and(start <= src_indices, src_indices < end)
if np.any(on_iproc):
# This converts from iproc-then-ivar to ivar-then-iproc ordering
# Subtract off part of previous procs
# Then add all variables on previous procs
# Then all previous variables on this proc
# - np.sum(out_sizes[:iproc, idx_out])
# + np.sum(out_sizes[:iproc, :])
# + np.sum(out_sizes[iproc, :idx_out])
# + inds
offset = offsets_out[iproc, idx_out] - start
output_inds[on_iproc] = src_indices[on_iproc] + offset
start = end
# 2. Compute the input indices
input_inds = np.arange(offsets_in[myproc, idx_in],
offsets_in[myproc, idx_in] +
sizes_in[myproc, idx_in], dtype=INT_DTYPE)
# Now the indices are ready - input_inds, output_inds
xfer_in.append(input_inds)
xfer_out.append(output_inds)
isub = abs2isub[abs_in]
fwd_xfer_in[isub].append(input_inds)
fwd_xfer_out[isub].append(output_inds)
if rev and abs_out in abs2isub:
isub = abs2isub[abs_out]
rev_xfer_in[isub].append(input_inds)
rev_xfer_out[isub].append(output_inds)
xfer_in = merge(xfer_in)
xfer_out = merge(xfer_out)
for isub in range(nsub_allprocs):
fwd_xfer_in[isub] = merge(fwd_xfer_in[isub])
fwd_xfer_out[isub] = merge(fwd_xfer_out[isub])
if rev:
rev_xfer_in[isub] = merge(rev_xfer_in[isub])
rev_xfer_out[isub] = merge(rev_xfer_out[isub])
out_vec = vectors['output'][vec_name]
transfers[vec_name] = {}
xfer_all = PETScTransfer(vectors['input'][vec_name], out_vec,
xfer_in, xfer_out, group.comm)
transfers[vec_name]['fwd', None] = xfer_all
if rev:
transfers[vec_name]['rev', None] = xfer_all
for isub in range(nsub_allprocs):
transfers[vec_name]['fwd', isub] = PETScTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
fwd_xfer_in[isub], fwd_xfer_out[isub], group.comm)
if rev:
transfers[vec_name]['rev', isub] = PETScTransfer(
vectors['input'][vec_name], vectors['output'][vec_name],
rev_xfer_in[isub], rev_xfer_out[isub], group.comm)
if group._use_derivatives:
transfers['nonlinear'] = transfers['linear']