Statistics
| Branch: | Revision:

root / fiddle / heuristic-betweenness-centrality / utility.py @ 739fe075

History | View | Annotate | Download (1.96 KB)

1
import os
2
import networkx as nx
3
import matplotlib.pyplot as plt
4
import betweenness_centrality as centrality
5

    
6
MAIN_CODE_DIR = os.environ.get('MAIN_CODE_DIR', '')
7

    
8
def initialize():
9
    """
10
    Create all necessary folder
11
    """
12
    output_dir = os.path.join(MAIN_CODE_DIR, 'output')
13
    if not os.path.isdir(output_dir):
14
        os.makedirs(output_dir)
15

    
16
    visualization_dir = os.path.join(output_dir, 'visualization')
17
    if not os.path.isdir(visualization_dir):
18
        os.makedirs(visualization_dir)
19

    
20
def weight_betweenness_centrality(filepath, file_suffix):
21
    graph = nx.read_weighted_edgelist(filepath)
22

    
23
    score = centrality.weight_betweenness_centrality(graph, rescale=True)
24
    filepath = '/output/weight_basic_%s.csv'  % file_suffix
25

    
26
    with open(MAIN_CODE_DIR + filepath, 'w') as output:
27
        for node, bc in score.iteritems():
28
            output.write('%s, %s\n' % (node, bc))
29

    
30
def brandes_betweenness_centrality(filepath, file_suffix):
31
    graph = nx.read_weighted_edgelist(filepath)
32

    
33
    score = nx.betweenness_centrality(graph)
34
    filepath = '/output/networkx_%s.csv'  % file_suffix
35

    
36
    with open(MAIN_CODE_DIR + filepath, 'w') as output:
37
        for node, bc in score.iteritems():
38
            output.write('%s, %s\n' % (node, bc))
39

    
40
def brandes_betweenness_centrality_endpoints(filepath, file_suffix):
41
    graph = nx.read_weighted_edgelist(filepath)
42

    
43
    score = nx.betweenness_centrality(graph, endpoints=True)
44
    filepath = '/output/heuristic_%s.csv'  % file_suffix
45

    
46
    with open(MAIN_CODE_DIR + filepath, 'w') as output:
47
        for node, bc in score.iteritems():
48
            output.write('%s, %s\n' % (node, bc))
49

    
50

    
51
def plot_graph(output_filepath, graph, title=""):
52
    plt.figure(1, figsize=(15, 10))
53
    plt.title(title)
54
    pos=nx.spring_layout(graph)
55
    colors=range(20)
56
    nx.draw(graph, pos, node_size=300, with_label=True)
57
    nx.draw_networkx_labels(graph,pos,font_size=10,font_family='sans-serif')
58
    plt.savefig(output_filepath)
59
    plt.close()
60
    # plt.show()