Add node graph generation capabilities

pykat/__main__.py

@@ -46,8 +46,10 @@ from .finesse import kat as katparser
               help="Display results as figure.")
 @click.option("--save-figure", type=click.File("wb", lazy=False),
               help="Save image of figure to file.")
+@click.option("--display-graph", is_flag=True, help="Generate and display model node graph.")
 @click.version_option(version=__version__, prog_name="Pykat")
-def cli(file, xstart, xstop, xsteps, xscale, trace, maxtem, ignored_blocks, plot, save_figure):
+def cli(file, xstart, xstop, xsteps, xscale, trace, maxtem, ignored_blocks, plot, save_figure,
+        display_graph):
     """Base CLI command group"""
     kat = katparser()
     kat.load(file.name)
@@ -106,3 +108,8 @@ def cli(file, xstart, xstop, xsteps, xscale, trace, maxtem, ignored_blocks, plot
             click.echo("Cannot plot or save figure without an xaxis defined in FILE.",
                        err=True)
             sys.exit(1)
+
+    if display_graph:
+        from .tools.plotting.graph import NodeGraph
+        nodegraph = NodeGraph(kat)
+        nodegraph.view_pdf()
\ No newline at end of file

pykat/tools/plotting/graph.py

+"""Model node graph visualisation with Graphviz.
+
+Based on node_graph project by Sebastian Steinlechner: https://github.com/sestei/node_graph/
+"""
+
+import tempfile
+from graphviz import Digraph
+from ...components import laser, space
+
+
+class NodeGraph(object):
+    """Pykat node graph plotter."""
+    def __init__(self, kat):
+        self.kat = kat
+
+    def view_pdf(self):
+        """View the graph as a PDF"""
+        return self.node_graph().view(directory=tempfile.gettempdir(), cleanup=True)
+
+    def _repr_svg_(self):
+        """Graphviz rendering for Jupyter notebooks."""
+        return self.node_graph()._repr_svg_()
+
+    def node_graph(self, engine="fdp", node_style="filled", node_font="Helvetica",
+                   node_font_size=10, node_width=0.2, node_shape="point", node_color="red",
+                   node_dump_color="black", edge_arrowhead="dot", graph_splines="compound",
+                   graph_font="Helvetica", graph_font_size=8, component_color="LightSkyBlue",
+                   laser_color="Orange", space_color="MediumSlateBlue",
+                   detector_color="YellowGreen"):
+        """Create Graphviz node graph"""
+        graph = Digraph(engine=engine)
+        graph.attr("node", style=node_style, fontname=node_font, fontsize=str(node_font_size))
+        graph.attr("edge", arrowhead=edge_arrowhead)
+        graph.attr("graph", splines=graph_splines, fontname=graph_font,
+                   fontsize=str(graph_font_size))
+
+        def add_kat_node(from_id, node):
+            node_id = str(node.id)
+            color = node_dump_color if node.name == 'dump' else node_color
+            graph.node(node_id, label='', xlabel=node.name, width=str(node_width),
+                       shape=node_shape, color=color)
+            graph.edge(from_id, node_id)
+
+        for comp in self.kat.components.values():
+            if isinstance(comp, laser):
+                attr = {'fillcolor': laser_color, 'shape': 'box'}
+            elif isinstance(comp, space):
+                attr = {'fillcolor': space_color, 'shape': 'diamond'}
+            else:
+                attr = {'fillcolor': component_color, 'shape': 'box'}
+            graph.node(comp.name, **attr)
+            for node in comp.nodes:
+                add_kat_node(comp.name, node)
+
+        for det in self.kat.detectors.values():
+            if len(det._nodes) > 0:
+                graph.node(det.name, fillcolor=detector_color, shape='ellipse')
+                for node in det._nodes:
+                    add_kat_node(det.name, node)
+
+        return graph
\ No newline at end of file

requirements.txt

@@ -7,3 +7,4 @@ pandas
 matplotlib
 tabulate
 click >= 7.0
+graphviz
\ No newline at end of file