Python | 自定义colormap | 设置对称colormap

写在前面

本文记录一个常用的自定义colormap的函数，主要有两个功能实现：

• 设置第一个颜色为白色
• 设置中间的颜色为白色

主要实现函数

以下是主要的实现函数：

import cmaps
import matplotlib as mpl   
import matplotlib.pyplot as plt

def get_cmap_pr(n_colors: int, newcmap = cmaps.MPL_RdYlBu, method: int = None) -> list:
    """
    根据指定的颜色数量和方法生成一个颜色列表。

    参数:
    n_colors (int): 需要生成的颜色数量。
    newcmap (callable): 用于生成颜色的 colormap 函数，默认为 cmaps.MPL_RdYlBu。
    method (int, 可选): 修改颜色列表的方法。
        - 如果为 1，第一个颜色设为白色。
        - 如果为 2，将中间的颜色设为白色。

    返回:
    list: 生成的颜色列表。
    """
    index = list(range(1, n_colors + 1))
    
    color_list = [newcmap(i / n_colors) for i in index]
    
    if method == 1:
        color_list[0] = [1., 1., 1.]  # 第一个颜色设为白色
    elif method == 2:
        mid_index = len(color_list) // 2
        color_list[mid_index] = [1., 1., 1.]  # 将中间的颜色设为白色
        color_list[mid_index - 1] = [1., 1., 1.]
    
    return color_list

效果

如果使用的是默认的colormap，随便传入一个生成的颜色数量，则显示如下效果:

cmap1 = get_cmap_pr(10)

fig, ax = plt.subplots(figsize=(6, 1))
fig.subplots_adjust(bottom=0.5)

# 创建一个 colormap 对象
cmap = mpl.colors.ListedColormap(cmap1)

# 创建一个颜色条
norm = mpl.colors.Normalize(vmin=-10, vmax=10)
cb1 = mpl.colorbar.ColorbarBase(ax, cmap=cmap, norm=norm, orientation='horizontal')

plt.show()

使用方法1，则设置第一个颜色为白色：

cmap1 = get_cmap_pr(10,method=1)

使用方法2，则设置中间的两个颜色为白色：

cmap1 = get_cmap_pr(10,method=2)

总结

以上函数适用于设置的colorbar的颜色数量为偶数的设置情况，常用于以0为中心的正负对称colormap以及希望设置第一个颜色为白色的colormap的情况。