自定义Matplotlib热力图色条 区分扑杀前后缺失值颜色

我用Matplotlib绘制热力图展示不同动物在各天数的得分，X轴为天数，Y轴为动物ID，单元格数值代表得分，无数据的时间点标记为NaN（Not a Number）。目前所有NaN值统一显示为灰色，但我希望将扑杀前和扑杀后的缺失值分别用两种不同颜色表示，可通过两个不同变量区分。

现有代码 #

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# Load the data from the CSV file
data = pd.read_csv("WB_1.csv")

# Step 1: Replace "NA" with pandas' NA representation
data = data.replace('NaN', pd.NA)

# Step 2: Set the pig IDs as the index and remove any empty rows or columns
data = data.set_index(data.columns[0])
data = data.dropna(axis=0, how='all')
data = data.dropna(axis=1, how='all')

# Generate the heatmap
plt.figure(figsize=(12, 8)) # Adjust the figure size as needed

# Define the colormap for valid data points (YlGnBu) and the color for NaN values (grey)
cmap_valid = cm.OrRd # You can choose a different colormap for valid data points
cmap_valid.set_bad(color='lightgrey') # Set NaN values to be displayed in grey

# Plot the heatmap using Seaborn with the specified colormaps
sns.heatmap(data, cmap=cmap_valid, annot=True, cbar_kws={'label': 'Log 10 viral load (Mean CN/mL)'},
            annot_kws={'size': 8}) # Adjust the 'size' value as needed for font size

# Add labels and title
plt.xlabel("Days After Virus Challenge")
plt.ylabel("Animal IDs")
plt.title("Viraemia in Whole blood Post challenge")

# Display the heatmap
plt.show()

数据集 #

NaN NC Both.csv #

NaN only.csv #

核心思路是先标记两类缺失值的位置，再通过分层绘制热力图实现差异化颜色显示，以下是具体步骤：

步骤1：预处理数据并标记两类缺失值 #

假设：

• 扑杀前缺失值：Day1、Day2的NC或nan（未采样/未发病阶段）
• 扑杀后缺失值：Day4、Day5的nan（动物扑杀后无数据）

先统一缺失值标记，再生成两类缺失值的掩码：

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np

# 读取本地数据（替换为你的文件路径）
data = pd.read_csv("WB_1.csv")
# 将NC和字符串NaN替换为pandas标准缺失值
data = data.replace(['NC', 'NaN'], pd.NA)
# 设置动物ID为索引并转换为数值类型
data = data.set_index(data.columns[0])
data = data.apply(pd.to_numeric, errors='coerce')
# 移除全空的行/列
data = data.dropna(axis=0, how='all')
data = data.dropna(axis=1, how='all')

# 生成两类缺失值的掩码
# 掩码1：扑杀前缺失（Day1、Day2的NaN）
mask_pre = pd.DataFrame(False, index=data.index, columns=data.columns)
mask_pre[['Day 1', 'Day 2']] = data[['Day 1', 'Day 2']].isna()
# 掩码2：扑杀后缺失（Day4、Day5的NaN）
mask_post = pd.DataFrame(False, index=data.index, columns=data.columns)
mask_post[['Day 4', 'Day 5']] = data[['Day 4', 'Day 5']].isna()

步骤2：分层绘制热力图 #

通过三次叠加绘制，分别为两类缺失值设置背景色，最后覆盖有效数据的热力图：

plt.figure(figsize=(12, 8))

# 1. 绘制扑杀前缺失值的背景色（示例用浅蓝）
sns.heatmap(data, mask=~mask_pre, cmap=['#cce5ff'], cbar=False, annot=False)
# 2. 绘制扑杀后缺失值的背景色（示例用浅红）
sns.heatmap(data, mask=~mask_post, cmap=['#ffcccc'], cbar=False, annot=False)
# 3. 绘制有效数据的热力图，未被标记的NaN设为白色
cmap_valid = cm.OrRd
cmap_valid.set_bad(color='white')
sns.heatmap(data, cmap=cmap_valid, annot=True, 
            cbar_kws={'label': 'Log 10 viral load (Mean CN/mL)'},
            annot_kws={'size': 8}, 
            mask=mask_pre | mask_post, 
            cbar=True)

# 设置标签和标题
plt.xlabel("Days After Virus Challenge")
plt.ylabel("Animal IDs")
plt.title("Viraemia in Whole blood Post challenge")

plt.show()

可选方案：自定义颜色映射（适用于复杂缺失分布） #

如果两类缺失值的位置不固定，可将其替换为极端值后自定义颜色映射：

# 复制数据并替换缺失值
data_custom = data.copy()
min_val = data_custom.min().min()
max_val = data_custom.max().max()
# 扑杀前缺失值替换为远小于最小值的数
data_custom[mask_pre] = min_val - 10
# 扑杀后缺失值替换为远大于最大值的数
data_custom[mask_post] = max_val + 10

# 自定义颜色映射：有效数据用OrRd，扑杀前用浅蓝，扑杀后用浅红
from matplotlib.colors import ListedColormap
cmap_base = cm.get_cmap('OrRd', 256)
new_colors = ['#cce5ff']  # 扑杀前缺失色
new_colors.extend(cmap_base(np.linspace(0, 1, 254)))  # 有效数据色
new_colors.append('#ffcccc')  # 扑杀后缺失色
cmap_custom = ListedColormap(new_colors)

plt.figure(figsize=(12,8))
sns.heatmap(data_custom, cmap=cmap_custom, annot=True, 
            cbar_kws={'label': 'Log 10 viral load (Mean CN/mL)'},
            annot_kws={'size': 8}, 
            vmin=min_val-10, vmax=max_val+10)

plt.xlabel("Days After Virus Challenge")
plt.ylabel("Animal IDs")
plt.title("Viraemia in Whole blood Post challenge")
plt.show()

内容的提问来源于stack exchange，提问作者Rakshith Kumar