CODEWITHSUNDEEP
pythonpandasnetworkx
still_learning
still_learning

Reputation: 806

How to color nodes within networkx using a column in Pandas

I have this dataset:

    User    Val     Color   
92  Laura   NaN      red
100 Laura   John    red
148 Laura   Mike    red
168 Laura   Mirk    red
293 Laura   Sara    red
313 Laura   Sim     red
440 Martyn  Pierre  orange
440 Martyn  Hugh    orange
440 Martyn  Lauren  orange
440 Martyn  Sim     orange

I would like to assign to each User (no duplicates) the corresponding colour: in this example, the node called Laura should be red; the node called Martyn should be orange; the other nodes (John, Mike, Mirk, Sara, Sim, Pierrre, Hugh and Lauren) should be in green. I have tried to use this column (Color) to define a set of colours within my code by using networkx, but the approach seems to be wrong, since the nodes are not coloured as I previously described, i.e. as I would expect. Please see below the code I have used:

I am using the following code:

G = nx.from_pandas_edgelist(df, 'User', 'Val')
labels = [i for i in dict(G.nodes).keys()]
labels = {i:i for i in dict(G.nodes).keys()}
colors = df[["User", "Color"]].drop_duplicates()["Color"]

plt.figure(3,figsize=(30,50)) 
pos = nx.spring_layout(G) 
nx.draw(G, node_color = df.Color, pos = pos)
net = nx.draw_networkx_labels(G, pos = pos)

Upvotes: 1

Views: 572

Answers (1)

yatu
yatu

Reputation: 88236

Looks like you're in the right track, but got a couple of things wrong. Along with using drop_duplicates, build a dictionary and use it to lookup the color in nx.draw. Also, you don't need to construct a labels dictionary, nx.draw can handle that for you.

G = nx.from_pandas_edgelist(df, 'User', 'Val')

d = dict(df.drop_duplicates(subset=['User','Color'])[['User','Color']]
           .to_numpy().tolist())
# {'Laura': 'red', 'Martyn': 'orange'}

nodes = G.nodes()
plt.figure(figsize=(10,6)) 
pos = nx.draw(G, with_labels=True, 
              nodelist=nodes,
              node_color=[d.get(i,'lightgreen') for i in nodes], 
              node_size=1000)

enter image description here

Upvotes: 3

Related Questions