Python中的多色散点图图例 [英] Multicolor scatter plot legend in Python

问题描述

我有一些基本的汽车发动机尺寸，马力和车身类型数据(如下所示)

I have some basic car engine size, horsepower and body type data (sample shown below)

         body-style  engine-size  horsepower
0   convertible          130       111.0
2     hatchback          152       154.0
3         sedan          109       102.0
7         wagon          136       110.0
69      hardtop          183       123.0

在其中我制作了一个散点图，在x轴上具有马力，在y轴上具有引擎尺寸，并使用body-style作为配色方案来区分车身类别. 我还使用了单独数据框中每辆车的压缩比"来指示点的大小

Out of which I made a scatter plot with horsepower on x axis, engine size on y axis and using body-style as a color scheme to differentiate body classes and. I also used "compression ratio" of each car from a seperate dataframe to dictate the point point size

效果很好，除了我无法为我的情节显示颜色图例.我是初学者，需要帮助.

This worked out well except I cant display color legends for my plot. Help needed as i'm a beginner.

这是我的代码:

dict = {'convertible':'red' ,  'hatchback':'blue' , 'sedan':'purple' , 'wagon':'yellow' , 'hardtop':'green'}

wtf["colour column"] = wtf["body-style"].map(dict)
wtf["comp_ratio_size"] = df['compression-ratio'].apply ( lambda x : x*x)

fig = plt.figure(figsize=(8,8),dpi=75)
ax = fig.gca()
plt.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")
ax.legend()

推荐答案

在matplotlib中，您可以轻松生成

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()

输出:

这篇关于Python中的多色散点图图例的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！