def test_common_subpath(self):
plists = [
([], ''),
(['foo.bar.baz'], 'foo.bar.baz'),
(['a.b.c', 'a.b'], 'a.b'),
(['a.b', 'a.b.c'], 'a.b'),
(['foo.foo', 'foo.bar'], 'foo'),
(['foo.foo', 'bar'], ''),
(['xx.yy.zz.foo.bar', 'xx.yy.zz.blah',
'xx.yy.zz.blah.ho'], 'xx.yy.zz'),
]
for plist, ans in plists:
self.assertEqual(common_subpath(plist), ans)
def view_dyn_shapes(root,
outfile='shape_dep_graph.png',
show=True,
title=None):
"""
Generate a plot file containing the dynamic shape dependency graph.
Optionally displays the plot.
Parameters
----------
root : system or Problem
The top level system or Problem.
outfile : str, optional
The name of the plot file. Defaults to 'shape_dep_graph.png'.
show : bool, optional
If True, display the plot. Defaults to True.
title : str, optional
Sets the title of the plot.
"""
if MPI and MPI.COMM_WORLD.rank != 0:
return
if isinstance(root, Problem):
system = root.model
else:
system = root
if root.pathname != '':
raise RuntimeError(
"view_dyn_shapes cannot be called on a subsystem of the model. "
"Call it with the Problem or the model.")
try:
import matplotlib.pyplot as plt
except ImportError:
raise RuntimeError("The view_dyn_shapes command requires matplotlib.")
graph = system._shapes_graph
knowns = system._shape_knowns
if graph is None:
raise RuntimeError(
"Can't plot dynamic shape dependency graph because it hasn't been "
"computed yet. view_dyn_shapes must be called after problem setup()."
)
if graph.order() == 0:
print("The model has no dynamically shaped variables.")
return
if title is None:
# keep the names from being super long by removing any common subpath
common = common_subpath(graph.nodes())
if common:
title = f"Dynamic shape dependencies in group '{common}'"
common_idx = len(common) + 1 if common else 0
else:
title = "Dynamic shape dependencies"
common_idx = 0
abs2meta = system._var_allprocs_abs2meta
# label variables with known shape at the start of the algorithm in green, unknowns in red.
# prepend the shape onto the variable name
node_colors = []
node_labels = {}
for n in graph:
meta = abs2meta['input'][n] if n in abs2meta['input'] else abs2meta[
'output'][n]
shape = meta['shape']
if shape is None:
shape = '?'
node_colors.append('red')
else:
if meta.get('shape_by_conn', False) or meta.get(
'copy_shape', False):
node_colors.append('blue')
else:
node_colors.append('green')
node_labels[n] = f"{shape}: {n[common_idx:]}"
nx.draw_networkx(graph,
with_labels=True,
node_color=node_colors,
labels=node_labels)
plt.axis('off') # turn of axis
plt.title(title)
plt.savefig(outfile)
if show:
plt.show()