Setting up the legend for a scatterplot in Python

Viewed 95

I have imported some data from an excel file onto Python and I wanted to do a scatterplot. The problem I'm having is with the legend. I include the code below:

data_2 = xlrd.open_workbook('/Users/.../.../Book1.xlsx')
sheet = data_2.sheet_by_index(0)
sheet.cell_value(0,0)
fig = plt.figure(figsize = (11,8), facecolor='white', edgecolor='black')
plt.xlabel('some x label')
plt.ylabel('some y label')
plt.xticks(np.arange(0.0, 30, 1.0))
plt.yticks(np.arange(-9.50, 1.0, 0.50))
plt.axhline(y=0.0, color='grey', linestyle=':')
plt.margins(0.01)
plt.savefig('plot.png')

for i in range(1, sheet.nrows):
    xval = sheet.cell_value(i,0)
    for j in range(1, sheet.ncols):
        yval = sheet.cell_value(i,j)
        plt.scatter(xval,yval, label = 'i = {}'.format(j))
plt.legend(loc='center left', bbox_to_anchor=(1, 0.5)) #the legend will appear on the right of the figure

While the scatterplot is fine, the legend entries repeat 31 times (31 being the number of columns in the dataset). I cannot figure out what is causing the problem. I have tried different placement of the line plt.legend(..) but unsuccessfully. I have tried building the legend from scratch but extracting the colours of the dots from the plot was hard! How can I modify the code to output the legend with sheet.ncols - 1(The first column in the dataset are the xaxis values) entries only on the right?

If it helps, this is the graph I get when I use Excel to plot the data (but I don't want to use Excel for plots) enter image description here

I want the code to reproduce such a plot in python. Note the legend placement is at the bottom but it would be better to have it on the right of the figure

2 Answers

Try only labeling one row's worth of entries, with something like plt.scatter(xval,yval, label = 'i = {}'.format(j) if i==1 else None). Matplotlib will automatically create a legend entry for each labeled object on the plot regardless of the label's content.

import xlrd
import numpy as np
from matplotlib import pyplot as plt

data_2 = xlrd.open_workbook(".\Book1.xls")
sheet = data_2.sheet_by_index(0)
sheet.cell_value(0,0)
fig = plt.figure(figsize = (11,8), facecolor="white", edgecolor="black")
plt.xlabel("x")
plt.ylabel("y")
plt.xticks(np.arange(0.0, 30, 1.0))
plt.yticks(np.arange(-9.50, 1.0, 0.50))
plt.axhline(y=0.0, color="grey", linestyle=":")
plt.margins(0.01)

for i in range(1, sheet.nrows):
    xval = sheet.cell_value(i, 0)
    for j in range(1, sheet.ncols):
        yval = sheet.cell_value(i, j)
        plt.scatter(xval, yval, label = "i={}".format(j))
plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))

plt.savefig("plot.png")

I'm not sure what you did wrong, because I basically used your code, only shifting the for loop up so the plot didn't come up empty, but didn't get the error you got.

Diagram with only one legend

The dataset is simple:

0, -2, -0.5, 1, 0, -1.5
1, -4, -2, -3, 1, 0.7
2, 1, 0.2, 0.3, 0.88, 1
Related