python - 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_pa​​tch 等格式，但它似乎不采用数值(或数字字符串)。那么我怎样才能使用 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()

两种方法产生相同的图: