From 8d8ff0e3d9a7b8a209bc8842490c8ca5a596933f Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Thu, 21 Mar 2024 12:00:08 -0400 Subject: [PATCH 01/15] added graph plotting for dataflow and system tree using pydot --- openmdao/api.py | 1 + openmdao/core/component.py | 13 ++ openmdao/core/group.py | 373 +++++++++++++++++++++++++++++++++- openmdao/core/system.py | 14 ++ openmdao/utils/graph_utils.py | 105 ++++++++++ openmdao/utils/om.py | 8 +- 6 files changed, 508 insertions(+), 6 deletions(-) diff --git a/openmdao/api.py b/openmdao/api.py index ab3138f6f1..1b63cace7c 100644 --- a/openmdao/api.py +++ b/openmdao/api.py @@ -65,6 +65,7 @@ from openmdao.utils.spline_distributions import cell_centered from openmdao.utils.spline_distributions import sine_distribution from openmdao.utils.spline_distributions import node_centered +from openmdao.utils.graph_utils import write_graph # Vectors from openmdao.vectors.default_vector import DefaultVector diff --git a/openmdao/core/component.py b/openmdao/core/component.py index 64fd6fdece..3a9c4cbfa5 100644 --- a/openmdao/core/component.py +++ b/openmdao/core/component.py @@ -1766,6 +1766,19 @@ def _has_fast_rel_lookup(self): """ return True + def _get_graph_node_meta(self): + """ + Return metadata to add to this system's graph node. + + Returns + ------- + dict + Metadata for this system's graph node. + """ + meta = super()._get_graph_node_meta() + meta['base'] = 'ExplicitComponent' if self.is_explicit() else 'ImplicitComponent' + return meta + class _DictValues(object): """ diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 94da4ac205..35258e6b82 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -30,7 +30,7 @@ meta2src_iter, get_rev_conns, _contains_all from openmdao.utils.units import is_compatible, unit_conversion, _has_val_mismatch, _find_unit, \ _is_unitless, simplify_unit -from openmdao.utils.graph_utils import get_sccs_topo, get_out_of_order_nodes +from openmdao.utils.graph_utils import get_out_of_order_nodes, write_graph from openmdao.utils.mpi import MPI, check_mpi_exceptions, multi_proc_exception_check import openmdao.utils.coloring as coloring_mod from openmdao.utils.indexer import indexer, Indexer @@ -774,7 +774,7 @@ def _setup(self, comm, prob_meta): # determine which connections are managed by which group, and check validity of connections self._setup_connections() - def _get_dataflow_graph(self): + def _get_dataflow_graph(self, include_display_data=False): """ Return a graph of all variables and components in the model. @@ -4215,6 +4215,307 @@ def compute_sys_graph(self, comps_only=False, add_edge_info=True): return graph + def _get_graph_display_info(self, display_map=None): + base_display_map = { + 'ExplicitComponent': { + 'fillcolor': '"aquamarine3:aquamarine"', + 'style': 'filled', + 'shape': 'box', + }, + 'ImplicitComponent': { + 'fillcolor': '"lightblue:lightslateblue"', + 'style': 'filled', + 'shape': 'box', + }, + 'IndepVarComp': { + 'fillcolor': '"chartreuse2:chartreuse4"', + 'style': 'filled', + 'shape': 'box', + }, + 'Group': { + 'fillcolor': 'gray75', + 'style': 'filled', + 'shape': 'octagon', + }, + } + + node_info = {} + for s in self.system_iter(recurse=True, include_self=True): + meta = s._get_graph_node_meta() + if display_map and 'classname' in meta and meta['classname'] in display_map: + meta.update(display_map[meta['classname']]) + elif display_map and 'base' in meta and meta['base'] in display_map: + meta.update(display_map[meta['base']]) + elif 'base' in meta and meta['base'] in base_display_map: + meta.update(base_display_map[meta['base']]) + node_info[s.pathname] = meta.copy() + + if self.comm.size > 1: + abs2prom = self._var_abs2prom + all_abs2prom = self._var_allprocs_abs2prom + if (len(all_abs2prom['input']) != len(abs2prom['input']) or + len(all_abs2prom['output']) != len(abs2prom['output'])): + # not all systems exist in all procs, so must gather info from all procs + if self._gather_full_data(): + all_node_info = self.comm.allgather(node_info) + else: + all_node_info = self.comm.allgather({}) + + for info in all_node_info: + for pathname, meta in info.items(): + if pathname not in node_info: + node_info[pathname] = meta + + return node_info + + def _get_graph_node_meta(self): + """ + Return metadata to add to this system's graph node. + + Returns + ------- + dict + Metadata for this system's graph node. + """ + meta = super()._get_graph_node_meta() + # TODO: maybe set 'implicit' based on whether there are any implicit comps anywhere + # inside of the group or its children. + meta['base'] = 'Group' + return meta + + def _get_cluster_tree(self, node_info): + """ + Create a nested collection of pydot Cluster objects to represent the tree of groups. + + Parameters + ---------- + node_info : dict + A dict of metadata keyed by pathname. + + Returns + ------- + pydot.Dot, dict + The pydot graph and a dict of groups keyed by pathname. + """ + try: + import pydot + except ImportError: + issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") + return + + groups = {} + pydot_graph = pydot.Dot(graph_type='digraph') + prefix = self.pathname + '.' if self.pathname else '' + for varpath in chain(self._var_allprocs_abs2prom['input'], + self._var_allprocs_abs2prom['output']): + group = varpath.rpartition('.')[0].rpartition('.')[0] + if group not in groups: + # reverse the list so parents will exist before children + ancestor_list = list(all_ancestors(group))[::-1] + for path in ancestor_list: + if path.startswith(prefix): + if path not in groups: + parent, _, name = path.rpartition('.') + if path in node_info: + ttip = f"{node_info[path]['classname']} {path}" + else: + ttip = path + groups[path] = pydot.Cluster(path, label=name, tooltip=ttip, + fillcolor=_cluster_color(path), + style='filled') + if parent and parent.startswith(prefix): + groups[parent].add_subgraph(groups[path]) + else: + pydot_graph.add_subgraph(groups[path]) + + return pydot_graph, groups + + def _get_tree_graph(self, display_map=None): + """ + Create a pydot graph of the system tree (without clusters). + + Parameters + ---------- + display_map : dict or None + A map of classnames to pydot node attributes. + + Returns + ------- + pydot.Dot + The pydot tree graph. + """ + try: + import pydot + except ImportError: + issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") + return + + node_info = self._get_graph_display_info(display_map) + + systems = {} + pydot_graph = pydot.Dot(graph_type='graph') + prefix = self.pathname + '.' if self.pathname else '' + label = self.name if self.name else 'Model' + top_node = pydot.Node(label, label=label, + tooltip=f"{node_info[self.pathname]['classname']} {self.pathname}", + fillcolor='gray95', style='filled') + pydot_graph.add_node(top_node) + systems[self.pathname] = top_node + + for varpath in chain(self._var_allprocs_abs2prom['input'], + self._var_allprocs_abs2prom['output']): + system = varpath.rpartition('.')[0] + if system not in systems: + # reverse the list so parents will exist before children + ancestor_list = list(all_ancestors(system))[::-1] + for path in ancestor_list: + if path.startswith(prefix): + if path not in systems: + parent, _, name = path.rpartition('.') + kwargs = _filter_meta4dot(node_info[path]) + ttip = f"{node_info[path]['classname']} {path}" + systems[path] = pydot.Node(path, label=name, tooltip=ttip, **kwargs) + pydot_graph.add_node(systems[path]) + if parent.startswith(prefix) or parent == self.pathname: + pydot_graph.add_edge(pydot.Edge(systems[parent], systems[path])) + + return pydot_graph + + def _decorate_graph_for_display(self, G, exclude=()): + try: + import pydot + except ImportError: + issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") + return G, {} + + node_info = self._get_graph_display_info() + + exclude = set(exclude) + + # dot doesn't like ':' in node names, so if we find any, we have to put explicit quotes + # around the node name and label. + replace = {} + for node, meta in G.nodes(data=True): + if node in node_info: + meta.update(_filter_meta4dot(node_info[node])) + meta['tooltip'] = f"{node_info[node]['classname']} {node}" + quoted = node.rpartition('.')[2] + meta['label'] = f'"{quoted}"' + if 'type_' in meta: # variable node + if node.rpartition('.')[0] in exclude: + exclude.add(node) # remove all variables of excluded components + if ':' in node and node not in exclude: + replace[node] = node.replace(':', ';') + + if replace: + G = nx.relabel_nodes(G, replace, copy=True) + + if exclude: + if not replace: + G = G.copy() + G.remove_nodes_from(exclude) + + return G, node_info + + def _apply_clusters(self, G, node_info): + try: + import pydot + except ImportError: + issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") + return G + + pydot_graph, groups = self._get_cluster_tree(node_info) + prefix = self.pathname + '.' if self.pathname else '' + pydot_nodes = {} + skip = {'type_', 'local', 'base'} + for node, meta in G.nodes(data=True): + kwargs = {n:v for n, v in meta.items() if n not in skip} + if 'type_' in meta: # variable node + group = node.rpartition('.')[0].rpartition('.')[0] + kwargs['shape'] = 'plain' # just text for variables, otherwise too busy + else: + group = node.rpartition('.')[0] + + pdnode = pydot_nodes[node] = pydot.Node(node, **kwargs) + + if group and group.startswith(prefix): + groups[group].add_node(pdnode) + else: + pydot_graph.add_node(pdnode) + + for u, v in G.edges(): + node1 = pydot_nodes[u] + node2 = pydot_nodes[v] + pydot_graph.add_edge(pydot.Edge(node1, node2, arrowhead='lnormal')) + + # layout graph from left to right + pydot_graph.set_rankdir('LR') + + return pydot_graph + + def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, + display=True, exclude=(), outfile=None): + """ + Use pydot to create a graphical representation of the specified graph. + + Parameters + ---------- + gtype : str + The type of graph to create. Options include 'system', 'component', 'nested', + and 'dataflow'. + recurse : bool + If True, recurse into subsystems when gtype is 'dataflow'. + show_vars : bool + If True, show all variables in the graph. Only relevant when gtype is 'dataflow'. + display : bool + If True, pop up a window to view the graph. + exclude : iter of str + Iter of pathnames to exclude from the generated graph. + outfile : str or None + The name of the file to write the graph to. The format is inferred from the extension. + Default is None, which writes to '__graph.svg', where system_path + is 'model' for the top level group, and any '.' in the pathname is replaced with '_'. + + Returns + ------- + pydot.Dot or None + The pydot graph that was created. + """ + if gtype == 'tree': + G = self._get_tree_graph() + elif gtype == 'dataflow': + if show_vars: + if self.pathname == '': + G = self._dataflow_graph + else: + G = self._problem_meta['model_ref']()._dataflow_graph + # we're not the top level group, so get our subgraph of the top level graph + prefix = self.pathname + '.' + ournodes = {n for n in G.nodes() if n.startswith(prefix)} + G = nx.subgraph(G, ournodes) + + G, node_info = self._decorate_graph_for_display(G, exclude=exclude) + G = self._apply_clusters(G, node_info) + elif recurse: + G = self.compute_sys_graph(comps_only=True, add_edge_info=False) + G, node_info = self._decorate_graph_for_display(G, exclude=exclude) + G = self._apply_clusters(G, node_info) + else: + G = self.compute_sys_graph(comps_only=False, add_edge_info=False) + G, _ = self._decorate_graph_for_display(G, exclude=exclude) + else: + raise ValueError(f"unrecognized graph type '{gtype}'. Allowed types are ['tree', " + "'dataflow'].") + + if G is None: + return + + if outfile is None: + name = self.pathname.replace('.', '_') if self.pathname else 'model' + outfile = f"{name}_{gtype}_graph.svg" + + return write_graph(G, prog='dot', display=display, outfile=outfile) + def _get_auto_ivc_out_val(self, tgts, vars_to_gather): # all tgts are continuous variables # only called from top level group @@ -5186,3 +5487,71 @@ def _active_responses(self, user_response_names, responses=None): meta['remote'] = meta['source'] not in self._var_abs2meta['output'] return active_resps + + +def _vars2groups(varnameiter): + """ + Return a set of all groups containing the given variables. + + Parameters + ---------- + varnameiter : iter of str + Iterator of variable pathnames. + + Returns + ------- + set + Set of group pathnames. + """ + groups = {} + for name in varnameiter: + gname = name.rpartition('.')[0].rpartition('.')[0] + if gname not in groups: + groups.update(all_ancestors(gname)) + + return groups + + +def _cluster_color(path): + """ + Return the color of the cluster that contains the given path. + + The idea here is to make nested clusters stand out wrt their parent cluster. + + Parameters + ---------- + path : str + Pathname of a variable. + + Returns + ------- + int + The color of the cluster that contains the given path. + """ + depth = path.count('.') + 1 if path else 0 + + ncols = 10 + maxcol = 98 + mincol = 40 + + # allow 8 nesting levels of difference + col = maxcol - (depth % ncols) * (maxcol - mincol) // ncols + return f"gray{col}" + + +def _filter_meta4dot(meta): + """ + Remove unnecessary metadata from the given metadata dict before passing to pydot. + + Parameters + ---------- + meta : dict + Metadata dict. + + Returns + ------- + dict + Metadata dict with unnecessary items removed. + """ + skip = {'type_', 'local', 'base', 'classname'} + return {k: v for k, v in meta.items() if k not in skip} diff --git a/openmdao/core/system.py b/openmdao/core/system.py index 8c340a7e3f..286707d161 100644 --- a/openmdao/core/system.py +++ b/openmdao/core/system.py @@ -1380,6 +1380,20 @@ def get_source(self, name): raise KeyError(f"{self.msginfo}: source for '{name}' not found.") + def _get_graph_node_meta(self): + """ + Return metadata to add to this system's graph node. + + Returns + ------- + dict + Metadata for this system's graph node. + """ + return { + 'classname': type(self).__name__, + 'implicit': not self.is_explicit(), + } + def _setup_check(self): """ Do any error checking on user's setup, before any other recursion happens. diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index 1c04fa0538..e42d6798e5 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -2,6 +2,8 @@ Various graph related utilities. """ import networkx as nx +from openmdao.utils.file_utils import _load_and_exec +import openmdao.utils.hooks as hooks def get_sccs_topo(graph): @@ -54,3 +56,106 @@ def get_out_of_order_nodes(graph, orders): out_of_order.append((u, v)) return strongcomps, out_of_order + + +def write_graph(G, prog='dot', display=True, outfile='graph.svg'): + """ + Write the graph to a file and optionally display it. + + Parameters + ---------- + G : nx.DiGraph or pydot.Dot + The graph to be written. + prog : str + The graphviz program to use for layout. + display : bool + If True, display the graph after writing it. + outfile : str + The name of the file to write. + + Returns + ------- + pydot.Dot + The graph that was written. + """ + from openmdao.utils.webview import webview + + try: + import pydot + except ImportError: + raise RuntimeError("graph requires the pydot package. You can install it using " + "'pip install pydot'.") + + ext = outfile.rpartition('.')[2] + + if isinstance(G, nx.Graph): + pydot_graph = nx.drawing.nx_pydot.to_pydot(G) + else: + pydot_graph = G + + try: + pstr = getattr(pydot_graph, f"create_{ext}")(prog=prog) + except AttributeError: + raise AttributeError(f"pydot graph has no 'create_{ext}' method.") + + with open(outfile, 'wb') as f: + f.write(pstr) + + if display: + webview(outfile) + + return pydot_graph + + +def _graph_setup_parser(parser): + """ + Set up the openmdao subparser for the 'openmdao graph' command. + + Parameters + ---------- + parser : argparse subparser + The parser we're adding options to. + """ + parser.add_argument('file', nargs=1, help='Python file containing the model.') + parser.add_argument('-p', '--problem', action='store', dest='problem', help='Problem name') + parser.add_argument('-o', action='store', dest='outfile', help='file containing graph output.') + parser.add_argument('--group', action='store', dest='group', help='pathname of group to graph.') + parser.add_argument('--type', action='store', dest='type', default='dataflow', + help='type of graph (dataflow, tree). Default is dataflow.') + parser.add_argument('--no-display', action='store_false', dest='show', + help="don't display the graph.") + parser.add_argument('--no-recurse', action='store_false', dest='recurse', + help="don't recurse from the specified group down. This only applies to " + "the dataflow graph type.") + parser.add_argument('--show-vars', action='store_true', dest='show_vars', + help="show variables in the graph. This only applies to the dataflow graph." + " Default is False.") + parser.add_argument('--autoivc', action='store_true', dest='auto_ivc', + help="include the _auto_ivc component in the graph. This applies to " + "graphs of the top level group only. Default is False.") + + +def _graph_cmd(options, user_args): + """ + Return the post_setup hook function for 'openmdao graph'. + + Parameters + ---------- + options : argparse Namespace + Command line options. + user_args : list of str + Args to be passed to the user script. + """ + def _view_graph(problem): + group = problem.model._get_subsystem(options.group) if options.group else problem.model + if not options.group and not options.auto_ivc: + exclude = {'_auto_ivc'} + else: + exclude = set() + group.write_graph(gtype=options.type, recurse=options.recurse, + show_vars=options.show_vars, display=options.show, exclude=exclude, + outfile=options.outfile) + + # register the hooks + hooks._register_hook('final_setup', 'Problem', post=_view_graph, exit=True) + _load_and_exec(options.file[0], user_args) \ No newline at end of file diff --git a/openmdao/utils/om.py b/openmdao/utils/om.py index 0e8c7f4464..9cd6b5075e 100644 --- a/openmdao/utils/om.py +++ b/openmdao/utils/om.py @@ -61,6 +61,7 @@ _find_repos_setup_parser, _find_repos_exec from openmdao.utils.reports_system import _list_reports_setup_parser, _list_reports_cmd, \ _view_reports_setup_parser, _view_reports_cmd +from openmdao.utils.graph_utils import _graph_setup_parser, _graph_cmd def _view_connections_setup_parser(parser): @@ -538,14 +539,14 @@ def _set_dyn_hook(prob): 'Display connection information for variables across multiple MPI processes.'), 'find_repos': (_find_repos_setup_parser, _find_repos_exec, 'Find repos on github having openmdao topics.'), + 'graph': (_graph_setup_parser, _graph_cmd, 'Generate a graph for a group.'), 'iprof': (_iprof_setup_parser, _iprof_exec, 'Profile calls to particular object instances.'), 'iprof_totals': (_iprof_totals_setup_parser, _iprof_totals_exec, 'Generate total timings of calls to particular object instances.'), 'list_installed': (_list_installed_setup_parser, _list_installed_cmd, 'List installed types recognized by OpenMDAO.'), - 'list_reports': (_list_reports_setup_parser, _list_reports_cmd, - 'List available reports.'), + 'list_reports': (_list_reports_setup_parser, _list_reports_cmd, 'List available reports.'), 'mem': (_mem_prof_setup_parser, _mem_prof_exec, 'Profile memory used by OpenMDAO related functions.'), 'mempost': (_mempost_setup_parser, _mempost_exec, 'Post-process memory profile output.'), @@ -570,8 +571,7 @@ def _set_dyn_hook(prob): 'view_dyn_shapes': (_view_dyn_shapes_setup_parser, _view_dyn_shapes_cmd, 'View the dynamic shape dependency graph.'), 'view_mm': (_meta_model_parser, _meta_model_cmd, "View a metamodel."), - 'view_reports': (_view_reports_setup_parser, _view_reports_cmd, - 'View existing reports.'), + 'view_reports': (_view_reports_setup_parser, _view_reports_cmd, 'View existing reports.'), } From 9ac75e8070ff5d28f48ba226cd9a4828f993ec40 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Thu, 21 Mar 2024 13:06:35 -0400 Subject: [PATCH 02/15] cleanup: --- openmdao/api.py | 1 - openmdao/core/group.py | 34 +++++----------------------------- openmdao/utils/graph_utils.py | 2 ++ 3 files changed, 7 insertions(+), 30 deletions(-) diff --git a/openmdao/api.py b/openmdao/api.py index 1b63cace7c..ab3138f6f1 100644 --- a/openmdao/api.py +++ b/openmdao/api.py @@ -65,7 +65,6 @@ from openmdao.utils.spline_distributions import cell_centered from openmdao.utils.spline_distributions import sine_distribution from openmdao.utils.spline_distributions import node_centered -from openmdao.utils.graph_utils import write_graph # Vectors from openmdao.vectors.default_vector import DefaultVector diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 35258e6b82..39bd1e5741 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -774,7 +774,7 @@ def _setup(self, comm, prob_meta): # determine which connections are managed by which group, and check validity of connections self._setup_connections() - def _get_dataflow_graph(self, include_display_data=False): + def _get_dataflow_graph(self): """ Return a graph of all variables and components in the model. @@ -5489,29 +5489,6 @@ def _active_responses(self, user_response_names, responses=None): return active_resps -def _vars2groups(varnameiter): - """ - Return a set of all groups containing the given variables. - - Parameters - ---------- - varnameiter : iter of str - Iterator of variable pathnames. - - Returns - ------- - set - Set of group pathnames. - """ - groups = {} - for name in varnameiter: - gname = name.rpartition('.')[0].rpartition('.')[0] - if gname not in groups: - groups.update(all_ancestors(gname)) - - return groups - - def _cluster_color(path): """ Return the color of the cluster that contains the given path. @@ -5530,12 +5507,11 @@ def _cluster_color(path): """ depth = path.count('.') + 1 if path else 0 - ncols = 10 - maxcol = 98 - mincol = 40 + ncolors = 10 + maxcolor = 98 + mincolor = 40 - # allow 8 nesting levels of difference - col = maxcol - (depth % ncols) * (maxcol - mincol) // ncols + col = maxcolor - (depth % ncolors) * (maxcolor - mincolor) // ncolors return f"gray{col}" diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index e42d6798e5..4e75006760 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -87,6 +87,8 @@ def write_graph(G, prog='dot', display=True, outfile='graph.svg'): "'pip install pydot'.") ext = outfile.rpartition('.')[2] + if not ext: + ext = 'svg' if isinstance(G, nx.Graph): pydot_graph = nx.drawing.nx_pydot.to_pydot(G) From 43e0a02bbd5696acd68e076a0e7ffc74c0f16a55 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Thu, 21 Mar 2024 23:53:26 -0400 Subject: [PATCH 03/15] added non-recurse mode to dataflow with vars --- openmdao/core/group.py | 210 ++++++++++++------ .../test_suite/scripts/circuit_analysis.py | 2 + openmdao/utils/graph_utils.py | 3 +- 3 files changed, 141 insertions(+), 74 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 39bd1e5741..d84da08e2a 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -11,6 +11,11 @@ import numpy as np import networkx as nx +try: + import pydot +except ImportError: + pydot = None + from openmdao.core.configinfo import _ConfigInfo from openmdao.core.system import System, collect_errors from openmdao.core.component import Component, _DictValues @@ -44,6 +49,31 @@ namecheck_rgx = re.compile('[a-zA-Z][_a-zA-Z0-9]*') +# mapping of system type to graph display properties +_base_display_map = { + 'ExplicitComponent': { + 'fillcolor': '"aquamarine3:aquamarine"', + 'style': 'filled', + 'shape': 'box', + }, + 'ImplicitComponent': { + 'fillcolor': '"lightblue:lightslateblue"', + 'style': 'filled', + 'shape': 'box', + }, + 'IndepVarComp': { + 'fillcolor': '"chartreuse2:chartreuse4"', + 'style': 'filled', + 'shape': 'box', + }, + 'Group': { + 'fillcolor': 'gray75', + 'style': 'filled', + 'shape': 'octagon', + }, +} + + # use a class with slots instead of a namedtuple so that we can # change index after creation if needed. class _SysInfo(object): @@ -4216,38 +4246,16 @@ def compute_sys_graph(self, comps_only=False, add_edge_info=True): return graph def _get_graph_display_info(self, display_map=None): - base_display_map = { - 'ExplicitComponent': { - 'fillcolor': '"aquamarine3:aquamarine"', - 'style': 'filled', - 'shape': 'box', - }, - 'ImplicitComponent': { - 'fillcolor': '"lightblue:lightslateblue"', - 'style': 'filled', - 'shape': 'box', - }, - 'IndepVarComp': { - 'fillcolor': '"chartreuse2:chartreuse4"', - 'style': 'filled', - 'shape': 'box', - }, - 'Group': { - 'fillcolor': 'gray75', - 'style': 'filled', - 'shape': 'octagon', - }, - } - node_info = {} for s in self.system_iter(recurse=True, include_self=True): meta = s._get_graph_node_meta() - if display_map and 'classname' in meta and meta['classname'] in display_map: + if display_map and meta['classname'] in display_map: meta.update(display_map[meta['classname']]) - elif display_map and 'base' in meta and meta['base'] in display_map: + elif display_map and meta['base'] in display_map: meta.update(display_map[meta['base']]) - elif 'base' in meta and meta['base'] in base_display_map: - meta.update(base_display_map[meta['base']]) + elif meta['base'] in _base_display_map: + meta.update(_base_display_map[meta['base']]) + meta['tooltip'] = f"{meta['classname']} {s.pathname}" node_info[s.pathname] = meta.copy() if self.comm.size > 1: @@ -4297,12 +4305,6 @@ def _get_cluster_tree(self, node_info): pydot.Dot, dict The pydot graph and a dict of groups keyed by pathname. """ - try: - import pydot - except ImportError: - issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") - return - groups = {} pydot_graph = pydot.Dot(graph_type='digraph') prefix = self.pathname + '.' if self.pathname else '' @@ -4316,11 +4318,8 @@ def _get_cluster_tree(self, node_info): if path.startswith(prefix): if path not in groups: parent, _, name = path.rpartition('.') - if path in node_info: - ttip = f"{node_info[path]['classname']} {path}" - else: - ttip = path - groups[path] = pydot.Cluster(path, label=name, tooltip=ttip, + groups[path] = pydot.Cluster(path, label=name, + tooltip=node_info[path]['tooltip'], fillcolor=_cluster_color(path), style='filled') if parent and parent.startswith(prefix): @@ -4344,12 +4343,6 @@ def _get_tree_graph(self, display_map=None): pydot.Dot The pydot tree graph. """ - try: - import pydot - except ImportError: - issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") - return - node_info = self._get_graph_display_info(display_map) systems = {} @@ -4373,8 +4366,7 @@ def _get_tree_graph(self, display_map=None): if path not in systems: parent, _, name = path.rpartition('.') kwargs = _filter_meta4dot(node_info[path]) - ttip = f"{node_info[path]['classname']} {path}" - systems[path] = pydot.Node(path, label=name, tooltip=ttip, **kwargs) + systems[path] = pydot.Node(path, label=name, **kwargs) pydot_graph.add_node(systems[path]) if parent.startswith(prefix) or parent == self.pathname: pydot_graph.add_edge(pydot.Edge(systems[parent], systems[path])) @@ -4382,12 +4374,6 @@ def _get_tree_graph(self, display_map=None): return pydot_graph def _decorate_graph_for_display(self, G, exclude=()): - try: - import pydot - except ImportError: - issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") - return G, {} - node_info = self._get_graph_display_info() exclude = set(exclude) @@ -4398,14 +4384,14 @@ def _decorate_graph_for_display(self, G, exclude=()): for node, meta in G.nodes(data=True): if node in node_info: meta.update(_filter_meta4dot(node_info[node])) - meta['tooltip'] = f"{node_info[node]['classname']} {node}" - quoted = node.rpartition('.')[2] - meta['label'] = f'"{quoted}"' + quoted = f'"{node.rpartition(".")[2]}"' + meta['label'] = quoted if 'type_' in meta: # variable node if node.rpartition('.')[0] in exclude: exclude.add(node) # remove all variables of excluded components - if ':' in node and node not in exclude: - replace[node] = node.replace(':', ';') + if ':' in node and node not in exclude: # fix ':' in node names for use in dot + replace[node] = f'"{node}"' + meta['shape'] = 'plain' # just text for variables, otherwise too busy if replace: G = nx.relabel_nodes(G, replace, copy=True) @@ -4417,24 +4403,20 @@ def _decorate_graph_for_display(self, G, exclude=()): return G, node_info - def _apply_clusters(self, G, node_info): - try: - import pydot - except ImportError: - issue_warning("write_graph requires pydot. Install pydot using 'pip install pydot'.") - return G + def _add_boundary_nodes(self, G): + return G + def _apply_clusters(self, G, node_info): pydot_graph, groups = self._get_cluster_tree(node_info) prefix = self.pathname + '.' if self.pathname else '' pydot_nodes = {} - skip = {'type_', 'local', 'base'} for node, meta in G.nodes(data=True): - kwargs = {n:v for n, v in meta.items() if n not in skip} + noquote_node = node.strip('"') + kwargs = _filter_meta4dot(meta) if 'type_' in meta: # variable node - group = node.rpartition('.')[0].rpartition('.')[0] - kwargs['shape'] = 'plain' # just text for variables, otherwise too busy + group = noquote_node.rpartition('.')[0].rpartition('.')[0] else: - group = node.rpartition('.')[0] + group = noquote_node.rpartition('.')[0] pdnode = pydot_nodes[node] = pydot.Node(node, **kwargs) @@ -4453,8 +4435,33 @@ def _apply_clusters(self, G, node_info): return pydot_graph + def _get_boundary_conns(self): + """ + Return lists of incoming and outgoing boundary connections. + + Returns + ------- + tuple + A tuple of (incoming, outgoing) boundary connections. + """ + if not self.pathname: + return ([], []) + + top = self._problem_meta['model_ref']() + prefix = self.pathname + '.' + + incoming = [] + outgoing = [] + for abs_in, abs_out in top._conn_global_abs_in2out.items(): + if abs_in.startswith(prefix) and not abs_out.startswith(prefix): + incoming.append((abs_in, abs_out)) + if abs_out.startswith(prefix) and not abs_in.startswith(prefix): + outgoing.append((abs_in, abs_out)) + + return incoming, outgoing + def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, - display=True, exclude=(), outfile=None): + display=True, show_boundary=False, exclude=(), outfile=None): """ Use pydot to create a graphical representation of the specified graph. @@ -4469,6 +4476,8 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, If True, show all variables in the graph. Only relevant when gtype is 'dataflow'. display : bool If True, pop up a window to view the graph. + show_boundary : bool + If True, include connections to variables outside the boundary of the Group. exclude : iter of str Iter of pathnames to exclude from the generated graph. outfile : str or None @@ -4481,21 +4490,73 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, pydot.Dot or None The pydot graph that was created. """ + if pydot is None: + issue_warning("To view the system graph, you must install pydot. " + "You can install it via 'pip install pydot'.") + return + if gtype == 'tree': G = self._get_tree_graph() elif gtype == 'dataflow': if show_vars: - if self.pathname == '': + prefix = self.pathname + '.' if self.pathname else '' + lenpre = len(prefix) + + if self.pathname: + G = self._problem_meta['model_ref']()._dataflow_graph + else: G = self._dataflow_graph + + if not recurse: + # keep all direct children and their variables + keep = {n for n in G.nodes() if n[lenpre:].count('.') == 0} + keep.update({n for n, d in G.nodes(data=True) if 'type_' in d and + n.rpartition('.')[0] in keep}) + + # keep all variables involved in connections owned by this group + inconnvars = set() + outconnvars = set() + for abs_in, abs_out in self._conn_abs_in2out.items(): + if abs_in not in keep: + inconnvars.add(abs_in) + if abs_out not in keep: + outconnvars.add(abs_out) + + for invar in inconnvars: + grp = prefix + invar[lenpre:].partition('.')[0] + if grp not in G: + G.add_node(grp, **_base_display_map['Group']) + G.add_edge(invar, grp) + keep.add(grp) + keep.add(invar) + + for outvar in outconnvars: + grp = prefix + outvar[lenpre:].partition('.')[0] + if grp not in G: + G.add_node(grp, **_base_display_map['Group']) + G.add_edge(grp, outvar) + keep.add(grp) + keep.add(outvar) + + if self.pathname == '': + if not recurse: + G = nx.subgraph(G, keep) else: - G = self._problem_meta['model_ref']()._dataflow_graph # we're not the top level group, so get our subgraph of the top level graph - prefix = self.pathname + '.' ournodes = {n for n in G.nodes() if n.startswith(prefix)} + + if not recurse: + ournodes.update(keep) + G = nx.subgraph(G, ournodes) G, node_info = self._decorate_graph_for_display(G, exclude=exclude) - G = self._apply_clusters(G, node_info) + if recurse: + G = self._apply_clusters(G, node_info) + else: + # layout graph from left to right + G.graph['graph'] = {'rankdir': 'LR'} + elif recurse: G = self.compute_sys_graph(comps_only=True, add_edge_info=False) G, node_info = self._decorate_graph_for_display(G, exclude=exclude) @@ -4503,6 +4564,9 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, else: G = self.compute_sys_graph(comps_only=False, add_edge_info=False) G, _ = self._decorate_graph_for_display(G, exclude=exclude) + + if show_boundary: + G = self._add_boundary_nodes(G) else: raise ValueError(f"unrecognized graph type '{gtype}'. Allowed types are ['tree', " "'dataflow'].") diff --git a/openmdao/test_suite/scripts/circuit_analysis.py b/openmdao/test_suite/scripts/circuit_analysis.py index 6aec1eb105..8a3f95b44b 100644 --- a/openmdao/test_suite/scripts/circuit_analysis.py +++ b/openmdao/test_suite/scripts/circuit_analysis.py @@ -156,3 +156,5 @@ def setup(self): p['circuit.n2.V'] = .7 p.run_model() + + p.model.write_graph(gtype='dataflow', recurse=True, show_vars=True) diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index 4e75006760..559bf62334 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -160,4 +160,5 @@ def _view_graph(problem): # register the hooks hooks._register_hook('final_setup', 'Problem', post=_view_graph, exit=True) - _load_and_exec(options.file[0], user_args) \ No newline at end of file + _load_and_exec(options.file[0], user_args) + From ebeec446b9b890dda6a641e2f42dbf6fed50468b Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Fri, 22 Mar 2024 07:38:46 -0400 Subject: [PATCH 04/15] fixed graph layout orientation --- openmdao/core/group.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index d84da08e2a..6f38b7e5a1 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4508,6 +4508,9 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, G = self._dataflow_graph if not recurse: + # layout graph from left to right + G.graph['graph'] = {'rankdir': 'LR'} + # keep all direct children and their variables keep = {n for n in G.nodes() if n[lenpre:].count('.') == 0} keep.update({n for n, d in G.nodes(data=True) if 'type_' in d and @@ -4553,9 +4556,6 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, G, node_info = self._decorate_graph_for_display(G, exclude=exclude) if recurse: G = self._apply_clusters(G, node_info) - else: - # layout graph from left to right - G.graph['graph'] = {'rankdir': 'LR'} elif recurse: G = self.compute_sys_graph(comps_only=True, add_edge_info=False) From d32de33a4a2df5eb150353fb0b7a7c3c5a3b0071 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Fri, 22 Mar 2024 12:43:54 -0400 Subject: [PATCH 05/15] input/output vars to a collapsed group are not labeled with promoted names --- openmdao/core/group.py | 5 ++++- openmdao/test_suite/scripts/circuit_analysis.py | 2 -- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 6f38b7e5a1..49620d6468 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4385,7 +4385,8 @@ def _decorate_graph_for_display(self, G, exclude=()): if node in node_info: meta.update(_filter_meta4dot(node_info[node])) quoted = f'"{node.rpartition(".")[2]}"' - meta['label'] = quoted + if not ('label' in meta and meta['label']): + meta['label'] = quoted if 'type_' in meta: # variable node if node.rpartition('.')[0] in exclude: exclude.add(node) # remove all variables of excluded components @@ -4532,6 +4533,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, G.add_edge(invar, grp) keep.add(grp) keep.add(invar) + G.nodes[invar]['label'] = self._var_allprocs_abs2prom['input'][invar] for outvar in outconnvars: grp = prefix + outvar[lenpre:].partition('.')[0] @@ -4540,6 +4542,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, G.add_edge(grp, outvar) keep.add(grp) keep.add(outvar) + G.nodes[outvar]['label'] = self._var_allprocs_abs2prom['output'][outvar] if self.pathname == '': if not recurse: diff --git a/openmdao/test_suite/scripts/circuit_analysis.py b/openmdao/test_suite/scripts/circuit_analysis.py index 8a3f95b44b..6aec1eb105 100644 --- a/openmdao/test_suite/scripts/circuit_analysis.py +++ b/openmdao/test_suite/scripts/circuit_analysis.py @@ -156,5 +156,3 @@ def setup(self): p['circuit.n2.V'] = .7 p.run_model() - - p.model.write_graph(gtype='dataflow', recurse=True, show_vars=True) From 93a525ae66cdc2b0d1a4d4e626f1a480994e739b Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Mon, 25 Mar 2024 10:47:40 -0400 Subject: [PATCH 06/15] interim --- openmdao/core/group.py | 193 ++++++++++++++---- .../test_suite/scripts/circuit_analysis.py | 2 + openmdao/utils/graph_utils.py | 5 +- openmdao/utils/tests/test_cmdline.py | 7 + 4 files changed, 165 insertions(+), 42 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 49620d6468..3246376312 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4307,7 +4307,8 @@ def _get_cluster_tree(self, node_info): """ groups = {} pydot_graph = pydot.Dot(graph_type='digraph') - prefix = self.pathname + '.' if self.pathname else '' + mypath = self.pathname + prefix = mypath + '.' if mypath else '' for varpath in chain(self._var_allprocs_abs2prom['input'], self._var_allprocs_abs2prom['output']): group = varpath.rpartition('.')[0].rpartition('.')[0] @@ -4315,15 +4316,18 @@ def _get_cluster_tree(self, node_info): # reverse the list so parents will exist before children ancestor_list = list(all_ancestors(group))[::-1] for path in ancestor_list: - if path.startswith(prefix): + if path.startswith(prefix) or path==mypath: if path not in groups: parent, _, name = path.rpartition('.') - groups[path] = pydot.Cluster(path, label=name, + groups[path] = pydot.Cluster(path, + label=path if path == mypath else name, tooltip=node_info[path]['tooltip'], fillcolor=_cluster_color(path), style='filled') if parent and parent.startswith(prefix): groups[parent].add_subgraph(groups[path]) + elif parent == mypath and parent in groups: + groups[parent].add_subgraph(groups[path]) else: pydot_graph.add_subgraph(groups[path]) @@ -4346,7 +4350,7 @@ def _get_tree_graph(self, display_map=None): node_info = self._get_graph_display_info(display_map) systems = {} - pydot_graph = pydot.Dot(graph_type='graph') + pydot_graph = pydot.Dot(graph_type='graph', center=True) prefix = self.pathname + '.' if self.pathname else '' label = self.name if self.name else 'Model' top_node = pydot.Node(label, label=label, @@ -4373,24 +4377,44 @@ def _get_tree_graph(self, display_map=None): return pydot_graph - def _decorate_graph_for_display(self, G, exclude=()): + def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True): + """ + Add metadata to the graph for display. + + Returned graph will have any variable nodes containing certain characters relabeled with + explicit quotes to avoid issues with dot. + + Parameters + ---------- + G : nx.DiGraph + The graph to be decorated. + exclude : iter of str + Iter of pathnames to exclude from the generated graph. + + Returns + ------- + nx.DiGraph, dict + The decorated graph and a dict of node metadata keyed by pathname. + """ node_info = self._get_graph_display_info() exclude = set(exclude) - # dot doesn't like ':' in node names, so if we find any, we have to put explicit quotes - # around the node name and label. + prefix = self.pathname + '.' if self.pathname else '' + replace = {} for node, meta in G.nodes(data=True): + if not abs_graph_names: + node = prefix + node if node in node_info: meta.update(_filter_meta4dot(node_info[node])) - quoted = f'"{node.rpartition(".")[2]}"' if not ('label' in meta and meta['label']): - meta['label'] = quoted + meta['label'] = f'"{node.rpartition(".")[2]}"' if 'type_' in meta: # variable node if node.rpartition('.')[0] in exclude: exclude.add(node) # remove all variables of excluded components - if ':' in node and node not in exclude: # fix ':' in node names for use in dot + # quote node names containing certain characters for use in dot + if (':' in node or '<' in node) and node not in exclude: replace[node] = f'"{node}"' meta['shape'] = 'plain' # just text for variables, otherwise too busy @@ -4404,12 +4428,42 @@ def _decorate_graph_for_display(self, G, exclude=()): return G, node_info - def _add_boundary_nodes(self, G): + def _add_boundary_nodes(self, G, incoming, outgoing, exclude=()): + lenpre = len(self.pathname) + 1 if self.pathname else 0 + if incoming: + tooltip = ['External Connections:'] + connstrs = set() + for in_abs, out_abs in incoming: + if in_abs in G: + connstrs.add(f"{out_abs} -> {in_abs[lenpre:]}") + tooltip += sorted(connstrs) + tooltip='\n'.join(tooltip) + G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='min') + for in_abs, _ in incoming: + if in_abs in G: + G.add_edge('_Incoming', in_abs, style='dashed', arrowsize=0.5) + + if outgoing: + tooltip = ['External Connections:'] + connstrs = set() + for in_abs, out_abs in outgoing: + if out_abs in G: + connstrs.add(f"{out_abs[lenpre:]} -> {in_abs}") + tooltip += sorted(connstrs) + tooltip='\n'.join(tooltip) + G.add_node('_Outgoing', label='Outgoing', shape='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='max') + for _, out_abs in outgoing: + if out_abs in G: + G.add_edge(out_abs, '_Outgoing', style='dashed', arrowsize=0.5) + return G def _apply_clusters(self, G, node_info): pydot_graph, groups = self._get_cluster_tree(node_info) prefix = self.pathname + '.' if self.pathname else '' + boundary_nodes = {'_Incoming', '_Outgoing'} pydot_nodes = {} for node, meta in G.nodes(data=True): noquote_node = node.strip('"') @@ -4423,13 +4477,35 @@ def _apply_clusters(self, G, node_info): if group and group.startswith(prefix): groups[group].add_node(pdnode) + elif self.pathname in groups and node not in boundary_nodes: + groups[self.pathname].add_node(pdnode) else: pydot_graph.add_node(pdnode) - for u, v in G.edges(): - node1 = pydot_nodes[u] - node2 = pydot_nodes[v] - pydot_graph.add_edge(pydot.Edge(node1, node2, arrowhead='lnormal')) + for u, v, meta in G.edges(data=True): + pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], + **_filter_meta4dot(meta, + arrowsize=0.5))) + + # layout graph from left to right + pydot_graph.set_rankdir('LR') + + return pydot_graph + + def _to_pydot_graph(self, G): + gmeta = G.graph.get('graph', {}).copy() + gmeta['graph_type'] = 'digraph' + pydot_graph = pydot.Dot(**gmeta) + pydot_nodes = {} + + for node, meta in G.nodes(data=True): + pdnode = pydot_nodes[node] = pydot.Node(node, **_filter_meta4dot(meta)) + pydot_graph.add_node(pdnode) + + for u, v, meta in G.edges(data=True): + pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], + **_filter_meta4dot(meta, + arrowsize=0.5))) # layout graph from left to right pydot_graph.set_rankdir('LR') @@ -4509,11 +4585,16 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, G = self._dataflow_graph if not recurse: + G = G.copy() + # layout graph from left to right - G.graph['graph'] = {'rankdir': 'LR'} + gname = 'model' if self.pathname == '' else self.pathname + G.graph['graph'] = {'rankdir': 'LR', 'label': f"Dataflow for '{gname}'", + 'center': 'True'} # keep all direct children and their variables - keep = {n for n in G.nodes() if n[lenpre:].count('.') == 0} + keep = {n for n in G.nodes() if n[lenpre:].count('.') == 0 and + n.startswith(prefix)} keep.update({n for n, d in G.nodes(data=True) if 'type_' in d and n.rpartition('.')[0] in keep}) @@ -4521,28 +4602,38 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, inconnvars = set() outconnvars = set() for abs_in, abs_out in self._conn_abs_in2out.items(): - if abs_in not in keep: - inconnvars.add(abs_in) - if abs_out not in keep: - outconnvars.add(abs_out) + inconnvars.add(abs_in) + outconnvars.add(abs_out) for invar in inconnvars: - grp = prefix + invar[lenpre:].partition('.')[0] - if grp not in G: - G.add_node(grp, **_base_display_map['Group']) - G.add_edge(invar, grp) - keep.add(grp) - keep.add(invar) - G.nodes[invar]['label'] = self._var_allprocs_abs2prom['input'][invar] + system = prefix + invar[lenpre:].partition('.')[0] + if invar in G: + if system not in G: + G.add_node(system, **_base_display_map['Group']) + G.add_edge(invar, system) + keep.add(system) + keep.add(invar) + prom_in = self._var_allprocs_abs2prom["input"][invar] + par, _, child = prom_in.partition('.') + if par == invar[lenpre:].partition('.')[0]: + G.nodes[invar]['label'] = f'"{child}"' + else: + G.nodes[invar]['label'] = f'"{self._var_allprocs_abs2prom["input"][invar]}"' for outvar in outconnvars: - grp = prefix + outvar[lenpre:].partition('.')[0] - if grp not in G: - G.add_node(grp, **_base_display_map['Group']) - G.add_edge(grp, outvar) - keep.add(grp) - keep.add(outvar) - G.nodes[outvar]['label'] = self._var_allprocs_abs2prom['output'][outvar] + system = prefix + outvar[lenpre:].partition('.')[0] + if outvar in G: + if system not in G: + G.add_node(system, **_base_display_map['Group']) + G.add_edge(system, outvar) + keep.add(system) + keep.add(outvar) + prom_out = self._var_allprocs_abs2prom["output"][outvar] + par, _, child = prom_out.partition('.') + if par == outvar[lenpre:].partition('.')[0]: + G.nodes[outvar]['label'] = f'"{child}"' + else: + G.nodes[outvar]['label'] = f'"{self._var_allprocs_abs2prom["output"][outvar]}"' if self.pathname == '': if not recurse: @@ -4556,20 +4647,34 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, G = nx.subgraph(G, ournodes) + if show_boundary and self.pathname: + incoming, outgoing = self._get_boundary_conns() + G = self._add_boundary_nodes(G.copy(), incoming, outgoing) + G, node_info = self._decorate_graph_for_display(G, exclude=exclude) + if recurse: G = self._apply_clusters(G, node_info) + else: + G = self._to_pydot_graph(G) elif recurse: G = self.compute_sys_graph(comps_only=True, add_edge_info=False) + if show_boundary and self.pathname: + incoming, outgoing = self._get_boundary_conns() + # convert var abs names to system abs names + incoming = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) + for in_abs, out_abs in incoming] + outgoing = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) + for in_abs, out_abs in outgoing] + G = self._add_boundary_nodes(G.copy(), incoming, outgoing) + G, node_info = self._decorate_graph_for_display(G, exclude=exclude) G = self._apply_clusters(G, node_info) else: G = self.compute_sys_graph(comps_only=False, add_edge_info=False) - G, _ = self._decorate_graph_for_display(G, exclude=exclude) - - if show_boundary: - G = self._add_boundary_nodes(G) + G, _ = self._decorate_graph_for_display(G, exclude=exclude, abs_graph_names=False) + G = self._to_pydot_graph(G) else: raise ValueError(f"unrecognized graph type '{gtype}'. Allowed types are ['tree', " "'dataflow'].") @@ -5582,7 +5687,7 @@ def _cluster_color(path): return f"gray{col}" -def _filter_meta4dot(meta): +def _filter_meta4dot(meta, **kwargs): """ Remove unnecessary metadata from the given metadata dict before passing to pydot. @@ -5590,6 +5695,8 @@ def _filter_meta4dot(meta): ---------- meta : dict Metadata dict. + kwargs : dict + Additional metadata that will be added only if they are not already present. Returns ------- @@ -5597,4 +5704,8 @@ def _filter_meta4dot(meta): Metadata dict with unnecessary items removed. """ skip = {'type_', 'local', 'base', 'classname'} - return {k: v for k, v in meta.items() if k not in skip} + dct = {k: v for k, v in meta.items() if k not in skip} + for k, v in kwargs.items(): + if k not in dct: + dct[k] = v + return dct diff --git a/openmdao/test_suite/scripts/circuit_analysis.py b/openmdao/test_suite/scripts/circuit_analysis.py index 6aec1eb105..12eacf1b9d 100644 --- a/openmdao/test_suite/scripts/circuit_analysis.py +++ b/openmdao/test_suite/scripts/circuit_analysis.py @@ -156,3 +156,5 @@ def setup(self): p['circuit.n2.V'] = .7 p.run_model() + + p.model.circuit.write_graph(gtype='dataflow', show_vars=True, display=True, show_boundary=True, recurse=False) \ No newline at end of file diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index 559bf62334..77d86e7ff3 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -132,6 +132,9 @@ def _graph_setup_parser(parser): parser.add_argument('--show-vars', action='store_true', dest='show_vars', help="show variables in the graph. This only applies to the dataflow graph." " Default is False.") + parser.add_argument('--show-boundary', action='store_true', dest='show_boundary', + help="show connections to variables outside of the graph. This only " + "applies to the dataflow graph. Default is False.") parser.add_argument('--autoivc', action='store_true', dest='auto_ivc', help="include the _auto_ivc component in the graph. This applies to " "graphs of the top level group only. Default is False.") @@ -156,7 +159,7 @@ def _view_graph(problem): exclude = set() group.write_graph(gtype=options.type, recurse=options.recurse, show_vars=options.show_vars, display=options.show, exclude=exclude, - outfile=options.outfile) + show_boundary=options.show_boundary, outfile=options.outfile) # register the hooks hooks._register_hook('final_setup', 'Problem', post=_view_graph, exit=True) diff --git a/openmdao/utils/tests/test_cmdline.py b/openmdao/utils/tests/test_cmdline.py index 2d8b9cf4e4..d769197cd6 100644 --- a/openmdao/utils/tests/test_cmdline.py +++ b/openmdao/utils/tests/test_cmdline.py @@ -55,6 +55,13 @@ def _test_func_name(func, num, param): ('openmdao comm_info {}'.format(os.path.join(scriptdir, 'circle_opt.py')), {}), ('openmdao cite {}'.format(os.path.join(scriptdir, 'circle_opt.py')), {}), ('openmdao compute_entry_points openmdao', {}), + ('openmdao graph --no-display {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), + ('openmdao graph --no-display --type=tree {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), + ('openmdao graph --no-display --show-vars {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), + ('openmdao graph --no-display --show-vars --no-recurse {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), + ('openmdao graph --no-display --group=circuit {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), + ('openmdao graph --no-display --group=circuit --show-vars {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), + ('openmdao graph --no-display --group=circuit --show-vars --no-recurse {}'.format(os.path.join(scriptdir, 'circuit_analysis.py')), {}), ('openmdao iprof --no_browser {}'.format(os.path.join(scriptdir, 'circle_opt.py')), {'tornado': tornado}), ('openmdao iprof_totals {}'.format(os.path.join(scriptdir, 'circle_opt.py')), {}), From b9b6c817d5eca7cae620fcb92529080442d0523a Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Mon, 25 Mar 2024 14:12:34 -0400 Subject: [PATCH 07/15] tweaking shapes and tooltips --- openmdao/core/group.py | 107 ++++++++++++++++++++++++++-------- openmdao/utils/graph_utils.py | 2 +- 2 files changed, 85 insertions(+), 24 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 3246376312..490c0cc388 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -49,7 +49,7 @@ namecheck_rgx = re.compile('[a-zA-Z][_a-zA-Z0-9]*') -# mapping of system type to graph display properties +# mapping of base system type to graph display properties _base_display_map = { 'ExplicitComponent': { 'fillcolor': '"aquamarine3:aquamarine"', @@ -69,7 +69,7 @@ 'Group': { 'fillcolor': 'gray75', 'style': 'filled', - 'shape': 'octagon', + 'shape': 'box', }, } @@ -4253,9 +4253,16 @@ def _get_graph_display_info(self, display_map=None): meta.update(display_map[meta['classname']]) elif display_map and meta['base'] in display_map: meta.update(display_map[meta['base']]) - elif meta['base'] in _base_display_map: - meta.update(_base_display_map[meta['base']]) - meta['tooltip'] = f"{meta['classname']} {s.pathname}" + else: + _get_node_display_meta(s, meta) + + ttlist = [f"Name: {s.pathname}"] + ttlist.append(f"Class: {meta['classname']}") + if _has_nondef_lin_solver(s): + ttlist.append(f"Linear Solver: {type(s.linear_solver).__name__}") + if _has_nondef_nl_solver(s): + ttlist.append(f"Nonlinear Solver: {type(s.nonlinear_solver).__name__}") + meta['tooltip'] = '\n'.join(ttlist) node_info[s.pathname] = meta.copy() if self.comm.size > 1: @@ -4354,8 +4361,7 @@ def _get_tree_graph(self, display_map=None): prefix = self.pathname + '.' if self.pathname else '' label = self.name if self.name else 'Model' top_node = pydot.Node(label, label=label, - tooltip=f"{node_info[self.pathname]['classname']} {self.pathname}", - fillcolor='gray95', style='filled') + **node_info[self.pathname]) pydot_graph.add_node(top_node) systems[self.pathname] = top_node @@ -4410,6 +4416,8 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True): meta.update(_filter_meta4dot(node_info[node])) if not ('label' in meta and meta['label']): meta['label'] = f'"{node.rpartition(".")[2]}"' + else: + meta['label'] = f'"{meta["label"]}"' if 'type_' in meta: # variable node if node.rpartition('.')[0] in exclude: exclude.add(node) # remove all variables of excluded components @@ -4430,33 +4438,43 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True): def _add_boundary_nodes(self, G, incoming, outgoing, exclude=()): lenpre = len(self.pathname) + 1 if self.pathname else 0 + for ex in exclude: + expre = ex + '.' + incoming = [(in_abs, out_abs) for in_abs, out_abs in incoming + if in_abs != ex and out_abs != ex and + not in_abs.startswith(expre) and not out_abs.startswith(expre)] + outgoing = [(in_abs, out_abs) for in_abs, out_abs in outgoing + if in_abs != ex and out_abs != ex and + not in_abs.startswith(expre) and not out_abs.startswith(expre)] + if incoming: tooltip = ['External Connections:'] connstrs = set() for in_abs, out_abs in incoming: if in_abs in G: - connstrs.add(f"{out_abs} -> {in_abs[lenpre:]}") + connstrs.add(f" {out_abs} -> {in_abs[lenpre:]}") tooltip += sorted(connstrs) tooltip='\n'.join(tooltip) - G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='min') - for in_abs, _ in incoming: - if in_abs in G: - G.add_edge('_Incoming', in_abs, style='dashed', arrowsize=0.5) + if connstrs: + G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='min') + for in_abs, out_abs in incoming: + if in_abs in G: + G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', arrowsize=0.5) if outgoing: tooltip = ['External Connections:'] connstrs = set() for in_abs, out_abs in outgoing: if out_abs in G: - connstrs.add(f"{out_abs[lenpre:]} -> {in_abs}") + connstrs.add(f" {out_abs[lenpre:]} -> {in_abs}") tooltip += sorted(connstrs) tooltip='\n'.join(tooltip) - G.add_node('_Outgoing', label='Outgoing', shape='rarrow', fillcolor='peachpuff3', + G.add_node('_Outgoing', label='Outgoing', arrowhead='rarrow', fillcolor='peachpuff3', style='filled', tooltip=f'"{tooltip}"', rank='max') - for _, out_abs in outgoing: + for in_abs, out_abs in outgoing: if out_abs in G: - G.add_edge(out_abs, '_Outgoing', style='dashed', arrowsize=0.5) + G.add_edge(out_abs, '_Outgoing', style='dashed', shape='lnormal', arrowsize=0.5) return G @@ -4485,6 +4503,7 @@ def _apply_clusters(self, G, node_info): for u, v, meta in G.edges(data=True): pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], **_filter_meta4dot(meta, + arrowhead='lnormal', arrowsize=0.5))) # layout graph from left to right @@ -4616,9 +4635,9 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, prom_in = self._var_allprocs_abs2prom["input"][invar] par, _, child = prom_in.partition('.') if par == invar[lenpre:].partition('.')[0]: - G.nodes[invar]['label'] = f'"{child}"' + G.nodes[invar]['label'] = child else: - G.nodes[invar]['label'] = f'"{self._var_allprocs_abs2prom["input"][invar]}"' + G.nodes[invar]['label'] = self._var_allprocs_abs2prom["input"][invar] for outvar in outconnvars: system = prefix + outvar[lenpre:].partition('.')[0] @@ -4631,9 +4650,9 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, prom_out = self._var_allprocs_abs2prom["output"][outvar] par, _, child = prom_out.partition('.') if par == outvar[lenpre:].partition('.')[0]: - G.nodes[outvar]['label'] = f'"{child}"' + G.nodes[outvar]['label'] = child else: - G.nodes[outvar]['label'] = f'"{self._var_allprocs_abs2prom["output"][outvar]}"' + G.nodes[outvar]['label'] = self._var_allprocs_abs2prom["output"][outvar] if self.pathname == '': if not recurse: @@ -4649,7 +4668,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if show_boundary and self.pathname: incoming, outgoing = self._get_boundary_conns() - G = self._add_boundary_nodes(G.copy(), incoming, outgoing) + G = self._add_boundary_nodes(G.copy(), incoming, outgoing, exclude) G, node_info = self._decorate_graph_for_display(G, exclude=exclude) @@ -4667,7 +4686,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, for in_abs, out_abs in incoming] outgoing = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) for in_abs, out_abs in outgoing] - G = self._add_boundary_nodes(G.copy(), incoming, outgoing) + G = self._add_boundary_nodes(G.copy(), incoming, outgoing, exclude) G, node_info = self._decorate_graph_for_display(G, exclude=exclude) G = self._apply_clusters(G, node_info) @@ -5709,3 +5728,45 @@ def _filter_meta4dot(meta, **kwargs): if k not in dct: dct[k] = v return dct + + +def _has_nondef_nl_solver(system): + """ + Return True if the given system has a nonlinear solver that is not the default. + + Parameters + ---------- + system : System + The system to check. + + Returns + ------- + bool + True if the system has a nonlinear solver that is not the default. + """ + return system.nonlinear_solver is not None and not isinstance(system.nonlinear_solver, + NonlinearRunOnce) + + +def _has_nondef_lin_solver(system): + """ + Return True if the given system has a linear solver that is not the default. + + Parameters + ---------- + system : System + The system to check. + + Returns + ------- + bool + True if the system has a linear solver that is not the default. + """ + return system.linear_solver is not None and not isinstance(system.linear_solver, LinearRunOnce) + + +def _get_node_display_meta(s, meta): + if meta['base'] in _base_display_map: + meta.update(_base_display_map[meta['base']]) + if _has_nondef_lin_solver(s) or _has_nondef_nl_solver(s): + meta['shape'] = 'box3d' diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index 77d86e7ff3..eb70a07396 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -153,7 +153,7 @@ def _graph_cmd(options, user_args): """ def _view_graph(problem): group = problem.model._get_subsystem(options.group) if options.group else problem.model - if not options.group and not options.auto_ivc: + if not options.auto_ivc: exclude = {'_auto_ivc'} else: exclude = set() From 24f752a742875cd03686bdbb39af3468cdfd2ddd Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Mon, 25 Mar 2024 17:07:13 -0400 Subject: [PATCH 08/15] working --- openmdao/core/group.py | 90 +++++++++++++++++++++++------------------- 1 file changed, 49 insertions(+), 41 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 490c0cc388..9924b47228 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4245,6 +4245,17 @@ def compute_sys_graph(self, comps_only=False, add_edge_info=True): return graph + def _get_prom_conns(self, conns): + abs2prom_in = self._var_allprocs_abs2prom['input'] + abs2prom_out = self._var_allprocs_abs2prom['output'] + prom2abs_in = self._var_allprocs_prom2abs_list['input'] + prefix = self.pathname + '.' if self.pathname else '' + prom_conns = {} + for inp, out in conns.items(): + prom = abs2prom_in[inp] + prom_conns[prom] = (out, [i for i in prom2abs_in[prom] if i.startswith(prefix)]) + return prom_conns + def _get_graph_display_info(self, display_map=None): node_info = {} for s in self.system_iter(recurse=True, include_self=True): @@ -4598,14 +4609,9 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, prefix = self.pathname + '.' if self.pathname else '' lenpre = len(prefix) - if self.pathname: - G = self._problem_meta['model_ref']()._dataflow_graph - else: - G = self._dataflow_graph + G = self._problem_meta['model_ref']()._get_dataflow_graph() if not recurse: - G = G.copy() - # layout graph from left to right gname = 'model' if self.pathname == '' else self.pathname G.graph['graph'] = {'rankdir': 'LR', 'label': f"Dataflow for '{gname}'", @@ -4617,42 +4623,44 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, keep.update({n for n, d in G.nodes(data=True) if 'type_' in d and n.rpartition('.')[0] in keep}) - # keep all variables involved in connections owned by this group - inconnvars = set() - outconnvars = set() - for abs_in, abs_out in self._conn_abs_in2out.items(): - inconnvars.add(abs_in) - outconnvars.add(abs_out) - - for invar in inconnvars: - system = prefix + invar[lenpre:].partition('.')[0] - if invar in G: - if system not in G: - G.add_node(system, **_base_display_map['Group']) - G.add_edge(invar, system) - keep.add(system) - keep.add(invar) - prom_in = self._var_allprocs_abs2prom["input"][invar] - par, _, child = prom_in.partition('.') - if par == invar[lenpre:].partition('.')[0]: - G.nodes[invar]['label'] = child - else: - G.nodes[invar]['label'] = self._var_allprocs_abs2prom["input"][invar] - - for outvar in outconnvars: - system = prefix + outvar[lenpre:].partition('.')[0] - if outvar in G: - if system not in G: - G.add_node(system, **_base_display_map['Group']) - G.add_edge(system, outvar) - keep.add(system) - keep.add(outvar) - prom_out = self._var_allprocs_abs2prom["output"][outvar] - par, _, child = prom_out.partition('.') - if par == outvar[lenpre:].partition('.')[0]: - G.nodes[outvar]['label'] = child + promconns = self._get_prom_conns(self._conn_abs_in2out) + + for prom_in, (abs_out, abs_ins) in promconns.items(): + nins = len(abs_ins) + the_in = abs_ins[0] if nins == 1 else prom_in + if the_in not in G: + G.add_node(the_in, type_='input', label=prom_in) + else: + l = prom_in[:lenpre] if prom_in.startswith(prefix) else prom_in + G.nodes[the_in]['label'] = l + + sysout = prefix + abs_out[lenpre:].partition('.')[0] + prom_out = self._var_allprocs_abs2prom['output'][abs_out] + + if sysout not in G: + G.add_node(sysout, **_base_display_map['Group']) + + l = prom_out[:lenpre] if prom_out.startswith(prefix) else prom_out + G.nodes[abs_out]['label'] = l + G.add_edge(sysout, abs_out) + + keep.add(sysout) + keep.add(abs_out) + + for abs_in in abs_ins: + sysin = prefix + abs_in[lenpre:].partition('.')[0] + if sysin not in G: + G.add_node(sysin, **_base_display_map['Group']) + if nins == 1: + G.add_edge(abs_ins[0], sysin) + keep.add(abs_ins[0]) else: - G.nodes[outvar]['label'] = self._var_allprocs_abs2prom["output"][outvar] + G.add_edge(prom_in, sysin) + keep.add(prom_in) + + if prom_in in G and nins > 1: + G.nodes[prom_in]['fontcolor'] = 'red' + G.nodes[prom_in]['tooltip'] = '\n'.join(abs_ins) if self.pathname == '': if not recurse: From be5590889a4e4957a420f9795d1493fc224cad6f Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Tue, 26 Mar 2024 09:38:03 -0400 Subject: [PATCH 09/15] added default exclusion of auto_ivc from tree view --- openmdao/core/group.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 9924b47228..3c14f44286 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4351,12 +4351,14 @@ def _get_cluster_tree(self, node_info): return pydot_graph, groups - def _get_tree_graph(self, display_map=None): + def _get_tree_graph(self, exclude, display_map=None): """ Create a pydot graph of the system tree (without clusters). Parameters ---------- + exclude : iter of str + Iter of pathnames to exclude from the generated graph. display_map : dict or None A map of classnames to pydot node attributes. @@ -4366,6 +4368,7 @@ def _get_tree_graph(self, display_map=None): The pydot tree graph. """ node_info = self._get_graph_display_info(display_map) + exclude = set(exclude) systems = {} pydot_graph = pydot.Dot(graph_type='graph', center=True) @@ -4379,7 +4382,7 @@ def _get_tree_graph(self, display_map=None): for varpath in chain(self._var_allprocs_abs2prom['input'], self._var_allprocs_abs2prom['output']): system = varpath.rpartition('.')[0] - if system not in systems: + if system not in systems and system not in exclude: # reverse the list so parents will exist before children ancestor_list = list(all_ancestors(system))[::-1] for path in ancestor_list: @@ -4603,9 +4606,13 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, return if gtype == 'tree': - G = self._get_tree_graph() + G = self._get_tree_graph(exclude) elif gtype == 'dataflow': if show_vars: + # get dvs and responses so we can color them differently + dvs = self.get_design_vars(recurse=True, get_sizes=False, use_prom_ivc=True) + resps = self.get_responses(recurse=True, get_sizes=False, use_prom_ivc=True) + prefix = self.pathname + '.' if self.pathname else '' lenpre = len(prefix) From 69cb97511a43f5401ba3c9b7269024e39913117d Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Tue, 26 Mar 2024 14:09:29 -0400 Subject: [PATCH 10/15] added visual display of dvs and responses --- openmdao/core/group.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 3c14f44286..2ef19e2417 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4397,7 +4397,8 @@ def _get_tree_graph(self, exclude, display_map=None): return pydot_graph - def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True): + def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True, dvs=None, + responses=None): """ Add metadata to the graph for display. @@ -4410,6 +4411,12 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True): The graph to be decorated. exclude : iter of str Iter of pathnames to exclude from the generated graph. + abs_graph_names : bool + If True, use absolute pathnames for nodes in the graph. + dvs : dict + Dict of design var metadata keyed on absolute name. + responses : list of str + Dict of response var metadata keyed on absolute name. Returns ------- @@ -4438,7 +4445,12 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True): # quote node names containing certain characters for use in dot if (':' in node or '<' in node) and node not in exclude: replace[node] = f'"{node}"' - meta['shape'] = 'plain' # just text for variables, otherwise too busy + if dvs and node in dvs: + meta['shape'] = 'cds' + elif responses and node in responses: + meta['shape'] = 'cds' + else: + meta['shape'] = 'plain' # just text for variables, otherwise too busy if replace: G = nx.relabel_nodes(G, replace, copy=True) @@ -4611,7 +4623,11 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if show_vars: # get dvs and responses so we can color them differently dvs = self.get_design_vars(recurse=True, get_sizes=False, use_prom_ivc=True) + desvars = set(dvs) + desvars.update(m['source'] for m in dvs.values()) resps = self.get_responses(recurse=True, get_sizes=False, use_prom_ivc=True) + responses = set(resps) + responses.update(m['source'] for m in resps.values()) prefix = self.pathname + '.' if self.pathname else '' lenpre = len(prefix) @@ -4685,7 +4701,8 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, incoming, outgoing = self._get_boundary_conns() G = self._add_boundary_nodes(G.copy(), incoming, outgoing, exclude) - G, node_info = self._decorate_graph_for_display(G, exclude=exclude) + G, node_info = self._decorate_graph_for_display(G, exclude=exclude, + dvs=desvars, responses=responses) if recurse: G = self._apply_clusters(G, node_info) From da61bce98a27dc4749fc79b586f1c4c96aae0e09 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Tue, 26 Mar 2024 15:25:06 -0400 Subject: [PATCH 11/15] cleanup --- openmdao/core/group.py | 106 +++++----------------------------- openmdao/utils/graph_utils.py | 97 ++++++++++++++++++++++++++++++- 2 files changed, 107 insertions(+), 96 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 2ef19e2417..f23d59c288 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -35,7 +35,8 @@ meta2src_iter, get_rev_conns, _contains_all from openmdao.utils.units import is_compatible, unit_conversion, _has_val_mismatch, _find_unit, \ _is_unitless, simplify_unit -from openmdao.utils.graph_utils import get_out_of_order_nodes, write_graph +from openmdao.utils.graph_utils import get_out_of_order_nodes, write_graph, _filter_meta4dot, \ + _to_pydot_graph, _add_boundary_nodes from openmdao.utils.mpi import MPI, check_mpi_exceptions, multi_proc_exception_check import openmdao.utils.coloring as coloring_mod from openmdao.utils.indexer import indexer, Indexer @@ -4247,7 +4248,6 @@ def compute_sys_graph(self, comps_only=False, add_edge_info=True): def _get_prom_conns(self, conns): abs2prom_in = self._var_allprocs_abs2prom['input'] - abs2prom_out = self._var_allprocs_abs2prom['output'] prom2abs_in = self._var_allprocs_prom2abs_list['input'] prefix = self.pathname + '.' if self.pathname else '' prom_conns = {} @@ -4323,10 +4323,16 @@ def _get_cluster_tree(self, node_info): pydot.Dot, dict The pydot graph and a dict of groups keyed by pathname. """ - groups = {} pydot_graph = pydot.Dot(graph_type='digraph') mypath = self.pathname prefix = mypath + '.' if mypath else '' + groups = {} + + if not mypath: + groups[''] = pydot.Cluster('', label='Model', style='filled', fillcolor=_cluster_color(''), + tooltip=node_info['']['tooltip']) + pydot_graph.add_subgraph(groups['']) + for varpath in chain(self._var_allprocs_abs2prom['input'], self._var_allprocs_abs2prom['output']): group = varpath.rpartition('.')[0].rpartition('.')[0] @@ -4462,48 +4468,6 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True, dvs=N return G, node_info - def _add_boundary_nodes(self, G, incoming, outgoing, exclude=()): - lenpre = len(self.pathname) + 1 if self.pathname else 0 - for ex in exclude: - expre = ex + '.' - incoming = [(in_abs, out_abs) for in_abs, out_abs in incoming - if in_abs != ex and out_abs != ex and - not in_abs.startswith(expre) and not out_abs.startswith(expre)] - outgoing = [(in_abs, out_abs) for in_abs, out_abs in outgoing - if in_abs != ex and out_abs != ex and - not in_abs.startswith(expre) and not out_abs.startswith(expre)] - - if incoming: - tooltip = ['External Connections:'] - connstrs = set() - for in_abs, out_abs in incoming: - if in_abs in G: - connstrs.add(f" {out_abs} -> {in_abs[lenpre:]}") - tooltip += sorted(connstrs) - tooltip='\n'.join(tooltip) - if connstrs: - G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='min') - for in_abs, out_abs in incoming: - if in_abs in G: - G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', arrowsize=0.5) - - if outgoing: - tooltip = ['External Connections:'] - connstrs = set() - for in_abs, out_abs in outgoing: - if out_abs in G: - connstrs.add(f" {out_abs[lenpre:]} -> {in_abs}") - tooltip += sorted(connstrs) - tooltip='\n'.join(tooltip) - G.add_node('_Outgoing', label='Outgoing', arrowhead='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='max') - for in_abs, out_abs in outgoing: - if out_abs in G: - G.add_edge(out_abs, '_Outgoing', style='dashed', shape='lnormal', arrowsize=0.5) - - return G - def _apply_clusters(self, G, node_info): pydot_graph, groups = self._get_cluster_tree(node_info) prefix = self.pathname + '.' if self.pathname else '' @@ -4537,26 +4501,6 @@ def _apply_clusters(self, G, node_info): return pydot_graph - def _to_pydot_graph(self, G): - gmeta = G.graph.get('graph', {}).copy() - gmeta['graph_type'] = 'digraph' - pydot_graph = pydot.Dot(**gmeta) - pydot_nodes = {} - - for node, meta in G.nodes(data=True): - pdnode = pydot_nodes[node] = pydot.Node(node, **_filter_meta4dot(meta)) - pydot_graph.add_node(pdnode) - - for u, v, meta in G.edges(data=True): - pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], - **_filter_meta4dot(meta, - arrowsize=0.5))) - - # layout graph from left to right - pydot_graph.set_rankdir('LR') - - return pydot_graph - def _get_boundary_conns(self): """ Return lists of incoming and outgoing boundary connections. @@ -4699,7 +4643,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if show_boundary and self.pathname: incoming, outgoing = self._get_boundary_conns() - G = self._add_boundary_nodes(G.copy(), incoming, outgoing, exclude) + G = _add_boundary_nodes(self.pathname, G.copy(), incoming, outgoing, exclude) G, node_info = self._decorate_graph_for_display(G, exclude=exclude, dvs=desvars, responses=responses) @@ -4707,7 +4651,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if recurse: G = self._apply_clusters(G, node_info) else: - G = self._to_pydot_graph(G) + G = _to_pydot_graph(G) elif recurse: G = self.compute_sys_graph(comps_only=True, add_edge_info=False) @@ -4718,14 +4662,14 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, for in_abs, out_abs in incoming] outgoing = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) for in_abs, out_abs in outgoing] - G = self._add_boundary_nodes(G.copy(), incoming, outgoing, exclude) + G = _add_boundary_nodes(self.pathname, G.copy(), incoming, outgoing, exclude) G, node_info = self._decorate_graph_for_display(G, exclude=exclude) G = self._apply_clusters(G, node_info) else: G = self.compute_sys_graph(comps_only=False, add_edge_info=False) G, _ = self._decorate_graph_for_display(G, exclude=exclude, abs_graph_names=False) - G = self._to_pydot_graph(G) + G = _to_pydot_graph(G) else: raise ValueError(f"unrecognized graph type '{gtype}'. Allowed types are ['tree', " "'dataflow'].") @@ -5738,30 +5682,6 @@ def _cluster_color(path): return f"gray{col}" -def _filter_meta4dot(meta, **kwargs): - """ - Remove unnecessary metadata from the given metadata dict before passing to pydot. - - Parameters - ---------- - meta : dict - Metadata dict. - kwargs : dict - Additional metadata that will be added only if they are not already present. - - Returns - ------- - dict - Metadata dict with unnecessary items removed. - """ - skip = {'type_', 'local', 'base', 'classname'} - dct = {k: v for k, v in meta.items() if k not in skip} - for k, v in kwargs.items(): - if k not in dct: - dct[k] = v - return dct - - def _has_nondef_nl_solver(system): """ Return True if the given system has a nonlinear solver that is not the default. diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index eb70a07396..7c87838393 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -5,6 +5,11 @@ from openmdao.utils.file_utils import _load_and_exec import openmdao.utils.hooks as hooks +try: + import pydot +except ImportError: + pydot = None + def get_sccs_topo(graph): """ @@ -80,9 +85,7 @@ def write_graph(G, prog='dot', display=True, outfile='graph.svg'): """ from openmdao.utils.webview import webview - try: - import pydot - except ImportError: + if pydot is None: raise RuntimeError("graph requires the pydot package. You can install it using " "'pip install pydot'.") @@ -165,3 +168,91 @@ def _view_graph(problem): hooks._register_hook('final_setup', 'Problem', post=_view_graph, exit=True) _load_and_exec(options.file[0], user_args) + +def _to_pydot_graph(G): + gmeta = G.graph.get('graph', {}).copy() + gmeta['graph_type'] = 'digraph' + pydot_graph = pydot.Dot(**gmeta) + pydot_nodes = {} + + for node, meta in G.nodes(data=True): + pdnode = pydot_nodes[node] = pydot.Node(node, **_filter_meta4dot(meta)) + pydot_graph.add_node(pdnode) + + for u, v, meta in G.edges(data=True): + pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], + **_filter_meta4dot(meta, + arrowsize=0.5))) + + # layout graph from left to right + pydot_graph.set_rankdir('LR') + + return pydot_graph + + +def _filter_meta4dot(meta, **kwargs): + """ + Remove unnecessary metadata from the given metadata dict before passing to pydot. + + Parameters + ---------- + meta : dict + Metadata dict. + kwargs : dict + Additional metadata that will be added only if they are not already present. + + Returns + ------- + dict + Metadata dict with unnecessary items removed. + """ + skip = {'type_', 'local', 'base', 'classname'} + dct = {k: v for k, v in meta.items() if k not in skip} + for k, v in kwargs.items(): + if k not in dct: + dct[k] = v + return dct + + +def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): + lenpre = len(pathname) + 1 if pathname else 0 + for ex in exclude: + expre = ex + '.' + incoming = [(in_abs, out_abs) for in_abs, out_abs in incoming + if in_abs != ex and out_abs != ex and + not in_abs.startswith(expre) and not out_abs.startswith(expre)] + outgoing = [(in_abs, out_abs) for in_abs, out_abs in outgoing + if in_abs != ex and out_abs != ex and + not in_abs.startswith(expre) and not out_abs.startswith(expre)] + + if incoming: + tooltip = ['External Connections:'] + connstrs = set() + for in_abs, out_abs in incoming: + if in_abs in G: + connstrs.add(f" {out_abs} -> {in_abs[lenpre:]}") + tooltip += sorted(connstrs) + tooltip='\n'.join(tooltip) + if connstrs: + G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='min') + for in_abs, out_abs in incoming: + if in_abs in G: + G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', arrowsize=0.5) + + if outgoing: + tooltip = ['External Connections:'] + connstrs = set() + for in_abs, out_abs in outgoing: + if out_abs in G: + connstrs.add(f" {out_abs[lenpre:]} -> {in_abs}") + tooltip += sorted(connstrs) + tooltip='\n'.join(tooltip) + G.add_node('_Outgoing', label='Outgoing', arrowhead='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='max') + for in_abs, out_abs in outgoing: + if out_abs in G: + G.add_edge(out_abs, '_Outgoing', style='dashed', shape='lnormal', arrowsize=0.5) + + return G + From 97c5696e52dd9618e01664ccbce568f3846e04ea Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Tue, 26 Mar 2024 15:32:24 -0400 Subject: [PATCH 12/15] cleanup --- openmdao/core/group.py | 72 ++----------------- .../test_suite/scripts/circuit_analysis.py | 2 - openmdao/utils/graph_utils.py | 25 +++++++ 3 files changed, 32 insertions(+), 67 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index f23d59c288..293569e600 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -36,7 +36,7 @@ from openmdao.utils.units import is_compatible, unit_conversion, _has_val_mismatch, _find_unit, \ _is_unitless, simplify_unit from openmdao.utils.graph_utils import get_out_of_order_nodes, write_graph, _filter_meta4dot, \ - _to_pydot_graph, _add_boundary_nodes + _to_pydot_graph, _add_boundary_nodes, _cluster_color from openmdao.utils.mpi import MPI, check_mpi_exceptions, multi_proc_exception_check import openmdao.utils.coloring as coloring_mod from openmdao.utils.indexer import indexer, Indexer @@ -4269,9 +4269,10 @@ def _get_graph_display_info(self, display_map=None): ttlist = [f"Name: {s.pathname}"] ttlist.append(f"Class: {meta['classname']}") - if _has_nondef_lin_solver(s): + if s.linear_solver is not None and not isinstance(s.linear_solver, NonlinearRunOnce): ttlist.append(f"Linear Solver: {type(s.linear_solver).__name__}") - if _has_nondef_nl_solver(s): + if s.nonlinear_solver is not None and not isinstance(s.nonlinear_solver, + NonlinearRunOnce): ttlist.append(f"Nonlinear Solver: {type(s.nonlinear_solver).__name__}") meta['tooltip'] = '\n'.join(ttlist) node_info[s.pathname] = meta.copy() @@ -5656,69 +5657,10 @@ def _active_responses(self, user_response_names, responses=None): return active_resps -def _cluster_color(path): - """ - Return the color of the cluster that contains the given path. - - The idea here is to make nested clusters stand out wrt their parent cluster. - - Parameters - ---------- - path : str - Pathname of a variable. - - Returns - ------- - int - The color of the cluster that contains the given path. - """ - depth = path.count('.') + 1 if path else 0 - - ncolors = 10 - maxcolor = 98 - mincolor = 40 - - col = maxcolor - (depth % ncolors) * (maxcolor - mincolor) // ncolors - return f"gray{col}" - - -def _has_nondef_nl_solver(system): - """ - Return True if the given system has a nonlinear solver that is not the default. - - Parameters - ---------- - system : System - The system to check. - - Returns - ------- - bool - True if the system has a nonlinear solver that is not the default. - """ - return system.nonlinear_solver is not None and not isinstance(system.nonlinear_solver, - NonlinearRunOnce) - - -def _has_nondef_lin_solver(system): - """ - Return True if the given system has a linear solver that is not the default. - - Parameters - ---------- - system : System - The system to check. - - Returns - ------- - bool - True if the system has a linear solver that is not the default. - """ - return system.linear_solver is not None and not isinstance(system.linear_solver, LinearRunOnce) - - def _get_node_display_meta(s, meta): if meta['base'] in _base_display_map: meta.update(_base_display_map[meta['base']]) - if _has_nondef_lin_solver(s) or _has_nondef_nl_solver(s): + if s.nonlinear_solver is not None and not isinstance(s.nonlinear_solver, NonlinearRunOnce): + meta['shape'] = 'box3d' + elif s.linear_solver is not None and not isinstance(s.linear_solver, LinearRunOnce): meta['shape'] = 'box3d' diff --git a/openmdao/test_suite/scripts/circuit_analysis.py b/openmdao/test_suite/scripts/circuit_analysis.py index 12eacf1b9d..6aec1eb105 100644 --- a/openmdao/test_suite/scripts/circuit_analysis.py +++ b/openmdao/test_suite/scripts/circuit_analysis.py @@ -156,5 +156,3 @@ def setup(self): p['circuit.n2.V'] = .7 p.run_model() - - p.model.circuit.write_graph(gtype='dataflow', show_vars=True, display=True, show_boundary=True, recurse=False) \ No newline at end of file diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index 7c87838393..a736aac28f 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -256,3 +256,28 @@ def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): return G + +def _cluster_color(path): + """ + Return the color of the cluster that contains the given path. + + The idea here is to make nested clusters stand out wrt their parent cluster. + + Parameters + ---------- + path : str + Pathname of a variable. + + Returns + ------- + int + The color of the cluster that contains the given path. + """ + depth = path.count('.') + 1 if path else 0 + + ncolors = 10 + maxcolor = 98 + mincolor = 40 + + col = maxcolor - (depth % ncolors) * (maxcolor - mincolor) // ncolors + return f"gray{col}" From d0107f1d68871c4c3a06d1a34ba30ddf9537cc08 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Tue, 26 Mar 2024 15:47:26 -0400 Subject: [PATCH 13/15] added docstrings --- openmdao/core/group.py | 45 +++++++++++++++++++++++++++++++++-- openmdao/utils/graph_utils.py | 21 ++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index 293569e600..fdea0a079e 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4247,6 +4247,19 @@ def compute_sys_graph(self, comps_only=False, add_edge_info=True): return graph def _get_prom_conns(self, conns): + """ + Return a dict of promoted connections. + + Parameters + ---------- + conns : dict + Dictionary containing absolute connections. + + Returns + ------- + dict + Dictionary of promoted connections. + """ abs2prom_in = self._var_allprocs_abs2prom['input'] prom2abs_in = self._var_allprocs_prom2abs_list['input'] prefix = self.pathname + '.' if self.pathname else '' @@ -4257,6 +4270,19 @@ def _get_prom_conns(self, conns): return prom_conns def _get_graph_display_info(self, display_map=None): + """ + Return display related metadata for this Group and all of its children. + + Parameters + ---------- + display_map : dict or None + A map of classnames to pydot node attributes. + + Returns + ------- + dict + Metadata keyed by system pathname. + """ node_info = {} for s in self.system_iter(recurse=True, include_self=True): meta = s._get_graph_node_meta() @@ -4421,9 +4447,9 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True, dvs=N abs_graph_names : bool If True, use absolute pathnames for nodes in the graph. dvs : dict - Dict of design var metadata keyed on absolute name. + Dict of design var metadata keyed on promoted name. responses : list of str - Dict of response var metadata keyed on absolute name. + Dict of response var metadata keyed on promoted name. Returns ------- @@ -4470,6 +4496,21 @@ def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True, dvs=N return G, node_info def _apply_clusters(self, G, node_info): + """ + Group nodes in the graph into clusters. + + Parameters + ---------- + G : nx.DiGraph + A pydot graph will be created based on this graph. + node_info : dict + A dict of metadata keyed by pathname. + + Returns + ------- + pydot.Graph + The corresponding pydot graph with clusters added. + """ pydot_graph, groups = self._get_cluster_tree(node_info) prefix = self.pathname + '.' if self.pathname else '' boundary_nodes = {'_Incoming', '_Outgoing'} diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index a736aac28f..449f715d48 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -215,6 +215,27 @@ def _filter_meta4dot(meta, **kwargs): def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): + """ + Add boundary nodes to the graph. + + Parameters + ---------- + pathname : str + Pathname of the current group. + G : nx.DiGraph + The graph. + incoming : list of (str, str) + List of incoming connections. + outgoing : list of (str, str) + List of outgoing connections. + exclude : list of str + List of pathnames to exclude from the graph. + + Returns + ------- + nx.DiGraph + The modified graph. + """ lenpre = len(pathname) + 1 if pathname else 0 for ex in exclude: expre = ex + '.' From e86a7f20b04e7b07313d91c96d9bf1636fccd418 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Wed, 27 Mar 2024 09:41:50 -0400 Subject: [PATCH 14/15] fixed pep8 errors --- openmdao/core/group.py | 17 +++++++++-------- openmdao/utils/graph_utils.py | 14 +++++++------- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/openmdao/core/group.py b/openmdao/core/group.py index fdea0a079e..ebd41af1ab 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -4356,7 +4356,8 @@ def _get_cluster_tree(self, node_info): groups = {} if not mypath: - groups[''] = pydot.Cluster('', label='Model', style='filled', fillcolor=_cluster_color(''), + groups[''] = pydot.Cluster('', label='Model', style='filled', + fillcolor=_cluster_color(''), tooltip=node_info['']['tooltip']) pydot_graph.add_subgraph(groups['']) @@ -4367,7 +4368,7 @@ def _get_cluster_tree(self, node_info): # reverse the list so parents will exist before children ancestor_list = list(all_ancestors(group))[::-1] for path in ancestor_list: - if path.startswith(prefix) or path==mypath: + if path.startswith(prefix) or path == mypath: if path not in groups: parent, _, name = path.rpartition('.') groups[path] = pydot.Cluster(path, @@ -4600,7 +4601,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, """ if pydot is None: issue_warning("To view the system graph, you must install pydot. " - "You can install it via 'pip install pydot'.") + "You can install it via 'pip install pydot'.") return if gtype == 'tree': @@ -4640,8 +4641,8 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if the_in not in G: G.add_node(the_in, type_='input', label=prom_in) else: - l = prom_in[:lenpre] if prom_in.startswith(prefix) else prom_in - G.nodes[the_in]['label'] = l + label = prom_in[:lenpre] if prom_in.startswith(prefix) else prom_in + G.nodes[the_in]['label'] = label sysout = prefix + abs_out[lenpre:].partition('.')[0] prom_out = self._var_allprocs_abs2prom['output'][abs_out] @@ -4649,8 +4650,8 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if sysout not in G: G.add_node(sysout, **_base_display_map['Group']) - l = prom_out[:lenpre] if prom_out.startswith(prefix) else prom_out - G.nodes[abs_out]['label'] = l + label = prom_out[:lenpre] if prom_out.startswith(prefix) else prom_out + G.nodes[abs_out]['label'] = label G.add_edge(sysout, abs_out) keep.add(sysout) @@ -4673,7 +4674,7 @@ def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, if self.pathname == '': if not recurse: - G = nx.subgraph(G, keep) + G = nx.subgraph(G, keep) else: # we're not the top level group, so get our subgraph of the top level graph ournodes = {n for n in G.nodes() if n.startswith(prefix)} diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index 449f715d48..de7bd5ee94 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -181,8 +181,7 @@ def _to_pydot_graph(G): for u, v, meta in G.edges(data=True): pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], - **_filter_meta4dot(meta, - arrowsize=0.5))) + **_filter_meta4dot(meta, arrowsize=0.5))) # layout graph from left to right pydot_graph.set_rankdir('LR') @@ -253,13 +252,14 @@ def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): if in_abs in G: connstrs.add(f" {out_abs} -> {in_abs[lenpre:]}") tooltip += sorted(connstrs) - tooltip='\n'.join(tooltip) + tooltip = '\n'.join(tooltip) if connstrs: G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='min') + style='filled', tooltip=f'"{tooltip}"', rank='min') for in_abs, out_abs in incoming: if in_abs in G: - G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', arrowsize=0.5) + G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', + arrowsize=0.5) if outgoing: tooltip = ['External Connections:'] @@ -268,9 +268,9 @@ def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): if out_abs in G: connstrs.add(f" {out_abs[lenpre:]} -> {in_abs}") tooltip += sorted(connstrs) - tooltip='\n'.join(tooltip) + tooltip = '\n'.join(tooltip) G.add_node('_Outgoing', label='Outgoing', arrowhead='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='max') + style='filled', tooltip=f'"{tooltip}"', rank='max') for in_abs, out_abs in outgoing: if out_abs in G: G.add_edge(out_abs, '_Outgoing', style='dashed', shape='lnormal', arrowsize=0.5) From 90f3a2e3721c846d6364cfed1c2ea742b7946a45 Mon Sep 17 00:00:00 2001 From: Bret Naylor Date: Wed, 27 Mar 2024 14:28:02 -0400 Subject: [PATCH 15/15] moved graph viewer stuff into its own class --- openmdao/core/group.py | 533 +---------------- openmdao/utils/graph_utils.py | 248 -------- openmdao/utils/om.py | 2 +- openmdao/visualization/graph_viewer.py | 792 +++++++++++++++++++++++++ 4 files changed, 807 insertions(+), 768 deletions(-) create mode 100644 openmdao/visualization/graph_viewer.py diff --git a/openmdao/core/group.py b/openmdao/core/group.py index ebd41af1ab..36f495a6fa 100644 --- a/openmdao/core/group.py +++ b/openmdao/core/group.py @@ -11,11 +11,6 @@ import numpy as np import networkx as nx -try: - import pydot -except ImportError: - pydot = None - from openmdao.core.configinfo import _ConfigInfo from openmdao.core.system import System, collect_errors from openmdao.core.component import Component, _DictValues @@ -35,8 +30,7 @@ meta2src_iter, get_rev_conns, _contains_all from openmdao.utils.units import is_compatible, unit_conversion, _has_val_mismatch, _find_unit, \ _is_unitless, simplify_unit -from openmdao.utils.graph_utils import get_out_of_order_nodes, write_graph, _filter_meta4dot, \ - _to_pydot_graph, _add_boundary_nodes, _cluster_color +from openmdao.utils.graph_utils import get_out_of_order_nodes from openmdao.utils.mpi import MPI, check_mpi_exceptions, multi_proc_exception_check import openmdao.utils.coloring as coloring_mod from openmdao.utils.indexer import indexer, Indexer @@ -50,31 +44,6 @@ namecheck_rgx = re.compile('[a-zA-Z][_a-zA-Z0-9]*') -# mapping of base system type to graph display properties -_base_display_map = { - 'ExplicitComponent': { - 'fillcolor': '"aquamarine3:aquamarine"', - 'style': 'filled', - 'shape': 'box', - }, - 'ImplicitComponent': { - 'fillcolor': '"lightblue:lightslateblue"', - 'style': 'filled', - 'shape': 'box', - }, - 'IndepVarComp': { - 'fillcolor': '"chartreuse2:chartreuse4"', - 'style': 'filled', - 'shape': 'box', - }, - 'Group': { - 'fillcolor': 'gray75', - 'style': 'filled', - 'shape': 'box', - }, -} - - # use a class with slots instead of a namedtuple so that we can # change index after creation if needed. class _SysInfo(object): @@ -4246,486 +4215,6 @@ def compute_sys_graph(self, comps_only=False, add_edge_info=True): return graph - def _get_prom_conns(self, conns): - """ - Return a dict of promoted connections. - - Parameters - ---------- - conns : dict - Dictionary containing absolute connections. - - Returns - ------- - dict - Dictionary of promoted connections. - """ - abs2prom_in = self._var_allprocs_abs2prom['input'] - prom2abs_in = self._var_allprocs_prom2abs_list['input'] - prefix = self.pathname + '.' if self.pathname else '' - prom_conns = {} - for inp, out in conns.items(): - prom = abs2prom_in[inp] - prom_conns[prom] = (out, [i for i in prom2abs_in[prom] if i.startswith(prefix)]) - return prom_conns - - def _get_graph_display_info(self, display_map=None): - """ - Return display related metadata for this Group and all of its children. - - Parameters - ---------- - display_map : dict or None - A map of classnames to pydot node attributes. - - Returns - ------- - dict - Metadata keyed by system pathname. - """ - node_info = {} - for s in self.system_iter(recurse=True, include_self=True): - meta = s._get_graph_node_meta() - if display_map and meta['classname'] in display_map: - meta.update(display_map[meta['classname']]) - elif display_map and meta['base'] in display_map: - meta.update(display_map[meta['base']]) - else: - _get_node_display_meta(s, meta) - - ttlist = [f"Name: {s.pathname}"] - ttlist.append(f"Class: {meta['classname']}") - if s.linear_solver is not None and not isinstance(s.linear_solver, NonlinearRunOnce): - ttlist.append(f"Linear Solver: {type(s.linear_solver).__name__}") - if s.nonlinear_solver is not None and not isinstance(s.nonlinear_solver, - NonlinearRunOnce): - ttlist.append(f"Nonlinear Solver: {type(s.nonlinear_solver).__name__}") - meta['tooltip'] = '\n'.join(ttlist) - node_info[s.pathname] = meta.copy() - - if self.comm.size > 1: - abs2prom = self._var_abs2prom - all_abs2prom = self._var_allprocs_abs2prom - if (len(all_abs2prom['input']) != len(abs2prom['input']) or - len(all_abs2prom['output']) != len(abs2prom['output'])): - # not all systems exist in all procs, so must gather info from all procs - if self._gather_full_data(): - all_node_info = self.comm.allgather(node_info) - else: - all_node_info = self.comm.allgather({}) - - for info in all_node_info: - for pathname, meta in info.items(): - if pathname not in node_info: - node_info[pathname] = meta - - return node_info - - def _get_graph_node_meta(self): - """ - Return metadata to add to this system's graph node. - - Returns - ------- - dict - Metadata for this system's graph node. - """ - meta = super()._get_graph_node_meta() - # TODO: maybe set 'implicit' based on whether there are any implicit comps anywhere - # inside of the group or its children. - meta['base'] = 'Group' - return meta - - def _get_cluster_tree(self, node_info): - """ - Create a nested collection of pydot Cluster objects to represent the tree of groups. - - Parameters - ---------- - node_info : dict - A dict of metadata keyed by pathname. - - Returns - ------- - pydot.Dot, dict - The pydot graph and a dict of groups keyed by pathname. - """ - pydot_graph = pydot.Dot(graph_type='digraph') - mypath = self.pathname - prefix = mypath + '.' if mypath else '' - groups = {} - - if not mypath: - groups[''] = pydot.Cluster('', label='Model', style='filled', - fillcolor=_cluster_color(''), - tooltip=node_info['']['tooltip']) - pydot_graph.add_subgraph(groups['']) - - for varpath in chain(self._var_allprocs_abs2prom['input'], - self._var_allprocs_abs2prom['output']): - group = varpath.rpartition('.')[0].rpartition('.')[0] - if group not in groups: - # reverse the list so parents will exist before children - ancestor_list = list(all_ancestors(group))[::-1] - for path in ancestor_list: - if path.startswith(prefix) or path == mypath: - if path not in groups: - parent, _, name = path.rpartition('.') - groups[path] = pydot.Cluster(path, - label=path if path == mypath else name, - tooltip=node_info[path]['tooltip'], - fillcolor=_cluster_color(path), - style='filled') - if parent and parent.startswith(prefix): - groups[parent].add_subgraph(groups[path]) - elif parent == mypath and parent in groups: - groups[parent].add_subgraph(groups[path]) - else: - pydot_graph.add_subgraph(groups[path]) - - return pydot_graph, groups - - def _get_tree_graph(self, exclude, display_map=None): - """ - Create a pydot graph of the system tree (without clusters). - - Parameters - ---------- - exclude : iter of str - Iter of pathnames to exclude from the generated graph. - display_map : dict or None - A map of classnames to pydot node attributes. - - Returns - ------- - pydot.Dot - The pydot tree graph. - """ - node_info = self._get_graph_display_info(display_map) - exclude = set(exclude) - - systems = {} - pydot_graph = pydot.Dot(graph_type='graph', center=True) - prefix = self.pathname + '.' if self.pathname else '' - label = self.name if self.name else 'Model' - top_node = pydot.Node(label, label=label, - **node_info[self.pathname]) - pydot_graph.add_node(top_node) - systems[self.pathname] = top_node - - for varpath in chain(self._var_allprocs_abs2prom['input'], - self._var_allprocs_abs2prom['output']): - system = varpath.rpartition('.')[0] - if system not in systems and system not in exclude: - # reverse the list so parents will exist before children - ancestor_list = list(all_ancestors(system))[::-1] - for path in ancestor_list: - if path.startswith(prefix): - if path not in systems: - parent, _, name = path.rpartition('.') - kwargs = _filter_meta4dot(node_info[path]) - systems[path] = pydot.Node(path, label=name, **kwargs) - pydot_graph.add_node(systems[path]) - if parent.startswith(prefix) or parent == self.pathname: - pydot_graph.add_edge(pydot.Edge(systems[parent], systems[path])) - - return pydot_graph - - def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True, dvs=None, - responses=None): - """ - Add metadata to the graph for display. - - Returned graph will have any variable nodes containing certain characters relabeled with - explicit quotes to avoid issues with dot. - - Parameters - ---------- - G : nx.DiGraph - The graph to be decorated. - exclude : iter of str - Iter of pathnames to exclude from the generated graph. - abs_graph_names : bool - If True, use absolute pathnames for nodes in the graph. - dvs : dict - Dict of design var metadata keyed on promoted name. - responses : list of str - Dict of response var metadata keyed on promoted name. - - Returns - ------- - nx.DiGraph, dict - The decorated graph and a dict of node metadata keyed by pathname. - """ - node_info = self._get_graph_display_info() - - exclude = set(exclude) - - prefix = self.pathname + '.' if self.pathname else '' - - replace = {} - for node, meta in G.nodes(data=True): - if not abs_graph_names: - node = prefix + node - if node in node_info: - meta.update(_filter_meta4dot(node_info[node])) - if not ('label' in meta and meta['label']): - meta['label'] = f'"{node.rpartition(".")[2]}"' - else: - meta['label'] = f'"{meta["label"]}"' - if 'type_' in meta: # variable node - if node.rpartition('.')[0] in exclude: - exclude.add(node) # remove all variables of excluded components - # quote node names containing certain characters for use in dot - if (':' in node or '<' in node) and node not in exclude: - replace[node] = f'"{node}"' - if dvs and node in dvs: - meta['shape'] = 'cds' - elif responses and node in responses: - meta['shape'] = 'cds' - else: - meta['shape'] = 'plain' # just text for variables, otherwise too busy - - if replace: - G = nx.relabel_nodes(G, replace, copy=True) - - if exclude: - if not replace: - G = G.copy() - G.remove_nodes_from(exclude) - - return G, node_info - - def _apply_clusters(self, G, node_info): - """ - Group nodes in the graph into clusters. - - Parameters - ---------- - G : nx.DiGraph - A pydot graph will be created based on this graph. - node_info : dict - A dict of metadata keyed by pathname. - - Returns - ------- - pydot.Graph - The corresponding pydot graph with clusters added. - """ - pydot_graph, groups = self._get_cluster_tree(node_info) - prefix = self.pathname + '.' if self.pathname else '' - boundary_nodes = {'_Incoming', '_Outgoing'} - pydot_nodes = {} - for node, meta in G.nodes(data=True): - noquote_node = node.strip('"') - kwargs = _filter_meta4dot(meta) - if 'type_' in meta: # variable node - group = noquote_node.rpartition('.')[0].rpartition('.')[0] - else: - group = noquote_node.rpartition('.')[0] - - pdnode = pydot_nodes[node] = pydot.Node(node, **kwargs) - - if group and group.startswith(prefix): - groups[group].add_node(pdnode) - elif self.pathname in groups and node not in boundary_nodes: - groups[self.pathname].add_node(pdnode) - else: - pydot_graph.add_node(pdnode) - - for u, v, meta in G.edges(data=True): - pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], - **_filter_meta4dot(meta, - arrowhead='lnormal', - arrowsize=0.5))) - - # layout graph from left to right - pydot_graph.set_rankdir('LR') - - return pydot_graph - - def _get_boundary_conns(self): - """ - Return lists of incoming and outgoing boundary connections. - - Returns - ------- - tuple - A tuple of (incoming, outgoing) boundary connections. - """ - if not self.pathname: - return ([], []) - - top = self._problem_meta['model_ref']() - prefix = self.pathname + '.' - - incoming = [] - outgoing = [] - for abs_in, abs_out in top._conn_global_abs_in2out.items(): - if abs_in.startswith(prefix) and not abs_out.startswith(prefix): - incoming.append((abs_in, abs_out)) - if abs_out.startswith(prefix) and not abs_in.startswith(prefix): - outgoing.append((abs_in, abs_out)) - - return incoming, outgoing - - def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, - display=True, show_boundary=False, exclude=(), outfile=None): - """ - Use pydot to create a graphical representation of the specified graph. - - Parameters - ---------- - gtype : str - The type of graph to create. Options include 'system', 'component', 'nested', - and 'dataflow'. - recurse : bool - If True, recurse into subsystems when gtype is 'dataflow'. - show_vars : bool - If True, show all variables in the graph. Only relevant when gtype is 'dataflow'. - display : bool - If True, pop up a window to view the graph. - show_boundary : bool - If True, include connections to variables outside the boundary of the Group. - exclude : iter of str - Iter of pathnames to exclude from the generated graph. - outfile : str or None - The name of the file to write the graph to. The format is inferred from the extension. - Default is None, which writes to '__graph.svg', where system_path - is 'model' for the top level group, and any '.' in the pathname is replaced with '_'. - - Returns - ------- - pydot.Dot or None - The pydot graph that was created. - """ - if pydot is None: - issue_warning("To view the system graph, you must install pydot. " - "You can install it via 'pip install pydot'.") - return - - if gtype == 'tree': - G = self._get_tree_graph(exclude) - elif gtype == 'dataflow': - if show_vars: - # get dvs and responses so we can color them differently - dvs = self.get_design_vars(recurse=True, get_sizes=False, use_prom_ivc=True) - desvars = set(dvs) - desvars.update(m['source'] for m in dvs.values()) - resps = self.get_responses(recurse=True, get_sizes=False, use_prom_ivc=True) - responses = set(resps) - responses.update(m['source'] for m in resps.values()) - - prefix = self.pathname + '.' if self.pathname else '' - lenpre = len(prefix) - - G = self._problem_meta['model_ref']()._get_dataflow_graph() - - if not recurse: - # layout graph from left to right - gname = 'model' if self.pathname == '' else self.pathname - G.graph['graph'] = {'rankdir': 'LR', 'label': f"Dataflow for '{gname}'", - 'center': 'True'} - - # keep all direct children and their variables - keep = {n for n in G.nodes() if n[lenpre:].count('.') == 0 and - n.startswith(prefix)} - keep.update({n for n, d in G.nodes(data=True) if 'type_' in d and - n.rpartition('.')[0] in keep}) - - promconns = self._get_prom_conns(self._conn_abs_in2out) - - for prom_in, (abs_out, abs_ins) in promconns.items(): - nins = len(abs_ins) - the_in = abs_ins[0] if nins == 1 else prom_in - if the_in not in G: - G.add_node(the_in, type_='input', label=prom_in) - else: - label = prom_in[:lenpre] if prom_in.startswith(prefix) else prom_in - G.nodes[the_in]['label'] = label - - sysout = prefix + abs_out[lenpre:].partition('.')[0] - prom_out = self._var_allprocs_abs2prom['output'][abs_out] - - if sysout not in G: - G.add_node(sysout, **_base_display_map['Group']) - - label = prom_out[:lenpre] if prom_out.startswith(prefix) else prom_out - G.nodes[abs_out]['label'] = label - G.add_edge(sysout, abs_out) - - keep.add(sysout) - keep.add(abs_out) - - for abs_in in abs_ins: - sysin = prefix + abs_in[lenpre:].partition('.')[0] - if sysin not in G: - G.add_node(sysin, **_base_display_map['Group']) - if nins == 1: - G.add_edge(abs_ins[0], sysin) - keep.add(abs_ins[0]) - else: - G.add_edge(prom_in, sysin) - keep.add(prom_in) - - if prom_in in G and nins > 1: - G.nodes[prom_in]['fontcolor'] = 'red' - G.nodes[prom_in]['tooltip'] = '\n'.join(abs_ins) - - if self.pathname == '': - if not recurse: - G = nx.subgraph(G, keep) - else: - # we're not the top level group, so get our subgraph of the top level graph - ournodes = {n for n in G.nodes() if n.startswith(prefix)} - - if not recurse: - ournodes.update(keep) - - G = nx.subgraph(G, ournodes) - - if show_boundary and self.pathname: - incoming, outgoing = self._get_boundary_conns() - G = _add_boundary_nodes(self.pathname, G.copy(), incoming, outgoing, exclude) - - G, node_info = self._decorate_graph_for_display(G, exclude=exclude, - dvs=desvars, responses=responses) - - if recurse: - G = self._apply_clusters(G, node_info) - else: - G = _to_pydot_graph(G) - - elif recurse: - G = self.compute_sys_graph(comps_only=True, add_edge_info=False) - if show_boundary and self.pathname: - incoming, outgoing = self._get_boundary_conns() - # convert var abs names to system abs names - incoming = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) - for in_abs, out_abs in incoming] - outgoing = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) - for in_abs, out_abs in outgoing] - G = _add_boundary_nodes(self.pathname, G.copy(), incoming, outgoing, exclude) - - G, node_info = self._decorate_graph_for_display(G, exclude=exclude) - G = self._apply_clusters(G, node_info) - else: - G = self.compute_sys_graph(comps_only=False, add_edge_info=False) - G, _ = self._decorate_graph_for_display(G, exclude=exclude, abs_graph_names=False) - G = _to_pydot_graph(G) - else: - raise ValueError(f"unrecognized graph type '{gtype}'. Allowed types are ['tree', " - "'dataflow'].") - - if G is None: - return - - if outfile is None: - name = self.pathname.replace('.', '_') if self.pathname else 'model' - outfile = f"{name}_{gtype}_graph.svg" - - return write_graph(G, prog='dot', display=display, outfile=outfile) - def _get_auto_ivc_out_val(self, tgts, vars_to_gather): # all tgts are continuous variables # only called from top level group @@ -5698,11 +5187,17 @@ def _active_responses(self, user_response_names, responses=None): return active_resps + def _get_graph_node_meta(self): + """ + Return metadata to add to this system's graph node. -def _get_node_display_meta(s, meta): - if meta['base'] in _base_display_map: - meta.update(_base_display_map[meta['base']]) - if s.nonlinear_solver is not None and not isinstance(s.nonlinear_solver, NonlinearRunOnce): - meta['shape'] = 'box3d' - elif s.linear_solver is not None and not isinstance(s.linear_solver, LinearRunOnce): - meta['shape'] = 'box3d' + Returns + ------- + dict + Metadata for this system's graph node. + """ + meta = super()._get_graph_node_meta() + # TODO: maybe set 'implicit' based on whether there are any implicit comps anywhere + # inside of the group or its children. + meta['base'] = 'Group' + return meta diff --git a/openmdao/utils/graph_utils.py b/openmdao/utils/graph_utils.py index de7bd5ee94..1c04fa0538 100644 --- a/openmdao/utils/graph_utils.py +++ b/openmdao/utils/graph_utils.py @@ -2,13 +2,6 @@ Various graph related utilities. """ import networkx as nx -from openmdao.utils.file_utils import _load_and_exec -import openmdao.utils.hooks as hooks - -try: - import pydot -except ImportError: - pydot = None def get_sccs_topo(graph): @@ -61,244 +54,3 @@ def get_out_of_order_nodes(graph, orders): out_of_order.append((u, v)) return strongcomps, out_of_order - - -def write_graph(G, prog='dot', display=True, outfile='graph.svg'): - """ - Write the graph to a file and optionally display it. - - Parameters - ---------- - G : nx.DiGraph or pydot.Dot - The graph to be written. - prog : str - The graphviz program to use for layout. - display : bool - If True, display the graph after writing it. - outfile : str - The name of the file to write. - - Returns - ------- - pydot.Dot - The graph that was written. - """ - from openmdao.utils.webview import webview - - if pydot is None: - raise RuntimeError("graph requires the pydot package. You can install it using " - "'pip install pydot'.") - - ext = outfile.rpartition('.')[2] - if not ext: - ext = 'svg' - - if isinstance(G, nx.Graph): - pydot_graph = nx.drawing.nx_pydot.to_pydot(G) - else: - pydot_graph = G - - try: - pstr = getattr(pydot_graph, f"create_{ext}")(prog=prog) - except AttributeError: - raise AttributeError(f"pydot graph has no 'create_{ext}' method.") - - with open(outfile, 'wb') as f: - f.write(pstr) - - if display: - webview(outfile) - - return pydot_graph - - -def _graph_setup_parser(parser): - """ - Set up the openmdao subparser for the 'openmdao graph' command. - - Parameters - ---------- - parser : argparse subparser - The parser we're adding options to. - """ - parser.add_argument('file', nargs=1, help='Python file containing the model.') - parser.add_argument('-p', '--problem', action='store', dest='problem', help='Problem name') - parser.add_argument('-o', action='store', dest='outfile', help='file containing graph output.') - parser.add_argument('--group', action='store', dest='group', help='pathname of group to graph.') - parser.add_argument('--type', action='store', dest='type', default='dataflow', - help='type of graph (dataflow, tree). Default is dataflow.') - parser.add_argument('--no-display', action='store_false', dest='show', - help="don't display the graph.") - parser.add_argument('--no-recurse', action='store_false', dest='recurse', - help="don't recurse from the specified group down. This only applies to " - "the dataflow graph type.") - parser.add_argument('--show-vars', action='store_true', dest='show_vars', - help="show variables in the graph. This only applies to the dataflow graph." - " Default is False.") - parser.add_argument('--show-boundary', action='store_true', dest='show_boundary', - help="show connections to variables outside of the graph. This only " - "applies to the dataflow graph. Default is False.") - parser.add_argument('--autoivc', action='store_true', dest='auto_ivc', - help="include the _auto_ivc component in the graph. This applies to " - "graphs of the top level group only. Default is False.") - - -def _graph_cmd(options, user_args): - """ - Return the post_setup hook function for 'openmdao graph'. - - Parameters - ---------- - options : argparse Namespace - Command line options. - user_args : list of str - Args to be passed to the user script. - """ - def _view_graph(problem): - group = problem.model._get_subsystem(options.group) if options.group else problem.model - if not options.auto_ivc: - exclude = {'_auto_ivc'} - else: - exclude = set() - group.write_graph(gtype=options.type, recurse=options.recurse, - show_vars=options.show_vars, display=options.show, exclude=exclude, - show_boundary=options.show_boundary, outfile=options.outfile) - - # register the hooks - hooks._register_hook('final_setup', 'Problem', post=_view_graph, exit=True) - _load_and_exec(options.file[0], user_args) - - -def _to_pydot_graph(G): - gmeta = G.graph.get('graph', {}).copy() - gmeta['graph_type'] = 'digraph' - pydot_graph = pydot.Dot(**gmeta) - pydot_nodes = {} - - for node, meta in G.nodes(data=True): - pdnode = pydot_nodes[node] = pydot.Node(node, **_filter_meta4dot(meta)) - pydot_graph.add_node(pdnode) - - for u, v, meta in G.edges(data=True): - pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], - **_filter_meta4dot(meta, arrowsize=0.5))) - - # layout graph from left to right - pydot_graph.set_rankdir('LR') - - return pydot_graph - - -def _filter_meta4dot(meta, **kwargs): - """ - Remove unnecessary metadata from the given metadata dict before passing to pydot. - - Parameters - ---------- - meta : dict - Metadata dict. - kwargs : dict - Additional metadata that will be added only if they are not already present. - - Returns - ------- - dict - Metadata dict with unnecessary items removed. - """ - skip = {'type_', 'local', 'base', 'classname'} - dct = {k: v for k, v in meta.items() if k not in skip} - for k, v in kwargs.items(): - if k not in dct: - dct[k] = v - return dct - - -def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): - """ - Add boundary nodes to the graph. - - Parameters - ---------- - pathname : str - Pathname of the current group. - G : nx.DiGraph - The graph. - incoming : list of (str, str) - List of incoming connections. - outgoing : list of (str, str) - List of outgoing connections. - exclude : list of str - List of pathnames to exclude from the graph. - - Returns - ------- - nx.DiGraph - The modified graph. - """ - lenpre = len(pathname) + 1 if pathname else 0 - for ex in exclude: - expre = ex + '.' - incoming = [(in_abs, out_abs) for in_abs, out_abs in incoming - if in_abs != ex and out_abs != ex and - not in_abs.startswith(expre) and not out_abs.startswith(expre)] - outgoing = [(in_abs, out_abs) for in_abs, out_abs in outgoing - if in_abs != ex and out_abs != ex and - not in_abs.startswith(expre) and not out_abs.startswith(expre)] - - if incoming: - tooltip = ['External Connections:'] - connstrs = set() - for in_abs, out_abs in incoming: - if in_abs in G: - connstrs.add(f" {out_abs} -> {in_abs[lenpre:]}") - tooltip += sorted(connstrs) - tooltip = '\n'.join(tooltip) - if connstrs: - G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='min') - for in_abs, out_abs in incoming: - if in_abs in G: - G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', - arrowsize=0.5) - - if outgoing: - tooltip = ['External Connections:'] - connstrs = set() - for in_abs, out_abs in outgoing: - if out_abs in G: - connstrs.add(f" {out_abs[lenpre:]} -> {in_abs}") - tooltip += sorted(connstrs) - tooltip = '\n'.join(tooltip) - G.add_node('_Outgoing', label='Outgoing', arrowhead='rarrow', fillcolor='peachpuff3', - style='filled', tooltip=f'"{tooltip}"', rank='max') - for in_abs, out_abs in outgoing: - if out_abs in G: - G.add_edge(out_abs, '_Outgoing', style='dashed', shape='lnormal', arrowsize=0.5) - - return G - - -def _cluster_color(path): - """ - Return the color of the cluster that contains the given path. - - The idea here is to make nested clusters stand out wrt their parent cluster. - - Parameters - ---------- - path : str - Pathname of a variable. - - Returns - ------- - int - The color of the cluster that contains the given path. - """ - depth = path.count('.') + 1 if path else 0 - - ncolors = 10 - maxcolor = 98 - mincolor = 40 - - col = maxcolor - (depth % ncolors) * (maxcolor - mincolor) // ncolors - return f"gray{col}" diff --git a/openmdao/utils/om.py b/openmdao/utils/om.py index 9cd6b5075e..443ecfd209 100644 --- a/openmdao/utils/om.py +++ b/openmdao/utils/om.py @@ -61,7 +61,7 @@ _find_repos_setup_parser, _find_repos_exec from openmdao.utils.reports_system import _list_reports_setup_parser, _list_reports_cmd, \ _view_reports_setup_parser, _view_reports_cmd -from openmdao.utils.graph_utils import _graph_setup_parser, _graph_cmd +from openmdao.visualization.graph_viewer import _graph_setup_parser, _graph_cmd def _view_connections_setup_parser(parser): diff --git a/openmdao/visualization/graph_viewer.py b/openmdao/visualization/graph_viewer.py new file mode 100644 index 0000000000..7bfa176d75 --- /dev/null +++ b/openmdao/visualization/graph_viewer.py @@ -0,0 +1,792 @@ +""" +Viewer graphs of a group's model hierarchy and connections. +""" +from itertools import chain + +try: + import pydot +except ImportError: + pydot = None + +import networkx as nx + +from openmdao.solvers.nonlinear.nonlinear_runonce import NonlinearRunOnce +from openmdao.solvers.linear.linear_runonce import LinearRunOnce +from openmdao.utils.om_warnings import issue_warning +from openmdao.utils.general_utils import all_ancestors +from openmdao.utils.file_utils import _load_and_exec +import openmdao.utils.hooks as hooks + + +# mapping of base system type to graph display properties +_base_display_map = { + 'ExplicitComponent': { + 'fillcolor': '"aquamarine3:aquamarine"', + 'style': 'filled', + 'shape': 'box', + }, + 'ImplicitComponent': { + 'fillcolor': '"lightblue:lightslateblue"', + 'style': 'filled', + 'shape': 'box', + }, + 'IndepVarComp': { + 'fillcolor': '"chartreuse2:chartreuse4"', + 'style': 'filled', + 'shape': 'box', + }, + 'Group': { + 'fillcolor': 'gray75', + 'style': 'filled', + 'shape': 'box', + }, +} + + +class GraphViewer(object): + """ + A class for viewing the model hierarchy and connections in a group. + + Parameters + ---------- + group : + The Group with graph to be viewed. + + Attributes + ---------- + _group : + The Group with graph to be viewed. + """ + + def __init__(self, group): + """ + Initialize the GraphViewer. + + Parameters + ---------- + group : + The Group with graph to be viewed. + """ + self._group = group + + def write_graph(self, gtype='dataflow', recurse=True, show_vars=False, + display=True, show_boundary=False, exclude=(), outfile=None): + """ + Use pydot to create a graphical representation of the specified graph. + + Parameters + ---------- + gtype : str + The type of graph to create. Options include 'system', 'component', 'nested', + and 'dataflow'. + recurse : bool + If True, recurse into subsystems when gtype is 'dataflow'. + show_vars : bool + If True, show all variables in the graph. Only relevant when gtype is 'dataflow'. + display : bool + If True, pop up a window to view the graph. + show_boundary : bool + If True, include connections to variables outside the boundary of the Group. + exclude : iter of str + Iter of pathnames to exclude from the generated graph. + outfile : str or None + The name of the file to write the graph to. The format is inferred from the extension. + Default is None, which writes to '__graph.svg', where system_path + is 'model' for the top level group, and any '.' in the pathname is replaced with '_'. + + Returns + ------- + pydot.Dot or None + The pydot graph that was created. + """ + if pydot is None: + issue_warning("To view the system graph, you must install pydot. " + "You can install it via 'pip install pydot'.") + return + + group = self._group + + if gtype == 'tree': + G = self._get_tree_graph(exclude) + elif gtype == 'dataflow': + if show_vars: + # get dvs and responses so we can color them differently + dvs = group.get_design_vars(recurse=True, get_sizes=False, use_prom_ivc=True) + desvars = set(dvs) + desvars.update(m['source'] for m in dvs.values()) + resps = group.get_responses(recurse=True, get_sizes=False, use_prom_ivc=True) + responses = set(resps) + responses.update(m['source'] for m in resps.values()) + + prefix = group.pathname + '.' if group.pathname else '' + lenpre = len(prefix) + + G = group._problem_meta['model_ref']()._get_dataflow_graph() + + if not recurse: + # layout graph from left to right + gname = 'model' if group.pathname == '' else group.pathname + G.graph['graph'] = {'rankdir': 'LR', 'label': f"Dataflow for '{gname}'", + 'center': 'True'} + + # keep all direct children and their variables + keep = {n for n in G.nodes() if n[lenpre:].count('.') == 0 and + n.startswith(prefix)} + keep.update({n for n, d in G.nodes(data=True) if 'type_' in d and + n.rpartition('.')[0] in keep}) + + promconns = self._get_prom_conns(group._conn_abs_in2out) + + for prom_in, (abs_out, abs_ins) in promconns.items(): + nins = len(abs_ins) + the_in = abs_ins[0] if nins == 1 else prom_in + if the_in not in G: + G.add_node(the_in, type_='input', label=prom_in) + else: + label = prom_in[:lenpre] if prom_in.startswith(prefix) else prom_in + G.nodes[the_in]['label'] = label + + sysout = prefix + abs_out[lenpre:].partition('.')[0] + prom_out = group._var_allprocs_abs2prom['output'][abs_out] + + if sysout not in G: + G.add_node(sysout, **_base_display_map['Group']) + + label = prom_out[:lenpre] if prom_out.startswith(prefix) else prom_out + G.nodes[abs_out]['label'] = label + G.add_edge(sysout, abs_out) + + keep.add(sysout) + keep.add(abs_out) + + for abs_in in abs_ins: + sysin = prefix + abs_in[lenpre:].partition('.')[0] + if sysin not in G: + G.add_node(sysin, **_base_display_map['Group']) + if nins == 1: + G.add_edge(abs_ins[0], sysin) + keep.add(abs_ins[0]) + else: + G.add_edge(prom_in, sysin) + keep.add(prom_in) + + if prom_in in G and nins > 1: + G.nodes[prom_in]['fontcolor'] = 'red' + G.nodes[prom_in]['tooltip'] = '\n'.join(abs_ins) + + if group.pathname == '': + if not recurse: + G = nx.subgraph(G, keep) + else: + # we're not the top level group, so get our subgraph of the top level graph + ournodes = {n for n in G.nodes() if n.startswith(prefix)} + + if not recurse: + ournodes.update(keep) + + G = nx.subgraph(G, ournodes) + + if show_boundary and group.pathname: + incoming, outgoing = self._get_boundary_conns() + G = _add_boundary_nodes(group.pathname, G.copy(), incoming, outgoing, exclude) + + G, node_info = self._decorate_graph_for_display(G, exclude=exclude, + dvs=desvars, responses=responses) + + if recurse: + G = self._apply_clusters(G, node_info) + else: + G = _to_pydot_graph(G) + + elif recurse: + G = group.compute_sys_graph(comps_only=True, add_edge_info=False) + if show_boundary and group.pathname: + incoming, outgoing = self._get_boundary_conns() + # convert var abs names to system abs names + incoming = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) + for in_abs, out_abs in incoming] + outgoing = [(in_abs.rpartition('.')[0], out_abs.rpartition('.')[0]) + for in_abs, out_abs in outgoing] + G = _add_boundary_nodes(group.pathname, G.copy(), incoming, outgoing, exclude) + + G, node_info = self._decorate_graph_for_display(G, exclude=exclude) + G = self._apply_clusters(G, node_info) + else: + G = group.compute_sys_graph(comps_only=False, add_edge_info=False) + G, _ = self._decorate_graph_for_display(G, exclude=exclude, abs_graph_names=False) + G = _to_pydot_graph(G) + else: + raise ValueError(f"unrecognized graph type '{gtype}'. Allowed types are ['tree', " + "'dataflow'].") + + if G is None: + return + + if outfile is None: + name = group.pathname.replace('.', '_') if group.pathname else 'model' + outfile = f"{name}_{gtype}_graph.svg" + + return write_graph(G, prog='dot', display=display, outfile=outfile) + + def _get_prom_conns(self, conns): + """ + Return a dict of promoted connections. + + Parameters + ---------- + conns : dict + Dictionary containing absolute connections. + + Returns + ------- + dict + Dictionary of promoted connections. + """ + group = self._group + abs2prom_in = group._var_allprocs_abs2prom['input'] + prom2abs_in = group._var_allprocs_prom2abs_list['input'] + prefix = group.pathname + '.' if group.pathname else '' + prom_conns = {} + for inp, out in conns.items(): + prom = abs2prom_in[inp] + prom_conns[prom] = (out, [i for i in prom2abs_in[prom] if i.startswith(prefix)]) + return prom_conns + + def _get_graph_display_info(self, display_map=None): + """ + Return display related metadata for this Group and all of its children. + + Parameters + ---------- + display_map : dict or None + A map of classnames to pydot node attributes. + + Returns + ------- + dict + Metadata keyed by system pathname. + """ + group = self._group + node_info = {} + for s in group.system_iter(recurse=True, include_self=True): + meta = s._get_graph_node_meta() + if display_map and meta['classname'] in display_map: + meta.update(display_map[meta['classname']]) + elif display_map and meta['base'] in display_map: + meta.update(display_map[meta['base']]) + else: + _get_node_display_meta(s, meta) + + ttlist = [f"Name: {s.pathname}"] + ttlist.append(f"Class: {meta['classname']}") + if s.linear_solver is not None and not isinstance(s.linear_solver, NonlinearRunOnce): + ttlist.append(f"Linear Solver: {type(s.linear_solver).__name__}") + if s.nonlinear_solver is not None and not isinstance(s.nonlinear_solver, + NonlinearRunOnce): + ttlist.append(f"Nonlinear Solver: {type(s.nonlinear_solver).__name__}") + meta['tooltip'] = '\n'.join(ttlist) + node_info[s.pathname] = meta.copy() + + if group.comm.size > 1: + abs2prom = group._var_abs2prom + all_abs2prom = group._var_allprocs_abs2prom + if (len(all_abs2prom['input']) != len(abs2prom['input']) or + len(all_abs2prom['output']) != len(abs2prom['output'])): + # not all systems exist in all procs, so must gather info from all procs + if group._gather_full_data(): + all_node_info = group.comm.allgather(node_info) + else: + all_node_info = group.comm.allgather({}) + + for info in all_node_info: + for pathname, meta in info.items(): + if pathname not in node_info: + node_info[pathname] = meta + + return node_info + + def _get_cluster_tree(self, node_info): + """ + Create a nested collection of pydot Cluster objects to represent the tree of groups. + + Parameters + ---------- + node_info : dict + A dict of metadata keyed by pathname. + + Returns + ------- + pydot.Dot, dict + The pydot graph and a dict of groups keyed by pathname. + """ + group = self._group + pydot_graph = pydot.Dot(graph_type='digraph') + mypath = group.pathname + prefix = mypath + '.' if mypath else '' + groups = {} + + if not mypath: + groups[''] = pydot.Cluster('', label='Model', style='filled', + fillcolor=_cluster_color(''), + tooltip=node_info['']['tooltip']) + pydot_graph.add_subgraph(groups['']) + + for varpath in chain(group._var_allprocs_abs2prom['input'], + group._var_allprocs_abs2prom['output']): + group = varpath.rpartition('.')[0].rpartition('.')[0] + if group not in groups: + # reverse the list so parents will exist before children + ancestor_list = list(all_ancestors(group))[::-1] + for path in ancestor_list: + if path.startswith(prefix) or path == mypath: + if path not in groups: + parent, _, name = path.rpartition('.') + groups[path] = pydot.Cluster(path, + label=path if path == mypath else name, + tooltip=node_info[path]['tooltip'], + fillcolor=_cluster_color(path), + style='filled') + if parent and parent.startswith(prefix): + groups[parent].add_subgraph(groups[path]) + elif parent == mypath and parent in groups: + groups[parent].add_subgraph(groups[path]) + else: + pydot_graph.add_subgraph(groups[path]) + + return pydot_graph, groups + + def _get_tree_graph(self, exclude, display_map=None): + """ + Create a pydot graph of the system tree (without clusters). + + Parameters + ---------- + exclude : iter of str + Iter of pathnames to exclude from the generated graph. + display_map : dict or None + A map of classnames to pydot node attributes. + + Returns + ------- + pydot.Dot + The pydot tree graph. + """ + node_info = self._get_graph_display_info(display_map) + exclude = set(exclude) + + group = self._group + systems = {} + pydot_graph = pydot.Dot(graph_type='graph', center=True) + prefix = group.pathname + '.' if group.pathname else '' + label = group.name if group.name else 'Model' + top_node = pydot.Node(label, label=label, + **node_info[group.pathname]) + pydot_graph.add_node(top_node) + systems[group.pathname] = top_node + + for varpath in chain(group._var_allprocs_abs2prom['input'], + group._var_allprocs_abs2prom['output']): + system = varpath.rpartition('.')[0] + if system not in systems and system not in exclude: + # reverse the list so parents will exist before children + ancestor_list = list(all_ancestors(system))[::-1] + for path in ancestor_list: + if path.startswith(prefix): + if path not in systems: + parent, _, name = path.rpartition('.') + kwargs = _filter_meta4dot(node_info[path]) + systems[path] = pydot.Node(path, label=name, **kwargs) + pydot_graph.add_node(systems[path]) + if parent.startswith(prefix) or parent == group.pathname: + pydot_graph.add_edge(pydot.Edge(systems[parent], systems[path])) + + return pydot_graph + + def _decorate_graph_for_display(self, G, exclude=(), abs_graph_names=True, dvs=None, + responses=None): + """ + Add metadata to the graph for display. + + Returned graph will have any variable nodes containing certain characters relabeled with + explicit quotes to avoid issues with dot. + + Parameters + ---------- + G : nx.DiGraph + The graph to be decorated. + exclude : iter of str + Iter of pathnames to exclude from the generated graph. + abs_graph_names : bool + If True, use absolute pathnames for nodes in the graph. + dvs : dict + Dict of design var metadata keyed on promoted name. + responses : list of str + Dict of response var metadata keyed on promoted name. + + Returns + ------- + nx.DiGraph, dict + The decorated graph and a dict of node metadata keyed by pathname. + """ + node_info = self._get_graph_display_info() + + exclude = set(exclude) + + prefix = self._group.pathname + '.' if self._group.pathname else '' + + replace = {} + for node, meta in G.nodes(data=True): + if not abs_graph_names: + node = prefix + node + if node in node_info: + meta.update(_filter_meta4dot(node_info[node])) + if not ('label' in meta and meta['label']): + meta['label'] = f'"{node.rpartition(".")[2]}"' + else: + meta['label'] = f'"{meta["label"]}"' + if 'type_' in meta: # variable node + if node.rpartition('.')[0] in exclude: + exclude.add(node) # remove all variables of excluded components + # quote node names containing certain characters for use in dot + if (':' in node or '<' in node) and node not in exclude: + replace[node] = f'"{node}"' + if dvs and node in dvs: + meta['shape'] = 'cds' + elif responses and node in responses: + meta['shape'] = 'cds' + else: + meta['shape'] = 'plain' # just text for variables, otherwise too busy + + if replace: + G = nx.relabel_nodes(G, replace, copy=True) + + if exclude: + if not replace: + G = G.copy() + G.remove_nodes_from(exclude) + + return G, node_info + + def _apply_clusters(self, G, node_info): + """ + Group nodes in the graph into clusters. + + Parameters + ---------- + G : nx.DiGraph + A pydot graph will be created based on this graph. + node_info : dict + A dict of metadata keyed by pathname. + + Returns + ------- + pydot.Graph + The corresponding pydot graph with clusters added. + """ + pydot_graph, groups = self._get_cluster_tree(node_info) + prefix = self._group.pathname + '.' if self._group.pathname else '' + boundary_nodes = {'_Incoming', '_Outgoing'} + pydot_nodes = {} + for node, meta in G.nodes(data=True): + noquote_node = node.strip('"') + kwargs = _filter_meta4dot(meta) + if 'type_' in meta: # variable node + group = noquote_node.rpartition('.')[0].rpartition('.')[0] + else: + group = noquote_node.rpartition('.')[0] + + pdnode = pydot_nodes[node] = pydot.Node(node, **kwargs) + + if group and group.startswith(prefix): + groups[group].add_node(pdnode) + elif self._group.pathname in groups and node not in boundary_nodes: + groups[self._group.pathname].add_node(pdnode) + else: + pydot_graph.add_node(pdnode) + + for u, v, meta in G.edges(data=True): + pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], + **_filter_meta4dot(meta, + arrowhead='lnormal', + arrowsize=0.5))) + + # layout graph from left to right + pydot_graph.set_rankdir('LR') + + return pydot_graph + + def _get_boundary_conns(self): + """ + Return lists of incoming and outgoing boundary connections. + + Returns + ------- + tuple + A tuple of (incoming, outgoing) boundary connections. + """ + if not self._group.pathname: + return ([], []) + + top = self._group._problem_meta['model_ref']() + prefix = self._group.pathname + '.' + + incoming = [] + outgoing = [] + for abs_in, abs_out in top._conn_global_abs_in2out.items(): + if abs_in.startswith(prefix) and not abs_out.startswith(prefix): + incoming.append((abs_in, abs_out)) + if abs_out.startswith(prefix) and not abs_in.startswith(prefix): + outgoing.append((abs_in, abs_out)) + + return incoming, outgoing + + +def _get_node_display_meta(s, meta): + if meta['base'] in _base_display_map: + meta.update(_base_display_map[meta['base']]) + if s.nonlinear_solver is not None and not isinstance(s.nonlinear_solver, NonlinearRunOnce): + meta['shape'] = 'box3d' + elif s.linear_solver is not None and not isinstance(s.linear_solver, LinearRunOnce): + meta['shape'] = 'box3d' + + +def write_graph(G, prog='dot', display=True, outfile='graph.svg'): + """ + Write the graph to a file and optionally display it. + + Parameters + ---------- + G : nx.DiGraph or pydot.Dot + The graph to be written. + prog : str + The graphviz program to use for layout. + display : bool + If True, display the graph after writing it. + outfile : str + The name of the file to write. + + Returns + ------- + pydot.Dot + The graph that was written. + """ + from openmdao.utils.webview import webview + + if pydot is None: + raise RuntimeError("graph requires the pydot package. You can install it using " + "'pip install pydot'.") + + ext = outfile.rpartition('.')[2] + if not ext: + ext = 'svg' + + if isinstance(G, nx.Graph): + pydot_graph = nx.drawing.nx_pydot.to_pydot(G) + else: + pydot_graph = G + + try: + pstr = getattr(pydot_graph, f"create_{ext}")(prog=prog) + except AttributeError: + raise AttributeError(f"pydot graph has no 'create_{ext}' method.") + + with open(outfile, 'wb') as f: + f.write(pstr) + + if display: + webview(outfile) + + return pydot_graph + + +def _to_pydot_graph(G): + gmeta = G.graph.get('graph', {}).copy() + gmeta['graph_type'] = 'digraph' + pydot_graph = pydot.Dot(**gmeta) + pydot_nodes = {} + + for node, meta in G.nodes(data=True): + pdnode = pydot_nodes[node] = pydot.Node(node, **_filter_meta4dot(meta)) + pydot_graph.add_node(pdnode) + + for u, v, meta in G.edges(data=True): + pydot_graph.add_edge(pydot.Edge(pydot_nodes[u], pydot_nodes[v], + **_filter_meta4dot(meta, arrowsize=0.5))) + + # layout graph from left to right + pydot_graph.set_rankdir('LR') + + return pydot_graph + + +def _filter_meta4dot(meta, **kwargs): + """ + Remove unnecessary metadata from the given metadata dict before passing to pydot. + + Parameters + ---------- + meta : dict + Metadata dict. + kwargs : dict + Additional metadata that will be added only if they are not already present. + + Returns + ------- + dict + Metadata dict with unnecessary items removed. + """ + skip = {'type_', 'local', 'base', 'classname'} + dct = {k: v for k, v in meta.items() if k not in skip} + for k, v in kwargs.items(): + if k not in dct: + dct[k] = v + return dct + + +def _add_boundary_nodes(pathname, G, incoming, outgoing, exclude=()): + """ + Add boundary nodes to the graph. + + Parameters + ---------- + pathname : str + Pathname of the current group. + G : nx.DiGraph + The graph. + incoming : list of (str, str) + List of incoming connections. + outgoing : list of (str, str) + List of outgoing connections. + exclude : list of str + List of pathnames to exclude from the graph. + + Returns + ------- + nx.DiGraph + The modified graph. + """ + lenpre = len(pathname) + 1 if pathname else 0 + for ex in exclude: + expre = ex + '.' + incoming = [(in_abs, out_abs) for in_abs, out_abs in incoming + if in_abs != ex and out_abs != ex and + not in_abs.startswith(expre) and not out_abs.startswith(expre)] + outgoing = [(in_abs, out_abs) for in_abs, out_abs in outgoing + if in_abs != ex and out_abs != ex and + not in_abs.startswith(expre) and not out_abs.startswith(expre)] + + if incoming: + tooltip = ['External Connections:'] + connstrs = set() + for in_abs, out_abs in incoming: + if in_abs in G: + connstrs.add(f" {out_abs} -> {in_abs[lenpre:]}") + tooltip += sorted(connstrs) + tooltip = '\n'.join(tooltip) + if connstrs: + G.add_node('_Incoming', label='Incoming', shape='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='min') + for in_abs, out_abs in incoming: + if in_abs in G: + G.add_edge('_Incoming', in_abs, style='dashed', arrowhead='lnormal', + arrowsize=0.5) + + if outgoing: + tooltip = ['External Connections:'] + connstrs = set() + for in_abs, out_abs in outgoing: + if out_abs in G: + connstrs.add(f" {out_abs[lenpre:]} -> {in_abs}") + tooltip += sorted(connstrs) + tooltip = '\n'.join(tooltip) + G.add_node('_Outgoing', label='Outgoing', arrowhead='rarrow', fillcolor='peachpuff3', + style='filled', tooltip=f'"{tooltip}"', rank='max') + for in_abs, out_abs in outgoing: + if out_abs in G: + G.add_edge(out_abs, '_Outgoing', style='dashed', shape='lnormal', arrowsize=0.5) + + return G + + +def _cluster_color(path): + """ + Return the color of the cluster that contains the given path. + + The idea here is to make nested clusters stand out wrt their parent cluster. + + Parameters + ---------- + path : str + Pathname of a variable. + + Returns + ------- + int + The color of the cluster that contains the given path. + """ + depth = path.count('.') + 1 if path else 0 + + ncolors = 10 + maxcolor = 98 + mincolor = 40 + + col = maxcolor - (depth % ncolors) * (maxcolor - mincolor) // ncolors + return f"gray{col}" + + +def _graph_setup_parser(parser): + """ + Set up the openmdao subparser for the 'openmdao graph' command. + + Parameters + ---------- + parser : argparse subparser + The parser we're adding options to. + """ + parser.add_argument('file', nargs=1, help='Python file containing the model.') + parser.add_argument('-p', '--problem', action='store', dest='problem', help='Problem name') + parser.add_argument('-o', action='store', dest='outfile', help='file containing graph output.') + parser.add_argument('--group', action='store', dest='group', help='pathname of group to graph.') + parser.add_argument('--type', action='store', dest='type', default='dataflow', + help='type of graph (dataflow, tree). Default is dataflow.') + parser.add_argument('--no-display', action='store_false', dest='show', + help="don't display the graph.") + parser.add_argument('--no-recurse', action='store_false', dest='recurse', + help="don't recurse from the specified group down. This only applies to " + "the dataflow graph type.") + parser.add_argument('--show-vars', action='store_true', dest='show_vars', + help="show variables in the graph. This only applies to the dataflow graph." + " Default is False.") + parser.add_argument('--show-boundary', action='store_true', dest='show_boundary', + help="show connections to variables outside of the graph. This only " + "applies to the dataflow graph. Default is False.") + parser.add_argument('--autoivc', action='store_true', dest='auto_ivc', + help="include the _auto_ivc component in the graph. This applies to " + "graphs of the top level group only. Default is False.") + + +def _graph_cmd(options, user_args): + """ + Return the post_setup hook function for 'openmdao graph'. + + Parameters + ---------- + options : argparse Namespace + Command line options. + user_args : list of str + Args to be passed to the user script. + """ + def _view_graph(problem): + group = problem.model._get_subsystem(options.group) if options.group else problem.model + if not options.auto_ivc: + exclude = {'_auto_ivc'} + else: + exclude = set() + GraphViewer(group).write_graph(gtype=options.type, recurse=options.recurse, + show_vars=options.show_vars, display=options.show, + exclude=exclude, show_boundary=options.show_boundary, + outfile=options.outfile) + + # register the hooks + hooks._register_hook('final_setup', 'Problem', post=_view_graph, exit=True) + _load_and_exec(options.file[0], user_args)