Search code examples
pythonpandasnetworkx

Edgelist from pandas dataframe with nodes of different colours

I have the following data frame:

Src  Dst
A    [A,B]
B    [B,A]
C    [C]
D    [D,E,F]
E    [E,D,F]
F    [F,D,E]
...

I would like to generate a network where Src is nodes, Dst are edges, and where a new column, Weight, can assign a different colour (green) to the node which is A or D, while the others are all the same (e.g. blue).

I have tried as follows:

Create a new column Weight

nd=["A","D"]
df['Weight'] = np.where(df.Src.isin(nd), 1, 0)

The problem is here that I do not know how to assign a colour, so I just tried to assign values 1 for A or D and 0 for all the other values, and separately to change colours.

For the graph, I used the following

G = nx.from_pandas_edgelist(df, 'Src', 'Dst')

The above line of code does not connect with lines the nodes in Dst, and I cannot understand the reason.
I have found something that might be helpful in assigning colours:

   colors=[]
for n in df.Src:
    if n in df.Weight:
        colors.append('g')
    else:
        colors.append('b')

# nodes
nx.draw_networkx_nodes(G,pos, node_color = colors)

but I have got this error:

ValueError: 'c' argument has 79 elements, which is inconsistent with 'x' and 'y' with size 76.

The image below would be something similar to my expected output (A and D nodes green, others blue, and links based on Dst data; please note that the image below does NOT currently reproduce neither the colour nor the edges expected).

enter image description here

Could you please help me giving me advice on how to do it?

Solution

  • Here's a way to do that:

    df["color"] = "blue"
df.loc[df.Src.isin(["A", "D"]), "color"] = "green"

# The following line is needed because, at least in the way my dataset 
# is created, 'Dst' is not a list but rather a string. 
# For example, Dst of 'A' is the string "[A,B]". Here, 
# I'm converting it to the list ["A", "B"]
# If your data doesn't need this, just comment this line out. 
df["Dst"] = df.Dst.apply(lambda x: x[1:-1].split(","))

G = nx.from_pandas_edgelist(df.explode("Dst"), 'Src', 'Dst')
nx.draw(G, node_color = df.color)

    The output is:

    enter image description here