Home > Mobile >  how to scatter plot in different colours and add a legend in python?
how to scatter plot in different colours and add a legend in python?

Time:09-10

from the following df I have extracted the desired rows

**name**  **event**        **X**         **Y**
peter       run             20             50
peter       jog             30             25
peter       swim            21             22 
peter       walk            28             32
jon         swim            20             51

I've run the following code

mask_activity = df.event.isin (['run','jog','swim']) & (df.name == "peter")
                                                     
df_activity = df.loc[mask_activity, ['X','Y']]

I want to plot those actions within a map called pitch with three different colours and add a legend to it


pitch = Pitch(line_color='black')
fig, ax = pitch.grid(grid_height=0.9, title_height=0.06, axis=False, 
                     endnote_height=0.04, title_space=0, endnote_space=0)

pitch.scatter(df_activity.x, df_activity.y, alpha = 1, s = 500, color = "blue", ax=ax['pitch'])
fig.suptitle("activity", fontsize = 30)

plt.show()

thanks a lot in advance for any help

CodePudding user response:

The event names are added to the data that forms the basis of the scatter plot, and the data frame is looped to draw a graph by associating events and colors. If the label is set to the event name, it will be displayed in the legend setting. Since the marker size is large, the spacing between the labels in the legend is adjusted.

mask_activity = df.event.isin (['run','jog','swim']) & (df.name == "peter")                                       
df_activity = df.loc[mask_activity, ['event','X','Y']]

df_activity
    event   X   Y
0   run     20  50
1   jog     30  25
2   swim    21  22

import matplotlib.pyplot as plt
from mplsoccer import Pitch, VerticalPitch

pitch = Pitch(line_color='black')
fig, ax = pitch.grid(grid_height=0.9, title_height=0.06, axis=False, 
                     endnote_height=0.04, title_space=0, endnote_space=0)

colors = {'run': 'blue', 'jog': 'green', 'swim': 'red'}
for row in df_activity.itertuples():
    c = colors[row[1]]
    pitch.scatter(row[2], row[3], alpha = 1, s = 500, color = c, label=row[1], ax=ax['pitch'])

fig.suptitle("activity", fontsize = 30)
ax['pitch'].legend(labelspacing=2)
plt.show()

enter image description here

  • Related