Maintenance will be performed on git.ligo.org, chat.ligo.org, and docs.ligo.org, starting at approximately 10am CDT Tuesday 20 August 2019. The maintenance is expected to take around an hour and here will be two short periods of downtime, one at the beginning of the maintenance and another at the end.

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