| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748 |
- def write_frame_png(self, filename="nix-tree.png"):
- """Dump the graph to a png file"""
- try:
- cmap = getattr(matplotlib.cm, self.config["color_map"])
- except AttributeError:
- raise util.TreeCLIError("Colormap {} does not exist".format(
- self.config["color_map"]))
- pos = {n: (n.x, n.y) for n in self.nodes}
- col_scale = 255.0/(self.depth+1.0)
- col = [(x.level+random.random()*self.config["color_scatter"])*col_scale
- for x in self.G.nodes()]
- col = [min([x,255]) for x in col]
- img_y_height=self.config["img_y_height_inches"]
- size_min = self.config["min_node_size"]
- size_max = self.config["max_node_size_over_min_node_size"] * size_min
- plt.figure(1, figsize=(img_y_height*self.config["aspect_ratio"],
- img_y_height))
- node_size = [min(size_min + (x.out_degree-1)*
- self.config["add_size_per_out_link"],
- size_max) if x.level > 0 else size_max for
- x in self.G.nodes()]
- # Draw edges
- nx.draw(self.G, pos, node_size=node_size, arrows=False,
- with_labels=self.config["show_labels"],
- edge_color=self.config["edge_color"],
- font_size=12*self.config["font_scale"],
- node_color=col, vmin=0, vmax=256,
- width=self.config["edge_width_scale"],
- alpha=self.config["edge_alpha"], nodelist=[])
- # Draw nodes
- nx.draw(self.G, pos, node_size=node_size, arrows=False,
- with_labels=self.config["show_labels"],
- font_size=12*self.config["font_scale"],
- node_color=col, vmin=0, vmax=255, edgelist=[],
- font_weight="light", cmap=cmap,
- font_color=self.config["font_color"])
- logger.info("Writing png file: {}".format(filename))
- plt.savefig(filename, dpi=self.config["dpi"])
- plt.close()
|
