bgneal@32: # -*- coding: utf-8 -*-
bgneal@32: """Chapter 5.5, exercise 6 in Allen Downey's Think Complexity book.
bgneal@32: 
bgneal@32: 1. Read Barabási and Albert’s paper and implement their algorithm for generating
bgneal@32:    graphs. See if you can replicate their Figure 2(A), which shows P(k) versus
bgneal@32:    k for a graph with 150 000 vertices.
bgneal@32: 
bgneal@32: """
bgneal@32: import random
bgneal@32: import sys
bgneal@32: 
bgneal@32: from matplotlib import pyplot
bgneal@32: 
bgneal@32: from Graph import Graph, Vertex, Edge
bgneal@32: 
bgneal@32: 
bgneal@32: class BAGraph(Graph):
bgneal@32:     """BAGraph implements the algorithm described in the Barabási and
bgneal@32:     Albert paper "Emergence of Scaling in Random Networks" from Science
bgneal@32:     Magazine Vol. 286, 15 October 1999.
bgneal@32: 
bgneal@32:     """
bgneal@32:     def __init__(self, m0, m):
bgneal@32:         """Create a graph with m0 vertices and m connecting edges.
bgneal@32: 
bgneal@32:         """
bgneal@32:         self.m0 = m0
bgneal@32:         self.m = m
bgneal@32:         self.histogram = []
bgneal@32: 
bgneal@32:         initial_vertices = [Vertex() for i in xrange(m0)]
bgneal@32: 
bgneal@32:         super(BAGraph, self).__init__(vs=initial_vertices)
bgneal@32: 
bgneal@32:         # Add initial edges between nodes randomly
bgneal@32:         n = 0
bgneal@32:         while n != m:
bgneal@32:             # pick two vertices at random
bgneal@32:             v = random.choice(initial_vertices)
bgneal@32:             w = random.choice(initial_vertices)
bgneal@32: 
bgneal@32:             # if they aren't the same and don't already have an edge, add one
bgneal@32:             if v is not w and not self.get_edge(v, w):
bgneal@32:                 self.add_edge(Edge(v, w))
bgneal@32:                 n += 1
bgneal@32: 
bgneal@32:     def add_edge(self, edge):
bgneal@32:         """Our version of add_edge() adds the two vertices on either end of the
bgneal@32:         edge to a histogram. This allows us to more easily pick popular vertices
bgneal@32:         when adding edges as part of the step() method.
bgneal@32: 
bgneal@32:         """
bgneal@32:         super(BAGraph, self).add_edge(edge)     # invoke base class version
bgneal@32: 
bgneal@32:         v, w = edge
bgneal@32:         self.histogram.append(v)
bgneal@32:         self.histogram.append(w)
bgneal@32: 
bgneal@32:     def step(self):
bgneal@32:         """This method adds a new vertex to the graph, then adds m edges that
bgneal@32:         link the new vertex to m different vertices already present in the
bgneal@32:         graph. Preferential treatment is given towards vertices that already
bgneal@32:         have high connectivity.
bgneal@32: 
bgneal@32:         The paper describes this preferential choosing as creating a probability
bgneal@32:         P that a new vertex will be connected to vertex i depends on the
bgneal@32:         connectivity ki of that vertex, so that P(ki) = ki / sum(kj).
bgneal@32: 
bgneal@32:         We implement that by keeping a list of vertices (histogram), where we
bgneal@32:         insert a vertex into this list whenever it gets an edge added to it. We
bgneal@32:         then randomly choose a vertex from this list. Thus the vertices with
bgneal@32:         more edges will be more likely chosen.
bgneal@32: 
bgneal@32:         """
bgneal@32:         # pick m unique vertices to attach to the new vertex
bgneal@32:         vs = set()
bgneal@32:         while len(vs) < self.m:
bgneal@32:             w = random.choice(self.histogram)
bgneal@32:             if w not in vs:
bgneal@32:                 vs.add(w)
bgneal@32: 
bgneal@32:         # Add the new vertex to the graph and create edges
bgneal@32:         v = Vertex()
bgneal@32:         self.add_vertex(v)
bgneal@32:         for w in vs:
bgneal@32:             self.add_edge(Edge(v, w))
bgneal@32: 
bgneal@32: 
bgneal@32: def main(script, m0, m, n):
bgneal@32: 
bgneal@32:     # create a BAGraph with parameters m0 & m:
bgneal@32:     m0, m, n = int(m0), int(m), int(n)
bgneal@32:     g = BAGraph(m0, m)
bgneal@32: 
bgneal@32:     # step the graph n times:
bgneal@32:     for i in xrange(n):
bgneal@32:         g.step()
bgneal@32: 
bgneal@32:     # retrieve probabilities
bgneal@32:     p = g.get_p()
bgneal@32: 
bgneal@32:     # plot P(k) versus k on a log-log scale
bgneal@32: 
bgneal@32:     vals = p.items()
bgneal@32:     vals.sort(key=lambda t: t[0])
bgneal@32:     x, y = zip(*vals)
bgneal@32: 
bgneal@34:     assert abs(sum(y) - 1.0) < 1e-6
bgneal@34: 
bgneal@32:     pyplot.clf()
bgneal@32:     pyplot.xscale('log')
bgneal@32:     pyplot.yscale('log')
bgneal@32:     pyplot.title('P(k) versus k')
bgneal@32:     pyplot.xlabel('k')
bgneal@32:     pyplot.ylabel('P(k)')
bgneal@32:     pyplot.plot(x, y, label='P(k) vs. k', color='green', linewidth=3)
bgneal@32:     pyplot.legend(loc='upper right')
bgneal@32:     pyplot.show()
bgneal@32: 
bgneal@32: 
bgneal@32: if __name__ == '__main__':
bgneal@32:     main(*sys.argv)