Commit 01c4d4fc authored by Sean Leavey's avatar Sean Leavey

Add node graph generation capabilities

parent 291303cb
......@@ -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
"""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
......@@ -7,3 +7,4 @@ pandas
matplotlib
tabulate
click >= 7.0
graphviz
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment