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" ...@@ -8,6 +8,7 @@ __version__ = "1.0.8"
# This flag is used to switch on the gui features in pkat at import time # This flag is used to switch on the gui features in pkat at import time
USE_GUI = False USE_GUI = False
HAS_OPTIVIS = False HAS_OPTIVIS = False
HAS_GRAPHVIZ = False
import imp import imp
...@@ -17,6 +18,13 @@ try: ...@@ -17,6 +18,13 @@ try:
except ImportError: except ImportError:
HAS_OPTIVIS = False HAS_OPTIVIS = False
try:
imp.find_module('pygraphviz')
HAS_GRAPHVIZ = True
except ImportError:
HAS_GRAPHVIZ = False
import pykat.exceptions as pkex import pykat.exceptions as pkex
NoGUIException = pkex.BasePyKatException("No PyQt4 module was found so cannot open a GUI") NoGUIException = pkex.BasePyKatException("No PyQt4 module was found so cannot open a GUI")
......
...@@ -91,7 +91,7 @@ import pykat.external.six as six ...@@ -91,7 +91,7 @@ import pykat.external.six as six
import pykat.exceptions as pkex 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: if HAS_OPTIVIS:
...@@ -2089,6 +2089,57 @@ class kat(object): ...@@ -2089,6 +2089,57 @@ class kat(object):
return out 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): def optivis(self):
if not HAS_OPTIVIS: if not HAS_OPTIVIS:
print("Optivis is not installed") 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