如何根据 Python 中的组突出显示/更改标签的颜色？

Question

我正在用 imshow 函数绘制一个 N×N 相似度矩阵。键变量是图形的标签。在另一个文件中，每个键都与一个组相关联（例如 A,B = group1、C = group2、E,F = group3）。我想根据它们所属的组用不同颜色突出显示图表中的标签。

fig, ax = plt.subplots(figsize = (50,50))
im = ax.imshow(matrix)
ax.set_xticks(np.arange(len(keys)))
ax.set_yticks(np.arange(len(keys)))
ax.set_xticklabels(keys)
ax.set_yticklabels(keys)

plt.setp(ax.get_xticklabels(), rotation=90, ha="right",
         rotation_mode="anchor")

fig.colorbar(im, ax=ax, shrink = 0.3, pad = 0.02)

plt.show()

Answer 1

如果您有一个与键长度相同的列表，并带有相应的组 id，您可以同时循环遍历刻度和组，并分配相应的颜色。

下面是一些示例代码：

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

keys = ['A', 'B', 'C', 'D', 'E', 'F']
groups = [1, 1, 2, 3, 3, 3]
colors = {1: 'navy', 2: 'crimson', 3: 'limegreen'}
fig, ax = plt.subplots(figsize=(12, 10))
im = ax.imshow(pd.DataFrame(np.random.rand(10, 6)).corr(), cmap='inferno', aspect='equal')
ax.set_xticks(np.arange(len(keys)))
ax.set_yticks(np.arange(len(keys)))
ax.set_xticklabels(keys, size=20)
ax.set_yticklabels(keys, size=20)

plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="anchor")

for xtick, ytick, group in zip(ax.get_xticklabels(), ax.get_yticklabels(), groups):
    xtick.set_color(colors[group])
    ytick.set_color(colors[group])

fig.colorbar(im, ax=ax, shrink=0.3, pad=0.02)

plt.show()