用 Plotly Python 绘制 Heatmap，自定义不等距 Colorbar

用 Matplotlib 绘制 pcolormesh 填色图

Matplotlib 的 pcolormesh 可以将二维数据绘制成填色图，可以一目了然地展示数据的分布情况。官方文档的 Demo 展示了这个函数的简单使用方式。

import matplotlib.pyplot as plt
import numpy as np

np.random.seed(19680801) # 设置随机数种子, 保证每次生成的 Z 数据都相同
Z = np.random.rand(6, 10) # 二维网格数据, shape 为 (6x10)
x = np.arange(0, 11, 1)  # 长度为 11 的 x 轴坐标
y = np.arange(0, 7, 1)  # 长度为 7 的 y 轴坐标

fig = plt.figure(figsize=(16, 9)) # 创建画布
ax = fig.add_subplot(111)
img = ax.pcolormesh(x, y, Z, cmap="jet") # 添加填色实例, 设置色标
ax.set_xticks(x) # 设置 x 轴坐标
ax.set_yticks(y) # 设置 y 轴坐标
fig.colorbar(img, ax=ax) # 添加色带
fig.savefig("00_test.png", bbox_inches='tight') # 执行绘图并导出成图片

生成的 Z 数据如下，形状为 6 行 10 列：

	0	1	2	3	4	5	6	7	8	9
0	0.70 0367	0. 742 751	0 .70 928	0. 566 746	0. 977 785	0. 706 335	0. 247 916	0. 157 883	0.69 7699	0. 719 957
1	0.25 7744	0. 341 547	0. 968 761	0. 694 507	0. 466 383	0. 702 813	0. 511 786	0. 928 741	0.73 9769	0. 622 439
2	0.65 1545	0. 396 808	0. 543 239	0.7 999	0. 721 545	0. 295 364	0. 160 946	0. 206 126	0.13 4325	0. 480 605
3	0.34 2522	0. 362 969	0. 972 918	0. 110 944	0. 388 264	0. 783 066	0. 972 897	0 .48 321	0.33 6421	0. 567 419
4	0 .047 9415	0. 388 937	0. 906 304	0. 161 018	0. 743 621	0. 632 974	0 .32 418	0. 922 377	0.23 7226	0. 823 946
5	0.75 0607	0. 113 784	0. 845 361	0. 923 932	0. 220 837	0. 933 054	0. 488 999	0. 474 719	0 .089 1675	0. 229 948

绘制出的图片如下：

Matplotlib pcolormesh Demo

从这张图中，结合右边的 Colorbar 我们可以清晰地看到数据的分布情况。

目前的 Colorbar 是像彩虹一样的渐变形式，色值分层也是默认的连续数值。我们可以做下面的修改，自定义色值分层，自定义每个区间的颜色，让 Colorbar 显示成「一格一格」的形式。

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import BoundaryNorm, ListedColormap # 新增

np.random.seed(19680801) # 设置随机数种子, 保证每次生成的 Z 数据都相同
Z = np.random.rand(6, 10) # 二维网格数据, shape 为 (6x10)
x = np.arange(0, 11, 1)  # 长度为 11 的 x 轴坐标
y = np.arange(0, 7, 1)  # 长度为 7 的 y 轴坐标

fig = plt.figure(figsize=(16, 9)) # 创建画布
ax = fig.add_subplot(111)

# 修改---
# 自定义色值 LEVEL
LEVEL = [0., 0.1, 0.3, 0.4, 0.8, 0.85, 0.9]
# 自定义每个区间的颜色
cmap = ListedColormap([ '#01A0F6', '#00ECEC', '#00D800', '#019000', '#FFFF00', '#E7C000', ])  # type: ignore
norm = BoundaryNorm(LEVEL, ncolors=cmap.N, clip=True)
img = ax.pcolormesh(x, y, Z, cmap=cmap, norm=norm) # 添加填色实例
ax.set_xticks(x) # 设置 x 轴坐标
ax.set_yticks(y) # 设置 y 轴坐标
fig.colorbar(img, ax=ax, ticks=LEVEL) # 添加色带
# ---

fig.savefig("01_test.png", dpi=200, bbox_inches='tight') # 执行绘图并导出成图片

Matplotlib pcolormesh 自定义 colorbar

可以看到，我们设置的 LEVEL 并不是连续的数值，但每个色块显示的长度都是相同的。 由此我们可以很灵活地根据业务需要去调整 Colorbar ，画出我们想要的图。

用 Plotly 绘制 Heatmap 填色图

Matplotlib 画出的图都是静态的图片，如果我们想实时看到图中每个方格里的实际值是多少，该怎么做呢？

有很多库可以绘制交互式的图表，既可以生成 HTML 网页，也可以实时在 Jupyter Notebook 中展示，本文介绍 Plotly 的使用。

与 pcolormesh 效果对应的就是 Heatmap，上文中的第一个 Demo 实现出来是这种效果：

绘图代码如下：

import plotly.graph_objects as go
import numpy as np
np.random.seed(1)

np.random.seed(19680801) # 设置随机数种子, 保证每次生成的 Z 数据都相同
Z = np.random.rand(6, 10) # 二维网格数据, shape 为 (6x10)
x = np.arange(0, 11, 1)  # 长度为 11 的 x 轴坐标
y = np.arange(0, 7, 1)  # 长度为 7 的 y 轴坐标

fig = go.Figure(data=go.Heatmap(
        z=Z,
        x=x,
        y=y,
        colorscale='jet'
    )
)

fig.update_layout(
    title='Plotly Heatmap Demo',
    xaxis=dict(tickvals=x, ticktext=x, title="X Axis"),
    yaxis=dict(tickvals=y, ticktext=y, title="Y Axis"),
)

fig.write_json("02_test.html")

这便是 Plotly 画出的交互图，把鼠标放上去可以看到每个方格的具体数值，并且可以自由放大、缩小，以及将图表保存成图片下载等等。

可以注意到这个图的 Colorbar 也是连续的彩虹状，而这个图要自定义 LEVEL 就不像 Matplotlib 那样简单了，两个库的实现思路完全不同。

这里给出实现的代码：

import plotly.graph_objects as go
import numpy as np
import pandas as pd


def trans_data(level: list, colors: list):
    """转换数据

    Args:
        level(list): 自定义层级列表
        colors(list): 自定义颜色列表, 长度应比 level 小 1

    Returns:
        (tuple):
            分组转换后的数据
            区间值列表
            色标分层列表
    """
    assert len(colors) + 1 == len(level), "colors 长度应比 level 小 1"

    labels = np.arange(len(LEVEL)+1)
    ratio_list = np.linspace(0, 1, len(color_list)+1).tolist()

    colorscale_list = [[ratio_list[0], color_list[0]]]
    for idx, _ in enumerate(color_list[1:]):
        colorscale_list.append([ratio_list[idx+1], color_list[idx]])
        colorscale_list.append([ratio_list[idx+1], color_list[idx+1]])
    colorscale_list.append([ratio_list[-1], color_list[-1]])

    Z_cut = pd.DataFrame(pd.cut(
        Z.flatten(),
        bins=[-np.infty] + LEVEL + [np.infty],
        labels=labels
    ).reshape(Z.shape))
    return Z_cut, labels, colorscale_list


np.random.seed(19680801)  # 设置随机数种子, 保证每次生成的 Z 数据都相同
Z = np.random.rand(6, 10)  # 二维网格数据, shape 为 (6x10)
x = np.arange(0, 11, 1)  # 长度为 11 的 x 轴坐标
y = np.arange(0, 7, 1)  # 长度为 7 的 y 轴坐标

LEVEL = [0., 0.1, 0.3, 0.4, 0.8, 0.85, 0.9]
color_list = ['#01A0F6', '#00ECEC', '#00D800', '#019000', '#FFFF00', '#E7C000']

Z_cut, labels, colorscale_list = trans_data(LEVEL, color_list)
Z_cut_values = np.unique(Z_cut.values.flatten())

fig = go.Figure(data=go.Heatmap(
    z=Z_cut,
    x=x,
    y=y,
    colorbar=dict(
        tickvals=np.arange(Z_cut_values.min(), Z_cut_values.max()+1),
        ticktext=LEVEL
    ),
    colorscale=colorscale_list,
    customdata=Z,
    hovertemplate='x: %{x}<br>y: %{y}<br>z:  %{customdata}<extra></extra>',
)
)

fig.update_layout(
    title='Plotly Heatmap Demo',
    xaxis=dict(tickvals=x, ticktext=x, title="X Axis"),
    yaxis=dict(tickvals=y, ticktext=y, title="Y Axis"),
)

fig.write_html("03_test.html")

图表效果是这样：