diff options
| -rw-r--r-- | rumba/elements/experimentation.py | 65 | ||||
| -rwxr-xr-x | setup.py | 3 |
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 @@ -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'] ) |