使用 np.where 在 matplotlib 中按类别着色答案

【问题标题】：Color by category in matplotlib using np.where使用 np.where 在 matplotlib 中按类别着色
【发布时间】：2020-12-24 04:54:19
【问题描述】：

我正在尝试创建一个包含 100 个数据点和三个变量的散点图：x 值、y 值和类别。此信息存储在 ndarray 中。

我可以创建散点图，但我不知道如何为每个类别使用不同的颜色。我使用以下代码进行绘图，似乎工作正常（虽然它还没有完成）：

def my_plot(data, color_map):
    f, ax = plt.subplots()
    ax.scatter(data.x, data.y, s = 150, edgecolors = "r")
    return f

在我的函数中，color_map 是一个参数，它引用我创建的用于为不同类别着色的字典（总共有四个）。这是字典：

color_map = {"winter":(15, 28, 75), "spring":(92, 57, 32), "summer":(255, 253, 211), "fall":(174, 12, 12)}

我想做的是以某种方式将此color_map 集成到我的函数中，以便我的绘图中的每个点接收不同的颜色。

我认为这可以使用np.where 创建一个掩码来完成，但我不确定如何继续...

【问题讨论】：

标签： python matplotlib

【解决方案1】：

颜色值需要除以 255，因为 matplotlib 喜欢它们介于 0 和 1 之间。使用此 dict，您可以为类别创建一组颜色：

from matplotlib import pyplot as plt
from matplotlib.lines import Line2D
import pandas as pd
import numpy as np

color_map = {"winter": (15, 28, 75), "spring": (92, 57, 32), "summer": (255, 253, 211), "fall": (174, 12, 12)}
color_map = {key: (r / 255, g / 255, b / 255,) for key, (r, g, b) in color_map.items()}
N = 200
data = pd.DataFrame({'x': np.random.uniform(1, 9, N), 'y': np.random.uniform(1, 5, N),
                     'cat': np.random.choice([*color_map.keys()], N)})

fig, ax = plt.subplots()
ax.scatter(data.x, data.y, s=150, color=[color_map[c] for c in data.cat], ec='r')

handles = [Line2D([], [], marker='o', ls='', color=col, markeredgecolor='r', label=label)
           for label, col in color_map.items()]
plt.legend(handles=handles, bbox_to_anchor=[1.02, 1.02], loc='upper left')
plt.tight_layout()
plt.show()

PS：用seaborn可以生成类似的情节，也会自动添加相应的图例。请注意，当前版本的 matplotlib (3.3.1) 的 hue 参数存在问题。通常您会将其添加为hue='cat'，但在此版本中，需要通过.to_list 添加workaround。

import seaborn as sns
ax = sns.scatterplot(x='x', y='y', hue=data['cat'].to_list(), s=150, palette=color_map, edgecolor='r', data=data)

【讨论】：