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之间的整数，表示颜色要点）。我将数据点分组为6个不同的簇，编号为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()

两种方法都产生相同的情节：

Answer 2

Altair在这里可以是一个很好的选择。

连续课程

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

df = pd.DataFrame(40*np.random.randn(10, 3), columns=['A', 'B','C'])

from altair import *
Chart(df).mark_circle().encode(x='A',y='B', color='C').configure_cell(width=200, height=150)

离散类

df = pd.DataFrame(10*np.random.randn(40, 2), columns=['A', 'B'])
df['C'] = np.random.choice(['alpha','beta','gamma','delta'], size=40)

from altair import *
Chart(df).mark_circle().encode(x='A',y='B', color='C').configure_cell(width=200, height=150)