ego_graph

ego_graph.py

import networkx as nx

# SAMPLE DATA FORMAT
#nodes = [('tensorflow', {'count': 13}),
# ('pytorch', {'count': 6}),
# ('keras', {'count': 6}),
# ('scikit', {'count': 2}),
# ('opencv', {'count': 5}),
# ('spark', {'count': 13}), ...]

#edges = [('pytorch', 'tensorflow', {'weight': 10, 'distance': 1}),
# ('keras', 'tensorflow', {'weight': 9, 'distance': 2}),
# ('scikit', 'tensorflow', {'weight': 8, 'distance': 3}),
# ('opencv', 'tensorflow', {'weight': 7, 'distance': 4}),
# ('spark', 'tensorflow', {'weight': 1, 'distance': 10}), ...]

#BUILD THE INITIAL FULL GRAPH
G=nx.Graph()
G.add_nodes_from(nodes)
G.add_edges_from(edges)

#BUILD THE EGO GRAPH FOR TENSORFLOW
EG = nx.ego_graph(G, 'tensorflow', distance = 'distance', radius = 22)

#FIND THE 2-CONNECTED SUBGRAPHS
subgraphs = nx.algorithms.connectivity.edge_kcomponents.k_edge_subgraphs(EG, k = 3)

#GET THE SUBGRAPH THAT CONTAINS TENSORFLOW
for s in subgraphs:
    if 'tensorflow' in s:
        break
pruned_EG = EG.subgraph(s)

ego_nodes = pruned_EG.nodes()
ego_edges = pruned_EG.edges()