elements: Add exporting connectivity graph to PDF

This allows for exporting the connectivity graph to a PDF, which can
be used in papers, deliverables, ...

2 files changed, 64 insertions, 4 deletions

diff --git a/rumba/elements/experimentation.py b/rumba/elements/experimentation.py
index 2c6f478..60a35e1 100644
--- a/rumba/elements/experimentation.py
+++ b/rumba/elements/experimentation.py
@@ -521,16 +521,75 @@ class Experiment(object):
         end = time.time()
         logger.info("Swap-out took %.2f seconds", end - start)
 
-    def to_dms_yaml(self, filename, clean=True):
+    def to_dms_yaml(self, filename):
+        """
+        Generate a YAML file of the experiment which can be fed to the
+        ARCFIRE DMS.
+
+        :param filename: The output YAML filename.
+        """
         mode = 'w'
-        if not clean:
-            mode += '+'
         with open(filename, mode) as f:
             for node in self.nodes:
                 f.write("---\n")
                 node.to_dms_yaml(f)
             f.write("...\n")
 
+        logger.info('Generated DMS YAML file')
+
+    def export_connectivity_graph(self, filename):
+        """
+        Generate a PDF of the physical connectivity graph.
+
+        :param filename: The output PDF filename.
+        """
+        try:
+            import pydot
+        except:
+            logger.error('Warning: pydot module not installed, '
+                         'cannot produce DIF graphs images')
+            return
+
+        colors = ['red', 'green', 'blue', 'orange', 'yellow']
+        fcolors = ['black', 'black', 'white', 'black', 'black']
+
+        gvizg = pydot.Dot(graph_type='graph')
+
+        i = 0
+        try:
+            for node in self.nodes:
+                g_node = pydot.Node(node.name,
+                                    label=node.name,
+                                    style="filled",
+                                    fillcolor=colors[i],
+                                    fontcolor=fcolors[i])
+
+                gvizg.add_node(g_node)
+                i += 1
+                if i == len(colors):
+                    i = 0
+        except:
+            logger.error('Failed to create pydot Node')
+            return
+
+        try:
+            for dif in self.dif_ordering:
+                if isinstance(dif, topology.ShimEthDIF):
+                    edge = pydot.Edge(dif.members[0].name,
+                                      dif.members[1].name,
+                                      label=dif.name,
+                                      color="black")
+                    gvizg.add_edge(edge)
+        except:
+            logger.error('Failed to create pydot Edge')
+            return
+
+        try:
+            gvizg.write_pdf(filename)
+            logger.info('Generated pdf of connectivity graph')
+        except:
+            logger.error('Failed to write PDF')
+
 class Executor:
     __metaclass__ = abc.ABCMeta
 
diff --git a/setup.py b/setup.py
index 77661a2..6c4fd62 100755
--- a/setup.py
+++ b/setup.py
@@ -30,6 +30,7 @@ setuptools.setup(
         'contextlib2; python_version<"3.0"',
         'enum34; python_version<"3.0"'
     ],
-    extras_require={'NumpyAcceleration': ['numpy']},
+    extras_require={'NumpyAcceleration': ['numpy'],
+                    'graphs' : ['pydot']},
     scripts=['tools/rumba-access']
 )