Skip to content

Instantly share code, notes, and snippets.

@nkthiebaut
Created January 19, 2020 21:06
Show Gist options
  • Save nkthiebaut/d3d8ca0b7b5f2f1b23fa45127d709ce7 to your computer and use it in GitHub Desktop.
Save nkthiebaut/d3d8ca0b7b5f2f1b23fa45127d709ce7 to your computer and use it in GitHub Desktop.
Plot a graph with Matplotlib
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import networkx as nx
def plot_graph(df, directed=False):
graph_engine = nx.DiGraph() if directed else None
G = nx.from_pandas_edgelist(df, source='source', target='target', edge_attr=True, create_using=graph_engine)
costs = G.edges.data('cost')
pos = nx.layout.spring_layout(G)
pos_labels = nx.layout.spring_layout(G, center=(0.05,0.05))
node_sizes = [10 for i in range(len(G))]
M = G.number_of_edges()
edge_colors = [x[2] for x in G.edges.data('cost')]
nodes = nx.draw_networkx_nodes(G, pos, node_size=node_sizes, node_color='blue')
labels = nx.draw_networkx_labels(G, pos=pos_labels)
edges = nx.draw_networkx_edges(G, pos, node_size=node_sizes, arrowstyle='->',
arrowsize=10, edge_color=edge_colors,
edge_vmin=0., edge_cmap=plt.cm.Blues, width=2)
# set alpha value for each edge
if directed:
edge_alphas = [0.5 for i in range(M)]
for i in range(M):
edges[i].set_alpha(edge_alphas[i])
pc = mpl.collections.PatchCollection(edges, cmap=plt.cm.Blues)
pc.set_array(edge_colors)
plt.colorbar(pc)
ax = plt.gca()
ax.set_axis_off()
plt.show()
df = pd.DataFrame({"source": [1, 2, 3], "target": [2, 3, 1], "cost": [4, 5, 6]})
plot_graph(df)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment