为图形添加图例

Question

这是代码

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
import numpy as np
 
fig, subs = plt.subplots(4,3) #setting the shape of the figure in one line as opposed to creating 12 variables 
 
iris = load_iris() ##code as per the example 
data = np.array(iris['data'])
targets = np.array(iris['target'])
 
cd = {0:'r',1:'b',2:"g"}
cols = np.array([cd[target] for target in targets])

 
# Row 1 
 
subs[0][0].scatter(data[:,0], data[:,1], c=cols)
subs[0][1].scatter(data[:,0], data[:,2], c=cols)
subs[0][2].scatter(data[:,0], data[:,3], c=cols)
 
# Row 2 
 
subs[1][0].scatter(data[:,1], data[:,0], c=cols)
subs[1][1].scatter(data[:,1], data[:,2], c=cols)
subs[1][2].scatter(data[:,1], data[:,3], c=cols)
 
# Row 3 
 
subs[2][0].scatter(data[:,2], data[:,0], c=cols)
subs[2][1].scatter(data[:,2], data[:,1], c=cols)
subs[2][2].scatter(data[:,2], data[:,3], c=cols)
 
#Row 4 
 
subs[3][0].scatter(data[:,3], data[:,0], c=cols)
subs[3][1].scatter(data[:,3], data[:,1], c=cols)
subs[3][2].scatter(data[:,3], data[:,2], c=cols)
 
plt.show()

我有兴趣添加一个图例，指示红点代表'setosa'、绿点代表'versicolor' 和蓝点代表'virginica'。该图例将位于上图的底部和中心。我该怎么做？

我想我必须和fig.legend 一起玩，但我完全不知道该怎么做。

Answer 1

您可以遍历其中一个子图中的目标，并使图例出现在该图之外。这是我用你的代码得到的：

代码如下：

import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
import numpy as np

fig, subs = plt.subplots(4,3, constrained_layout=True) #setting the shape of the figure in one line as opposed to creating 12 variables

iris = load_iris() ##code as per the example
data = np.array(iris['data'])
target_names = iris['target_names']
targets = np.array(iris['target'])

cd = {0:'r',1:'b',2:"g"}
cols = np.array([cd[target] for target in targets])


# Row 1

subs[0][0].scatter(data[:,0], data[:,1], c=cols)
subs[0][1].scatter(data[:,0], data[:,2], c=cols)
subs[0][2].scatter(data[:,0], data[:,3], c=cols)

# Row 2

subs[1][0].scatter(data[:,1], data[:,0], c=cols)
subs[1][1].scatter(data[:,1], data[:,2], c=cols)
subs[1][2].scatter(data[:,1], data[:,3], c=cols)

# Row 3

subs[2][0].scatter(data[:,2], data[:,0], c=cols)
subs[2][1].scatter(data[:,2], data[:,1], c=cols)
subs[2][2].scatter(data[:,2], data[:,3], c=cols)

# Row 4
subs[3][0].scatter(data[:,3], data[:,0], c=cols)

# loop for central subplot at last row
for t, name in zip(np.unique(targets), target_names):
    subs[3][1].scatter(data[targets==t,3], data[targets==t,1], c=cd[t], label=name)
subs[3][1].legend(bbox_to_anchor=(2, -.2), ncol=len(target_names))  # you can play with bbox_to_anchor for legend position

subs[3][2].scatter(data[:,3], data[:,2], c=cols)


plt.savefig('legend')

编辑：我还发现了这个post in the matplotlib documentation，您可以在其中直接从散点图中提取散点元素（不使用for 循环）。我在 IRIS 数据集上进行了尝试，但无法使其工作。

Answer 2

将label='versicor' 等添加到您的一个子图中：

# Row 1 
 
subs[0][0].scatter(data[:,0], data[:,1], c=cols, label='virginica')
subs[0][1].scatter(data[:,0], data[:,2], c=cols, label='setosa')
subs[0][2].scatter(data[:,0], data[:,3], c=cols, label='versicolor')

那么你可以在plt.show()之前调用fig.legend(loc='lower center', ncol=3)。

我已设置ncol=3 使其短而宽。

请参阅下面的示例：

import numpy as np
import matplotlib.pyplot as plt

a = range(1,11)
b = np.random.randn(10).cumsum()
c = np.random.randn(10).cumsum()
d = np.random.randn(10).cumsum()

fig, (ax1, ax2) = plt.subplots(2, figsize=(9,5))

ax1.plot(a,b, label = 'one')
ax1.plot(a,c, label = 'two')
ax1.plot(a,d, label = 'two')

ax2.plot(a,c)

fig.legend(loc='lower center', ncol=3)

plt.show()