Statistics
| Branch: | Revision:

mobicen / util / nx2gt.py @ c16015e2

History | View | Annotate | Download (5.47 KB)

1 1ef4948a LoreBz
2
import networkx as nx
3
import graph_tool as gt
4
import code  # code.interact(local=dict(globals(), **locals()))
5
6
7
8
def get_prop_type(value, key=None):
9
    """
10
    Performs typing and value conversion for the graph_tool PropertyMap class.
11
    If a key is provided, it also ensures the key is in a format that can be
12
    used with the PropertyMap. Returns a tuple, (type name, value, key)
13
    """
14
    if isinstance(key, unicode):
15
        # Encode the key as ASCII
16
        key = key.encode('ascii', errors='replace')
17
18
    # Deal with the value
19
    if isinstance(value, bool):
20
        tname = 'bool'
21
22
    elif isinstance(value, int):
23
        tname = 'float'
24
        value = float(value)
25
26
    elif isinstance(value, float):
27
        tname = 'float'
28
29
    elif isinstance(value, unicode):
30
        tname = 'string'
31
        value = value.encode('ascii', errors='replace')
32
33
    elif isinstance(value, dict):
34
        tname = 'object'
35
36
    else:
37
        tname = 'string'
38
        value = str(value)
39
40
    return tname, value, key
41
42
43
def nx2gt(nxG):
44
    """
45
    Converts a networkx graph to a graph-tool graph.
46
    """
47
    # Phase 0: Create a directed or undirected graph-tool Graph
48
    gtG = gt.Graph(directed=nxG.is_directed())
49
50
    # Add the Graph properties as "internal properties"
51
    for key, value in nxG.graph.items():
52
        # Convert the value and key into a type for graph-tool
53
        tname, value, key = get_prop_type(value, key)
54
55
        prop = gtG.new_graph_property(tname) # Create the PropertyMap
56
        gtG.graph_properties[key] = prop     # Set the PropertyMap
57
        gtG.graph_properties[key] = value    # Set the actual value
58
59
    # Phase 1: Add the vertex and edge property maps
60
    # Go through all nodes and edges and add seen properties
61
    # Add the node properties first
62
    nprops = set() # cache keys to only add properties once
63
    for node, data in nxG.nodes(data=True):
64
65
        # Go through all the properties if not seen and add them.
66
        for key, val in data.items():
67
            if key in nprops: continue # Skip properties already added
68
69
            # Convert the value and key into a type for graph-tool
70
            tname, _, key  = get_prop_type(val, key)
71
72
            prop = gtG.new_vertex_property(tname) # Create the PropertyMap
73
            gtG.vertex_properties[key] = prop     # Set the PropertyMap
74
75
            # Add the key to the already seen properties
76
            nprops.add(key)
77
78
    # Also add the node id: in NetworkX a node can be any hashable type, but
79
    # in graph-tool node are defined as indices. So we capture any strings
80
    # in a special PropertyMap called 'id' -- modify as needed!
81
    gtG.vertex_properties['id'] = gtG.new_vertex_property('string')
82
83
    # Add the edge properties second
84
    eprops = set() # cache keys to only add properties once
85
    for src, dst, data in nxG.edges(data=True):
86
87
        # Go through all the edge properties if not seen and add them.
88
        for key, val in data.items():
89
            if key in eprops: continue # Skip properties already added
90
91
            # Convert the value and key into a type for graph-tool
92
            tname, _, key = get_prop_type(val, key)
93
94
            prop = gtG.new_edge_property(tname) # Create the PropertyMap
95
            gtG.edge_properties[key] = prop     # Set the PropertyMap
96
97
            # Add the key to the already seen properties
98
            eprops.add(key)
99
100
    # Phase 2: Actually add all the nodes and vertices with their properties
101
    # Add the nodes
102
    vertices = {} # vertex mapping for tracking edges later
103
    for node, data in nxG.nodes(data=True):
104
105
        # Create the vertex and annotate for our edges later
106
        v = gtG.add_vertex()
107
        vertices[node] = v
108
109
        # Set the vertex properties, not forgetting the id property
110
        data['id'] = str(node)
111
        for key, value in data.items():
112
            gtG.vp[key][v] = value # vp is short for vertex_properties
113
114
    # Add the edges
115
    for src, dst, data in nxG.edges(data=True):
116
117
        # Look up the vertex structs from our vertices mapping and add edge.
118
        e = gtG.add_edge(vertices[src], vertices[dst])
119
120
        # Add the edge properties
121
        for key, value in data.items():
122
            gtG.ep[key][e] = value # ep is short for edge_properties
123
124
    # Done, finally!
125
    return gtG
126
127
128
if __name__ == '__main__':
129
130
    # Create the networkx graph
131
    nxG = nx.Graph(name="Undirected Graph")
132
    nxG.add_node("v1", name="alpha", color="red")
133
    nxG.add_node("v2", name="bravo", color="blue")
134
    nxG.add_node("v3", name="charlie", color="blue")
135
    nxG.add_node("v4", name="hub", color="purple")
136
    nxG.add_node("v5", name="delta", color="red")
137
    nxG.add_node("v6", name="echo", color="red")
138
139
    nxG.add_edge("v1", "v2", weight=0.5, label="follows")
140
    nxG.add_edge("v1", "v3", weight=0.25, label="follows")
141
    nxG.add_edge("v2", "v4", weight=0.05, label="follows")
142
    nxG.add_edge("v3", "v4", weight=0.35, label="follows")
143
    nxG.add_edge("v5", "v4", weight=0.65, label="follows")
144
    nxG.add_edge("v6", "v4", weight=0.53, label="follows")
145
    nxG.add_edge("v5", "v6", weight=0.21, label="follows")
146
147
    
148
    for item in nxG.edges(data=True):
149
        print item
150
151
    # Convert to graph-tool graph
152
    gtG = nx2gt(nxG)
153
    gtG.list_properties()
154
    
155
    bw_centrality = nx.betweenness_centrality(nxG, normalized=False,  weight='weight', endpoints=False)
156
    import graph_tool.centrality as gtc
157
    ws=gtG.edge_properties["weight"]
158
    vp, ep = gtc.betweenness(gtG, norm=False, weight=ws)
159
    #code.interact(local=dict(globals(), **locals()))
160
    print vp.a==bw_centrality.values()