0

dataframe with 3 columns

x       y     type
21/02   5     'a'
22/02   6     'b'
...   ...   ...

in total i have two type: 'a' and 'b'

based on a stackoverflow question and matplotlib documentation i came with following:

fig, ax = plt.subplots()   
xy = np.column_stack((df['x'],df['y'])) 
xy = xy.reshape(-1, 1, 2) 
segments = np.hstack([xy[:-1], xy[1:]]) 
coll = LineCollection(segments, color='r') 
ax.add_collection(coll)
plt.show()

This gives me following graph:

enter image description here

I think I have to do something with: coll.set_array(some_value) But I don't know how.

I got this far: I used same code as above but just changed one line:

coll = LineCollection(segments,cmap=plt.cm.gist_ncar)

enter image description here

But now the color of the segments are at random.

How do I base the color of the segments on the column 'type' in my dataframe?

J.A.Cado
  • 715
  • 5
  • 13
  • 24

1 Answers1

2

If you just have two colors, this should work:

colors = ["red" if type=="a" else "blue" for type in df["type"]] 
coll = LineCollection(segments, color=colors)

For more colors you can map the types to a list of colors in a similar way.

T. Kau
  • 593
  • 1
  • 4
  • 17