Source

privateer / privateer_pyevolve / plot.py

Full commit
import matplotlib.pyplot as plt
#from random import choice, shuffle
from pylab import clf

def draw_graph(self):

  #carrying_edges =[(u,v) for (u,v,d) in self.dag.edges(data=True) if d['carrying_info'] == True and d['sender'] == True]
  #noncarrying_edges =[(u,v) for (u,v,d) in self.dag.edges(data=True) if d['carrying_info'] == False and d['sender'] == True]
  #nonsending =[(u,v) for (u,v,d) in self.dag.edges(data=True) if d['sender'] == False]

  pos = nx.shell_layout(self.dag)

  nx.draw_networkx_nodes(self.dag,pos,node_size=150, nodelist=self.others, node_color='aquamarine')
  nx.draw_networkx_nodes(self.dag,pos,node_size=300, nodelist=[self.focal], node_color='red')
  nx.draw_networkx_nodes(self.dag,pos,node_size=300, nodelist=self.focal_neighbors, node_color='orange')

  sender_nodes = [node for node in self.dag.nodes() if self.dag.node[node]['sender'] == True]
  nx.draw_networkx_nodes(self.dag,pos,node_size=300, nodelist=sender_nodes, node_color='purple')

  true_gossip_nodes = []
  false_gossip_nodes = []
  for node in self.dag.nodes():
    try:
      if self.dag.node[node]['gossip'] == True:
        true_gossip_nodes.append(node)
      elif self.dag.node[node]['gossip'] == False:
        false_gossip_nodes.append(node)
    except KeyError:
      pass

  nx.draw_networkx_nodes(self.dag,pos,node_size=300, 
            nodelist=true_gossip_nodes, node_color='white',
            node_shape='d')
  nx.draw_networkx_nodes(self.dag,pos,node_size=300, 
            nodelist=false_gossip_nodes, node_color='grey',
            node_shape='d')

  nx.draw_networkx_edges(self.dag,pos,edgelist=self.dag.edges(),
                      width=1, edge_color='b', alpha=0.5, arrows=True)
  nx.draw_networkx_edges(self.dag,pos,edgelist=self.removed_edges,
                      width=1, edge_color='cyan', alpha=0.5, arrows=False, style='dotted')
  nx.draw_networkx_edges(self.dag,pos,edgelist=self.observer_edges,
                      width=1, edge_color='red', alpha=0.5, arrows=False, style='dotted')
  #nx.draw_networkx_edges(self.dag,pos,edgelist=carrying_edges,
  #                    width=1, edge_color='g', alpha=0.5, arrows=True)
  #nx.draw_networkx_edges(self.dag,pos,edgelist=noncarrying_edges,
  #                    width=1,alpha=0.5,edge_color='b',style='dashed', 
  #                    arrows=True)

  #nx.draw_networkx_edges(self.dag,pos,edgelist=nonsending,
  #                    width=3,alpha=0.5,edge_color='r',style='dotted',
  #                    arrows=True)

  nx.draw_networkx_labels(self.dag,pos,font_size=14,font_family='sans-serif')

  plt.axis('off')

def record_graph_intermediates(self, outfile, iter):
  clf()
  self.draw_graph()
  plt.savefig("%s-%03d.png" % (outfile, iter))

def animate_graph(self, movie, pngfiles='animation', fps=2):

  try:
    os.remove(movie)
  except OSError:
    pass

  command = 'ffmpeg -i %s-' % pngfiles + r'%03d.png ' + '-r %s ' % fps + movie
  print command
  os.system(command)