src/nepi/util/plotter.py

   1 #
   2 #    NEPI, a framework to manage network experiments
   3 #    Copyright (C) 2013 INRIA
   4 #
   5 #    This program is free software: you can redistribute it and/or modify
   6 #    it under the terms of the GNU General Public License as published by
   7 #    the Free Software Foundation, either version 3 of the License, or
   8 #    (at your option) any later version.
   9 #
  10 #    This program is distributed in the hope that it will be useful,
  11 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
  12 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13 #    GNU General Public License for more details.
  14 #
  15 #    You should have received a copy of the GNU General Public License
  16 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17 #
  18 # Author: Alina Quereilhac <alina.quereilhac@inria.fr>
  19
  20 import networkx
  21 import os
  22
  23 class PFormats:
  24     DOT = "dot"
  25     FIGURE = "figure"
  26
  27 class ECPlotter(object):
  28     def plot(self, ec, fpath = None, format= PFormats.FIGURE, persist = False):
  29         graph, labels = self._ec2graph(ec)
  30
  31         add_extension = False
  32
  33         if persist and not fpath:
  34             import tempfile
  35             dirpath = tempfile.mkdtemp()
  36             fpath = os.path.join(dirpath, "%s_%s" % (ec.exp_id, ec.run_id))
  37             add_extension = True
  38
  39         if format == PFormats.FIGURE:
  40             import matplotlib.pyplot as plt
  41             pos = networkx.graphviz_layout(graph, prog="neato")
  42             networkx.draw(graph, pos = pos, node_color="white",
  43                     node_size = 500, with_labels=True)
  44
  45             label = "\n".join(map(lambda v: "%s: %s" % (v[0], v[1]), labels.iteritems()))
  46             plt.annotate(label, xy=(0.0, 0.95), xycoords='axes fraction')
  47
  48             if persist:
  49                 if add_extension:
  50                     fpath += ".png"
  51
  52                 plt.savefig(fpath, bbox_inches="tight")
  53             else:
  54                 plt.show()
  55
  56         elif format == PFormats.DOT:
  57             if persist:
  58                 if add_extension:
  59                     fpath += ".dot"
  60
  61                 networkx.write_dot(graph, fpath)
  62             else:
  63                 import subprocess
  64                 subprocess.call(["dot", "-Tps", fpath, "-o", "%s.ps" % fpath])
  65                 subprocess.call(["evince","%s.ps" % fpath])
  66
  67         return fpath
  68
  69     def _ec2graph(self, ec):
  70         graph = networkx.Graph(graph = dict(overlap = "false"))
  71
  72         labels = dict()
  73         connections = set()
  74
  75         for guid, rm in ec._resources.iteritems():
  76             label = rm.get_rtype()
  77
  78             graph.add_node(guid,
  79                 label = "%d %s" % (guid, label),
  80                 width = 50/72.0, # 1 inch = 72 points
  81                 height = 50/72.0,
  82                 shape = "circle")
  83
  84             labels[guid] = label
  85
  86             for guid2 in rm.connections:
  87                 # Avoid adding a same connection twice
  88                 if (guid2, guid) not in connections:
  89                     connections.add((guid, guid2))
  90
  91         for (guid1, guid2) in connections:
  92             graph.add_edge(guid1, guid2)
  93
  94         return graph, labels