plot_utils.py 文件源码

def graph_plot(G,
               graph_name,
               nodes_color_fn=_pos_coloring,
               node_position_path="./data",
               node_position_file=True,
               show=True):

    if node_position_file:
        spring_pos = pickle.load(open(path_join(node_position_path, graph_name, 'node_pos.bin'), "rb"))
    else:
        spring_pos = nx.spring_layout(G)
        pickle.dump(spring_pos, open(path_join(node_position_path, graph_name, 'node_pos.bin'), "wb"))

    spring_pos_values = np.array(list(spring_pos.values()))
    norm_pos = np.linalg.norm(spring_pos_values, axis=1)
    nodes_color = nodes_color_fn(G, norm_pos)

    plt.figure(figsize=(5, 5))
    plt.axis("off")
    nx.draw_networkx(G, node_color=nodes_color, pos=spring_pos, camp=plt.get_cmap(CAMP), nodelist=sorted(G.nodes()))

    if show:
        plt.show()
    else:
        plt.clf()
        plt.close()

    return nodes_color