Matplotlib 基于现有颜色系列添加图例

Question

我使用散点图绘制了一些数据并将其指定为：

plt.scatter(rna.data['x'], rna.data['y'], s=size,
                    c=rna.data['colors'], edgecolors='none')

并且 rna.data 对象是一个 Pandas 数据框，其组织方式使得每一行代表一个数据点（'x' 和 'y' 代表坐标，'colors' 是一个 0-5 之间的整数，代表该点的颜色） . 我将数据点分组为六个不同的集群，编号为 0-5，并将集群编号放在每个集群的平均坐标处。

这将输出以下图表：

我想知道如何向该图中添加图例，指定颜色及其相应的簇号。plt.legend()要求样式代码采用诸如red_patch但它似乎不采用数字值（或数字字符串）的格式。那么如何使用 matplotlib 添加这个图例呢？有没有办法将我的数值颜色代码转换为所需的格式plt.legend()？非常感谢！

Answer 1

您可以使用空图创建图例句柄，其颜色基于散点图的颜色图和归一化。

import pandas as pd
import numpy as np; np.random.seed(1)
import matplotlib.pyplot as plt

x = [np.random.normal(5,2, size=20), np.random.normal(10,1, size=20),
     np.random.normal(5,1, size=20), np.random.normal(10,1, size=20)]
y = [np.random.normal(5,1, size=20), np.random.normal(5,1, size=20),
     np.random.normal(10,2, size=20), np.random.normal(10,2, size=20)]
c = [np.ones(20)*(i+1) for i in range(4)]

df = pd.DataFrame({"x":np.array(x).flatten(), 
                   "y":np.array(y).flatten(), 
                   "colors":np.array(c).flatten()})

size=81
sc = plt.scatter(df['x'], df['y'], s=size, c=df['colors'], edgecolors='none')

lp = lambda i: plt.plot([],color=sc.cmap(sc.norm(i)), ms=np.sqrt(size), mec="none",
                        label="Feature {:g}".format(i), ls="", marker="o")[0]
handles = [lp(i) for i in np.unique(df["colors"])]
plt.legend(handles=handles)
plt.show()

或者，您可以通过颜色列中的值过滤数据框，例如使用 groubpy，并为每个特征绘制一个散点图：

import pandas as pd
import numpy as np; np.random.seed(1)
import matplotlib.pyplot as plt

x = [np.random.normal(5,2, size=20), np.random.normal(10,1, size=20),
     np.random.normal(5,1, size=20), np.random.normal(10,1, size=20)]
y = [np.random.normal(5,1, size=20), np.random.normal(5,1, size=20),
     np.random.normal(10,2, size=20), np.random.normal(10,2, size=20)]
c = [np.ones(20)*(i+1) for i in range(4)]

df = pd.DataFrame({"x":np.array(x).flatten(), 
                   "y":np.array(y).flatten(), 
                   "colors":np.array(c).flatten()})

size=81
cmap = plt.cm.viridis
norm = plt.Normalize(df['colors'].values.min(), df['colors'].values.max())

for i, dff in df.groupby("colors"):
    plt.scatter(dff['x'], dff['y'], s=size, c=cmap(norm(dff['colors'])), 
                edgecolors='none', label="Feature {:g}".format(i))

plt.legend()
plt.show()

两种方法产生相同的图：