CODEWITHSUNDEEP
pythonmatplotlib
triki
triki

Reputation: 11

How to create a color bar using a dictionary in python

I am trying to create a plot on 2D plane where each coordinate represents a rectangle. The rectangle is filled with a color based on a category in the data at that point. There are unique and limited categories, I created a dictionary of {category:color} and successfully filled the rectangles. However, I want to show the color bar in the plot on the side indicating category-corresponding color.

I didn't use cmap available in the matplotlib to fill the colors, rather I used a user defined dictionary to map the colors to category. Hence, I couldn't use fig.colorbar to display the colorbar.

color_map = {'ClassA':'Red','ClassB':'Green','ClassC':'Yellow'}

def assign_colors(row):

return matplotlib.patches.Rectangle(row[x],row[y],width,height,color=color_map.get(row['Category'])

All I am looking for is help/ documentation which explains how can I create a color bar on the plot using the information from the dictionary. For example color bar in the plot specifies Class A - Red (Not String), Class B - Green and Class C- Yellow

Upvotes: 0

Views: 2318

Answers (1)

Max Kaha
Max Kaha

Reputation: 922

I have done something similar for a scatterplot I made for a project by building my own custom legend with this function:

def build_legend(data):
    """
    Build a legend for matplotlib plt from dict
    """
    legend_elements = []
    for key in data:
        legend_elements.append(Line2D([0], [0], marker='o', color='w', label=key,
                                        markerfacecolor=data[key], markersize=10))
    return legend_elements

You should then be able to use it when plotting your rectangle:

fig,ax = plt.subplots(1)
ax.add_patch(rect) # Add the patch to the Axes
legend_elements = build_legend(color_map)
ax.legend(handles=legend_elements, loc='upper left')
plt.show()

Output of this then looks similar to this legend I created:

enter image description here

Upvotes: 1

Related Questions