在 plt.plot() 命令之外定义绘图标签 - Matplotlib

Question

I'm making a scatter plot that uses two different symbols based on a condition in the data. In a for loop iterating through the rows of the data, if a condition is met a point is plotted with a circle, and if not met the point is plotted with a square:

for i in thick.index:
    if thick['Interest'][i] == 1:
        plt.scatter(thick['NiThickness'][i], thick['GdThickness'][i], marker = 'o', color = 'b')
    else:
        plt.scatter(thick['NiThickness'][i], thick['GdThickness'][i], marker = 's', color = 'r')

where 'Interest' is a column filled with ones and zeros(zeroes?).

I'd like to have one label in the legend for the circles, and one for the squares, but if I declare label = 'circle' in the plt.scatter(...) command I get as many rows in the legend as there rows in my data file.

Is there a simple trick I'm missing?

Thanks.

Answer 1

This is the pattern I use in such cases:

label_o = 'Circle'
label_s = 'Square'
for i in thick.index:
    if thick['Interest'][i] == 1:
        plt.scatter(thick['NiThickness'][i], thick['GdThickness'][i], marker=o', color='b', label=label_o)
        label_o = None
    else:
        plt.scatter(thick['NiThickness'][i], thick['GdThickness'][i], marker='s', color='r', label=label_s)
        label_s = None

This also nicely takes care of the case where only one of the categories is present.

Answer 2

if thick is a data-frame:

idx = thick['Interest'] == 1
ax = plt.subplot(111)
ax.scatter(thick['NiThickness'][idx], thick['GdThickness'][idx],
           marker='o', color='b', label='circle')
ax.scatter(thick['NiThickness'][~idx], thick['GdThickness'][~idx],
           marker='s', color='r', label='square')