Nodegraph plotting function

Added functionality to produce a graph of the node network. Relies
on pygraphviz.

pykat/__init__.py

@@ -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")

pykat/finesse.py

@@ -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")