CODEWITHSUNDEEP
matplotlib
BillyJo_rambler
BillyJo_rambler

Reputation: 583

Showing unique legend elements when building a plot from a loop

I am trying to plot an interval style plot, as shown in the code below. I am using a look to achieve this. However, when I try and view the legend, multiple legend elements are plotted from each catagory.

How can I plot it that only a single element of the legend shows i.e only one A and B entry?

Thanks. 

import matplotlib.pyplot as plt
import pandas as pd

data = pd.DataFrame({'code':['A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B'],
                     'd1':[1,2,3,4,5,1,2,3,4,5],
                     'd2':[2,3,4,5,6,2,3,4,5,6]})

f, ax = plt.subplots(1,1)
for lab, col in zip(['A', 'B'], ['red', 'green']):
    _ = data[data['code'] == lab]
    ax.plot((_['code'].values[0], _['code'].values[0]), (_['d1'], _['d2']),
            label = lab, color = col)
ax.legend()

Upvotes: 0

Views: 50

Answers (2)

Sheldore
Sheldore

Reputation: 39072

An alternate solution could look like

from collections import OrderedDict

# rest of the code

for lab, col in zip(['A', 'B'], ['red', 'green']):
    _ = data[data['code'] == lab]
    ax.plot((_['code'].values[0], _['code'].values[0]), (_['d1'], _['d2']),
            label = lab , color = col)
    i +=1 
ax.legend()

handles, labels = ax.get_legend_handles_labels()
unique = OrderedDict(zip(labels, handles))
ax.legend(unique.values(), unique.keys())

enter image description here

Upvotes: 1

jwalton
jwalton

Reputation: 5686

You were almost there. You were correctly extracting the x-values for the plotting, however, you'd forgotten to do the same for the y-values. Because of this you were creating a legend entry for every y-value.

Instead of using (_['d1'], _['d2']) in your call to ax.plot you probably wanted something like (_['d1'].values.min(), _['d2'].values.max()):

import matplotlib.pyplot as plt
import pandas as pd

data = pd.DataFrame({'code':['A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B'],
                     'd1':[1,2,3,4,5,1,2,3,4,5],
                     'd2':[2,3,4,5,6,2,3,4,5,6]})

f, ax = plt.subplots(1,1)

for lab, col in zip(['A', 'B'], ['red', 'green']):
    _ = data[data['code'] == lab]
    ax.plot((_['code'].values[0], _['code'].values[0]), (_['d1'].values[0], _['d2'].values[-1]),
            label = lab, color = col)
ax.legend()

enter image description here

Upvotes: 1

Related Questions