sim2.py

import networkx as nx
import matplotlib.pyplot as plt
from random import random,randint

N = 1000
K = 5
TMAX = 2

def make_small_world_network(n,k,p):
    assert p >= 0 and p <=1, "p is reconnect probability"
    assert n > 0, "n is num nodes"
    assert k > 0, "2*k is num of link per one node"

    G = nx.Graph()
    G.add_nodes_from(range(1,n+1))

    for i in range(1,n+1):
        for j in range(1,k+1):
            v = i
            u = i+j 
            if u > n:
                u = u % (n+1) + 1
            G.add_edge(v,u)

    for i in range(1,n+1):
        for j in range(1,k+1):
            if random() < p:
                # reconnect edge
                v = i
                old_u = i+j
                if old_u > n:
                    old_u = old_u % (n+1) + 1
                new_u = randint(1,n)
                while v != new_u and G.has_edge(v,new_u):
                    new_u = randint(1,n)
                G.remove_edge(v,old_u)
                G.add_edge(v,new_u)
    return G

def sim(G,r):
    time_count = 0
    infected_count = 1
    while True:
        time_count += 1

        if time_count == 1:
            pos = nx.spring_layout(G)

        plt.close()
        plt.axis('off')
        node_color = []
        for n in G.nodes_iter():
            if 'dead' in G.node[n]:
                node_color.append("r")
            elif 'infected' in G.node[n]:
                node_color.append("r")
            else:
                node_color.append("b")

        nx.draw_networkx(G,pos=pos,node_color=node_color,node_size=10,width=0.1,with_labels=False)
        plt.savefig("fig/{0}.png".format(time_count))

        remove_candidate = []
        for n in G.nodes_iter():
            if 'infected' in G.node[n]:
                G.node[n]['infected'] += 1
                if G.node[n]['infected'] >= TMAX:
                    # dead
                    remove_candidate.append(n)
        # remove dead node
        for i in remove_candidate:
            G.node[i]['dead'] = True

        if infected_count == N:
            # all node are infected and died
            return (time_count,infected_count)

        infected_nodes = 0
        infected_candidate = set()
        for n in G.nodes_iter():
            if 'infected' in G.node[n]:
                infected_nodes += 1
                for u,v in G.edges_iter(n):
                    if random() < r:
                        if not 'infected' in G.node[v] and not 'dead' in G.node[v]:
                            infected_candidate.add(v)
        for i in infected_candidate:
            G.node[i]['infected'] = 0
            infected_count += 1


def main():
    G = make_small_world_network(N,K,0.05)
    l = nx.average_shortest_path_length(G)
    c = nx.average_clustering(G)
    G.node[1]['infected'] = 0
    time_count,infected_count = sim(G,1)
    print("average shortest path length: {0}".format(l))
    print("average clustering coefficient: {0}".format(c))


if __name__ == "__main__":
    main()