diff --git a/muscle3/model_graph_pydot.py b/muscle3/model_graph_pydot.py new file mode 100644 index 00000000..398ed37e --- /dev/null +++ b/muscle3/model_graph_pydot.py @@ -0,0 +1,242 @@ +from typing import List, Union + +from ymmsl.configuration import PartialConfiguration +from ymmsl.identity import Reference +from ymmsl.component import Operator, Component + +import pydot + +COLORS = { + Operator.F_INIT: "#2998ba", + Operator.O_I: "#eddea1", + Operator.S: "#f1c40f", + Operator.O_F: "#e67e22", +} + + +def port_operator(port: Reference, component: Union[Component, None]): + """Look up the operator corresponding to a specific port""" + return component.ports.operator(port) if component else "normal" + + +def port_shape(operator: str, simple: bool=True): + """Given a port reference, find the component referred to, + look up the port type matching this name + and look up the shape corresponding to the port type.""" + + if simple: + if operator == Operator.F_INIT or operator == Operator.S: + return "normal" + return "none" + else: + # I think it is quite easy to misinterpret the direction of a MMSL diagram + # given that the 'weight' of the edge is towards the filled shape. + # I would consider making the sending_port smaller, but cannot find the setting + # in graphviz + # https://www.graphviz.org/docs/attr-types/arrowType/ + OPERATOR_SHAPES = { + Operator.NONE: "none", + Operator.F_INIT: "odiamond", + Operator.O_I: "dot", + Operator.S: "odot", + Operator.O_F: "diamond", + } + return OPERATOR_SHAPES[operator] + + +def find_component(name: Reference, components: List[Component]): + """Find a component by reference""" + return next( + (component for component in components if component.name == str(name)), None + ) + + +def set_style(graph, draw_ports: bool=False): + """set default properties to make for a more readable DOT file""" + graph.add_node( + pydot.Node( + "node", + shape="plain" if draw_ports else 'box', + style="rounded", + fixedsize="false", + width=2, + penwidth=2, + height=1, + labelloc="c", + ) + ) + + # set default edge properties + graph.add_node( + pydot.Node( + "edge", + dir="both", + labelfontsize=10, + fontsize=10, + penwidth=2, + len=2, + ) + ) + + +def trim_sending_port(identifier: str): + """Strip _out suffix from identifier""" + return identifier[:-4] if identifier.endswith("_out") else identifier + + +def trim_receiving_port(identifier: str): + """Strip _in, _init suffix from identifier""" + if identifier.endswith("_in"): + return identifier[:-3] + if identifier.endswith("_init"): + return identifier[:-5] + return identifier + + +def headport(identifier: Reference, component: Union[Component, None]): + """Given a reference, return the portPos. + https://www.graphviz.org/docs/attr-types/portPos/ + """ + return str(identifier) + + +def tailport(identifier: Reference, component: Union[Component, None]): + """Given a reference, return the portPos. + https://www.graphviz.org/docs/attr-types/portPos/ + """ + return str(identifier) + + +def port_shortname(identifier: Reference): + """Strip suffixes and summarize names to only a few characters""" + identifier = trim_sending_port(trim_receiving_port(str(identifier))) + return "".join([s[0] for s in identifier.split("_")]).upper() + + +def legend_html_label(): + return f"""< + + + + + + + +
F_INITO_IlegendSO_F
>""" + pass + + +def component_html_label(component: Component): + """Construct a HTML-like label (https://graphviz.org/doc/info/shapes.html#html)""" + + label = "<\n \n" + + for port in component.ports.f_init: + label += f" \n" + if len(component.ports.f_init) == 0: + label += f" \n" + for port in component.ports.o_i: + label += f" \n" + if len(component.ports.o_i) == 0: + label += f" \n" + + label += f" \n" + + for port in component.ports.s: + label += f" \n" + if len(component.ports.s) == 0: + label += f" \n" + for port in component.ports.o_f: + label += f" \n" + if len(component.ports.o_f) == 0: + label += f" \n" + + return label.replace("'", '"') + "
{port_shortname(port)}{port_shortname(port)}{component.name}{port_shortname(port)}{port_shortname(port)}
>" + + +def plot_model_graph(config: PartialConfiguration, simplify_edge_labels: bool=True, draw_ports: bool=False, simple_edges: bool=True, show_legend: bool=True) -> None: + """Convert a PartialConfiguration into DOT format.""" + graph = pydot.Dot( + config.model.name, + graph_type="digraph", + layout="dot", + pad=1, + # splines="ortho", + nodesep=0.6, + ranksep=0.75, + fontname="Sans-Serif", + ) + # be very careful with ortho splines, I have seen it put edges + # upside down and eating labels + set_style(graph, draw_ports=draw_ports) + + if draw_ports: + label_method = component_html_label + else: + label_method = lambda x: x.name + + # Start with a legend node + if draw_ports and show_legend: + graph.add_node(pydot.Node("legend", label=legend_html_label())) + + for component in config.model.components: + graph.add_node( + pydot.Node(str(component.name), label=label_method(component)) + ) + + # assume that conduits are ordered by port + last_port = None + for conduit in config.model.conduits: + # emit an edge[sametail=] config for this port name if it is changed from the previous one + if str(conduit.sending_port()) != last_port: + graph.add_node(pydot.Node('edge', sametail=str(conduit.sending_port()))) + last_port = str(conduit.sending_port()) + + # can we do this more elegantly? + sender = find_component(conduit.sending_component(), config.model.components) + receiver = find_component( + conduit.receiving_component(), config.model.components + ) + + # Due to yMMSL conventions we cannot have edges with the same head (since + # an input port can only have one conduit connected to it) + port_config = { + 'tailport': tailport(conduit.sending_port(), sender), + 'headport': headport(conduit.receiving_port(), receiver) + } + + + edge = pydot.Edge( + str(conduit.sending_component()), + str(conduit.receiving_component()), + arrowtail=port_shape( + port_operator( + conduit.sending_port(), + sender, + ), + simple_edges + ), + arrowhead=port_shape( + port_operator( + conduit.receiving_port(), + receiver, + ), + simple_edges + ), + **port_config + ) + + # if port names match exactly (optionally when removing an _in or _out suffix) + # we show the name on the conduit instead of on the port + if simplify_edge_labels and trim_sending_port( + str(conduit.sending_port()) + ) == trim_receiving_port(str(conduit.receiving_port())): + edge.set_label(trim_sending_port(str(conduit.sending_port()))) + else: + edge.set_taillabel(str(conduit.sending_port())) + edge.set_headlabel(str(conduit.receiving_port())) + + # we could consider setting a minlen based on the text length and font size + graph.add_edge(edge) + + return graph diff --git a/muscle3/muscle3.py b/muscle3/muscle3.py index f715d5fd..80d41ac0 100644 --- a/muscle3/muscle3.py +++ b/muscle3/muscle3.py @@ -1,17 +1,21 @@ import sys from collections import OrderedDict from pathlib import Path -from typing import Sequence +from typing import Sequence, Union import click import ymmsl from ymmsl import PartialConfiguration +import pydot +import subprocess -from libmuscle.planner.planner import ( - Planner, Resources, InsufficientResourcesAvailable) + +from libmuscle.planner.planner import Planner, Resources, InsufficientResourcesAvailable from libmuscle.snapshot_manager import SnapshotManager +from .model_graph_pydot import plot_model_graph + _RESOURCES_INCOMPLETE_MODEL = """ A model, implementations and resources must be given to be able to calculate @@ -32,20 +36,30 @@ def muscle3() -> None: pass -@muscle3.command(short_help='Calculate resources needed for a simulation') +@muscle3.command(short_help="Calculate resources needed for a simulation") @click.argument( - 'ymmsl_files', nargs=-1, required=True, type=click.Path( - exists=True, file_okay=True, dir_okay=False, readable=True, - allow_dash=True, resolve_path=True)) + "ymmsl_files", + nargs=-1, + required=True, + type=click.Path( + exists=True, + file_okay=True, + dir_okay=False, + readable=True, + allow_dash=True, + resolve_path=True, + ), +) @click.option( - '-c', '--cores-per-node', nargs=1, type=int, required=True, - help='Set number of cores per cluster node.') -@click.option( - '-v', '--verbose', is_flag=True, help='Show instance allocations.') -def resources( - ymmsl_files: Sequence[str], - cores_per_node: int, verbose: bool - ) -> None: + "-c", + "--cores-per-node", + nargs=1, + type=int, + required=True, + help="Set number of cores per cluster node.", +) +@click.option("-v", "--verbose", is_flag=True, help="Show instance allocations.") +def resources(ymmsl_files: Sequence[str], cores_per_node: int, verbose: bool) -> None: """Calculate the number of nodes needed to run the simulation. In order to run a MUSCLE3 simulation on a cluster, a batch job has @@ -81,7 +95,7 @@ def resources( click.echo(_RESOURCES_INCOMPLETE_MODEL, err=True) sys.exit(1) - resources = Resources({'node000001': set(range(cores_per_node))}) + resources = Resources({"node000001": set(range(cores_per_node))}) planner = Planner(resources) try: allocations = planner.allocate_all(config, True) @@ -94,33 +108,42 @@ def resources( if verbose: click.echo() if num_nodes == 1: - click.echo('A total of 1 node will be needed, as follows:') + click.echo("A total of 1 node will be needed, as follows:") else: - click.echo( - f'A total of {num_nodes} nodes will be needed,' - ' as follows:') + click.echo(f"A total of {num_nodes} nodes will be needed," " as follows:") click.echo() for instance in sorted(allocations): - click.echo(f'{instance}: {str(allocations[instance])}') + click.echo(f"{instance}: {str(allocations[instance])}") else: - click.echo(f'{num_nodes}', nl=False) + click.echo(f"{num_nodes}", nl=False) sys.exit(0) -@muscle3.command(short_help='Display details of a stored snapshot') +@muscle3.command(short_help="Display details of a stored snapshot") @click.argument( - 'snapshot_files', nargs=-1, required=True, type=click.Path( - exists=True, file_okay=True, dir_okay=False, readable=True, - allow_dash=True, resolve_path=True, path_type=Path)) -@click.option( - '-d', '--data', is_flag=True, - help='Display stored data. Note this may result in a lot of output!') + "snapshot_files", + nargs=-1, + required=True, + type=click.Path( + exists=True, + file_okay=True, + dir_okay=False, + readable=True, + allow_dash=True, + resolve_path=True, + path_type=Path, + ), +) @click.option( - '-v', '--verbose', is_flag=True, help='Display more metadata.') -def snapshot( - snapshot_files: Sequence[Path], data: bool, verbose: bool) -> None: + "-d", + "--data", + is_flag=True, + help="Display stored data. Note this may result in a lot of output!", +) +@click.option("-v", "--verbose", is_flag=True, help="Display more metadata.") +def snapshot(snapshot_files: Sequence[Path], data: bool, verbose: bool) -> None: """Display information about stored snapshots. Per provided snapshot, display metadata. Stored data can also be output by @@ -129,24 +152,30 @@ def snapshot( """ for file in snapshot_files: snapshot = SnapshotManager.load_snapshot_from_file(file) - click.echo(f'Snapshot at {file}:') - typ = 'Final' if snapshot.is_final_snapshot else 'Intermediate' - properties = OrderedDict([ - ('Snapshot type', typ), - ('Snapshot timestamp', - snapshot.message.timestamp if snapshot.message else float('-inf')), - ('Snapshot wallclock time', snapshot.wallclock_time), - ('Snapshot triggers', snapshot.triggers), - ]) + click.echo(f"Snapshot at {file}:") + typ = "Final" if snapshot.is_final_snapshot else "Intermediate" + properties = OrderedDict( + [ + ("Snapshot type", typ), + ( + "Snapshot timestamp", + snapshot.message.timestamp if snapshot.message else float("-inf"), + ), + ("Snapshot wallclock time", snapshot.wallclock_time), + ("Snapshot triggers", snapshot.triggers), + ] + ) if verbose: - properties.update([ - ('Internal: Port message counts', snapshot.port_message_counts), - ]) + properties.update( + [ + ("Internal: Port message counts", snapshot.port_message_counts), + ] + ) for prop_name, prop_value in properties.items(): - click.secho(f'{prop_name}: ', nl=False, bold=True) + click.secho(f"{prop_name}: ", nl=False, bold=True) click.echo(prop_value) if data: - click.secho('Snapshot data:', bold=True) + click.secho("Snapshot data:", bold=True) if snapshot.message is not None: click.echo(snapshot.message.data) else: @@ -154,14 +183,116 @@ def snapshot( click.echo() +@muscle3.command(short_help="Print a graphical representation of this workflow") +@click.argument( + "ymmsl_files", + nargs=-1, + required=True, + type=click.Path( + exists=True, + file_okay=True, + dir_okay=False, + readable=True, + allow_dash=True, + resolve_path=True, + ), +) +@click.option( + "-o", + "--out", + nargs=1, + required=False, + type=click.Path( + exists=False, + file_okay=True, + dir_okay=False, + readable=False, + allow_dash=True, + resolve_path=True, + ), + help="Output file (default ./output.format)", +) +@click.option( + "-f", + "--fmt", + type=click.Choice(pydot.Dot().formats, case_sensitive=False), + required=False, + help="Set output format (default svg).", +) +@click.option( + "-w", + "--viewer", + nargs=1, + required=False, + type=str, + help="Open with specified viewer (try xdg-open)", +) +@click.option("-v", "--verbose", is_flag=True, help="Show more info (prints the generated dot syntax)") +@click.option("-p", "--ports", is_flag=True, help="Explicitly draw component ports.") +@click.option("-l", "--legend", is_flag=True, help="Show a legend (only with --ports).") +@click.option("-s", "--simple-edges", is_flag=True, help="Only indicate conduit direction, not port types.") +@click.option("--portlabels", is_flag=True, help="Never simplify matching port labels along an edge.") +def graph( + ymmsl_files: Sequence[str], + out: Union[Path, None], + fmt: str, + viewer: str, + verbose: bool, + ports: bool, + legend: bool, + simple_edges: bool, + portlabels: bool +) -> None: + """Plot a graphical representation of the passed yMMSL files. + + To help develop or understand about a coupled simulation it may + be useful to view a graphical representation. + + If multiple yMMSL files are given, then they will be combined left + to right, i.e. if there are conflicting declarations, the one from + the file last given is used. + + Result: + + A graph is displayed or saved to disk, containing all the defined + components and the connections between them. + + Examples: + + muscle3 graph simulation.ymmsl + + """ + partial_config = _load_ymmsl_files(ymmsl_files) + + graph = plot_model_graph(partial_config, simplify_edge_labels=not portlabels, draw_ports=ports, simple_edges=simple_edges, show_legend=legend) + + if fmt is None: + fmt = "svg" + if out is None: + out = "output." + fmt + + if out == '-' and fmt == 'dot': + print(graph) + elif out == '-': + print(graph.create(format=fmt)) + else: + if verbose: + print(graph) + + graph.write(out, format=fmt) + + if viewer is not None and out != '-': + subprocess.run([viewer, out]) + + def _load_ymmsl_files(ymmsl_files: Sequence[str]) -> PartialConfiguration: """Loads and merges yMMSL files.""" configuration = PartialConfiguration() for path in ymmsl_files: - with open(path, 'r') as f: + with open(path, "r") as f: configuration.update(ymmsl.load(f)) return configuration -if __name__ == '__main__': +if __name__ == "__main__": muscle3() diff --git a/setup.py b/setup.py index 50cd2ab0..db73ea25 100644 --- a/setup.py +++ b/setup.py @@ -65,6 +65,9 @@ 'sphinx_rtd_theme', 'sphinx-fortran', 'tox' - ] + ], + 'ui': [ + 'pydot' + ], }, )