WS_Net_generation.py 文件源码

def save_topology(adj, sample_folder, dataset_name, graph_name):
    graph = nx.Graph()
    path = sample_folder+'/'+dataset_name

    if not os.path.isdir(path):
        os.makedirs(path)
    # 1. transfer adj to nx
    # adj_list = list(np.squeeze(adj[0,:,:,:]))
    adj_list = list(adj)

    for src in range(len(adj_list)):
        graph.add_node(src)
        for dst in range(len(adj_list[src])):
            if adj_list[src][dst] >= 0.2: # ?? sample ?? ?? [0,1]???
                graph.add_edge(src,dst)

    # 2. read position
    pos_file = glob.glob(path+'/*.pos')
    if pos_file == []:
        pos = nx.spring_layout(graph)
        pickle.dump(pos, open(path+'/graph.pos','wb'))
    else:
        pos = pickle.load(open(pos_file[0],'rb'))

    # 3. draw graph
    nx.draw_networkx_nodes(graph, pos, node_size=300, node_color='b', alpha=0.8)
    nx.draw_networkx_edges(graph, pos, width=1.5, alpha=0.8)
    nx.draw_networkx_labels(graph, pos, font_color='w')

    plt.savefig(path+'/'+graph_name+'.png')
    plt.savefig(path+'/'+graph_name+'.pdf')
    # plt.show()

    # 4. store graph
    pickle.dump(graph, open(path+'/'+graph_name+'.graph','wb'))