Commit 27ff9f40 authored by Sebastian Steinlechner's avatar Sebastian Steinlechner
Browse files

Nodegraph plotting function

Added functionality to produce a graph of the node network. Relies
on pygraphviz.
parent d3cb761a
......@@ -8,6 +8,7 @@ __version__ = "1.0.8"
# This flag is used to switch on the gui features in pkat at import time
USE_GUI = False
HAS_OPTIVIS = False
HAS_GRAPHVIZ = False
import imp
......@@ -17,6 +18,13 @@ try:
except ImportError:
HAS_OPTIVIS = False
try:
imp.find_module('pygraphviz')
HAS_GRAPHVIZ = True
except ImportError:
HAS_GRAPHVIZ = False
import pykat.exceptions as pkex
NoGUIException = pkex.BasePyKatException("No PyQt4 module was found so cannot open a GUI")
......
......@@ -91,7 +91,7 @@ import pykat.external.six as six
import pykat.exceptions as pkex
from pykat import USE_GUI, HAS_OPTIVIS, NoGUIException
from pykat import USE_GUI, HAS_OPTIVIS, HAS_GRAPHVIZ, NoGUIException
if HAS_OPTIVIS:
......@@ -2089,6 +2089,57 @@ class kat(object):
return out
def node_graph(self, output):
"""
node_graph(output)
Generate node visualisation of pykat node network using the
pygraphviz package. The generated graphics is saved in output, where the
output format is determined from the file extension. The exact formats
available depend on how graphviz was built, but generally .pdf and .png
should work.
"""
if not HAS_GRAPHVIZ:
print('pygraphviz is not installed.')
return
import pygraphviz as pgv
def add_kat_node(G, from_id, node):
color = 'black' if node.name == 'dump' else 'red'
G.add_node(node.id, label='', xlabel=node.name, width=0.2,
shape='point', fontsize=11, color=color)
G.add_edge(from_id, node.id)
G = pgv.AGraph(strict=False)
G.graph_attr['label'] = 'Created by PyKat'
G.graph_attr['pencolor'] = 'black'
G.node_attr['style'] = 'filled'
G.node_attr['fontname'] = 'helvetica'
# add components and associated nodes
for comp in self.components.itervalues():
if isinstance(comp, pykat.components.laser):
attr = {'fillcolor': 'Orange', 'shape': 'box'}
elif isinstance(comp, pykat.components.space):
attr = {'fillcolor': 'MediumSlateBlue', 'shape': 'diamond'}
else:
attr = {'fillcolor': 'LightSkyBlue', 'shape': 'box'}
G.add_node(comp.name, **attr)
for node in comp.nodes:
add_kat_node(G, comp.name, node)
# add detectors and associated nodes
for det in self.detectors.itervalues():
# slightly ugly, but unfortunately detectors don't have the .nodes
# property like components do.
if len(det._nodes) > 0:
G.add_node(det.name, fillcolor='YellowGreen', shape='ellipse')
for node in det._nodes:
add_kat_node(G, det.name, node)
G.draw(output, prog='neato')
def optivis(self):
if not HAS_OPTIVIS:
print("Optivis is not installed")
......
Supports Markdown
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