Search code examples
matplotliblegendscatter-plot

Matplotlib scatterplot legend


This is a repeat of a common question, but I've spent 4 hours reading responses and I still can't make sense of them and my code throws errors or returns unexpected results.

This task is so simple that it should be no more than a single line of code. Responses are either walls of code or they execute the scatter plot by doing a separate plt.scatter call for each class, which I don't want to do.

How do I make a legend for this? X and y data was fabricated with make_blobs.

map = {0: 'red', 1:'green', 2:"blue"}
plt.scatter(X[:,0], X[:,1], c=[map[i] for i in y])

All I want is a legend that looks like:
(red dot): 0
(green dot): 1
(blue dot): 2

ffs


Solution

  • as the full code is not available, I am just creating dummy data with 4 entries to demonstrate how to create scatter plot with the 3 colors you are talking about. This will create the legend based on what you enter in map.

    import matplotlib.pyplot as plt 
    data = {'x1': [1, 3, 6, 9], #Your X values
            'x2':[15, 11, 5, 8], #Your Y values 
            'y': [0,1,2,1]}  #The colors you want for each circle
    X=pd.DataFrame(data)
    
    map = {0: 'red', 1:'green', 2:"blue"}
    fig, ax = plt.subplots()
    for g in X.y.unique():  #For each group R/G/B
        ix = np.where(X.y == g)
        ax.scatter(X.x1[ix[0]], X.x2[ix[0]], c = map[g], label = map[g], s = 50) #Draw with said color
    ax.legend()
    plt.show()
    

    Output graph

    enter image description here