Multicolor scatter plot legend in Python

大兔子大兔子 提交于 2019-11-29 15:52:13

In matplotlib, you can easily generate custom legends. In your example, just retrieve the color-label combinations from your dictionary and create custom patches for the legend:

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import matplotlib.patches as mpatches
import pandas as pd

#this part just recreates your dataset
wtf =  pd.read_csv("test.csv", delim_whitespace=True)
col_dict = {'convertible':'red' ,  'hatchback':'blue' , 'sedan':'purple' , 'wagon':'yellow' , 'hardtop':'green'}
wtf["colour_column"] = wtf["body-style"].map(col_dict)
wtf["comp_ratio_size"] = np.square(wtf["horsepower"] - wtf["engine-size"])

fig = plt.figure(figsize=(8,8),dpi=75)
ax = fig.gca()
ax.scatter(wtf['engine-size'],wtf['horsepower'],c=wtf["colour_column"],s=wtf['comp_ratio_size'],alpha=0.4)
ax.set_xlabel('horsepower')
ax.set_ylabel("engine size")

#retrieve values from color dictionary and attribute it to corresponding labels
leg_el = [mpatches.Patch(facecolor = value, edgecolor = "black", label = key, alpha = 0.4) for key, value in col_dict.items()]
ax.legend(handles = leg_el)

plt.show()

Output:

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!