Search code examples
pythonmatplotliblegendscatter-plot

How to scatter plot in different colors and add a legend


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()

Solution

  • 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