如何用Matplotlib/Seaborn绘制带x-y位置与数值的子图颜色网格

问题背景 #

在R中可通过以下代码生成带栅格热力图的分面子图网格：

# R 代码
DF <- data.frame(expand.grid(seq(0, 7), seq(0, 7), seq(0, 5)))
names(DF) <- c("x", "y", "z")
DF$I <- runif(nrow(DF), 0, 1)
#      x y z          I
#   1: 0 0 0 0.70252977
#   2: 1 0 0 0.74346071
#  ---                 
# 383: 6 7 5 0.93409337
# 384: 7 7 5 0.14143277
library(ggplot2)
ggplot(DF, aes(x = x, y = y, fill = I)) + 
  facet_wrap(~z, ncol = 3) +
  geom_raster() + 
  scale_fill_viridis_c() +
  theme(legend.position = "bottom") # 图例置于底部

效果为：每个z值对应一个8×8的栅格热力图，按3列排列，底部有统一颜色条。

尝试用Python实现时，plt.imshow因需要矩阵格式数据无法直接在FacetGrid.map中使用，需自定义绘图函数处理数据格式、轴和颜色映射问题：

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

df = pd.DataFrame(list(itertools.product(range(8), range(8), range(6))), 
                  columns=['x', 'y', 'z'])
# 值的顺序和R不同，但不影响绘图
df['I'] = np.random.rand(df.shape[0])
#      x  y  z         I
# 0    0  0  0  0.076338
# 1    0  0  1  0.148386
# 2    0  0  2  0.481053
# ..  .. .. ..       ...
# 382  7  7  4  0.144188
# 383  7  7  5  0.700624
g = sns.FacetGrid(df, col='z', col_wrap=2, height=4, aspect=1)
g.map(plt.imshow, color = 'I') # <- 此处plt.imshow无法正常工作
# 如何修正？可能需要自定义绘图函数？
plt.show()

解决方案 #

核心思路是对每个z分组的长格式数据进行重塑，转为8×8矩阵后再绘制，同时统一颜色标尺、调整轴和图例位置。

方法1：自定义函数配合FacetGrid #

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

# 生成数据
df = pd.DataFrame(list(itertools.product(range(8), range(8), range(6))), 
                  columns=['x', 'y', 'z'])
df['I'] = np.random.rand(df.shape[0])

# 自定义栅格绘图函数
def plot_raster(x, y, I, **kwargs):
    # 将长格式数据重塑为8×8矩阵，origin='lower'保证y轴方向与ggplot一致
    raster_data = I.values.reshape(8, 8)
    im = plt.imshow(raster_data, cmap='viridis', origin='lower', **kwargs)
    # 设置轴刻度
    plt.xticks(range(8))
    plt.yticks(range(8))
    return im

# 创建分面网格，按3列排列
g = sns.FacetGrid(df, col='z', col_wrap=3, height=3, aspect=1)
# 映射自定义函数，传入对应列
g.map_dataframe(plot_raster, 'x', 'y', 'I')

# 添加底部统一颜色条
fig = g.fig
# 自定义颜色条位置：[左, 下, 宽, 高]
cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.03])
fig.colorbar(g.axes[0].images[0], cax=cbar_ax, orientation='horizontal')

# 调整布局避免颜色条与子图重叠
plt.subplots_adjust(bottom=0.15)
plt.show()

方法2：用seaborn.heatmap简化实现 #

heatmap可直接处理透视后的矩阵，自带轴标签控制：

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

df = pd.DataFrame(list(itertools.product(range(8), range(8), range(6))), 
                  columns=['x', 'y', 'z'])
df['I'] = np.random.rand(df.shape[0])

# 自定义热力图绘图函数
def plot_heatmap(x, y, I, **kwargs):
    # 将长格式数据转为透视表（矩阵格式）
    pivot_df = pd.pivot_table(pd.DataFrame({'x':x, 'y':y, 'I':I}), 
                              index='y', columns='x', values='I')
    # 绘制热力图，关闭子图单独颜色条
    sns.heatmap(pivot_df, cmap='viridis', cbar=False, **kwargs)

# 创建分面网格
g = sns.FacetGrid(df, col='z', col_wrap=3, height=3, aspect=1)
g.map_dataframe(plot_heatmap, 'x', 'y', 'I')

# 添加底部统一颜色条
fig = g.fig
cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.03])
fig.colorbar(g.axes[0].collections[0], cax=cbar_ax, orientation='horizontal')

plt.subplots_adjust(bottom=0.15)
plt.show()

关键说明 #

1. 数据重塑：每个z分组的长格式数据需转为8×8矩阵，reshape或pivot_table均可实现，注意origin='lower'保证y轴方向与R ggplot2一致。
2. 统一颜色标尺：所有分面共用同一颜色映射，避免子图单独显示颜色条，最后手动添加全局颜色条。
3. 图例位置：通过fig.add_axes指定底部颜色条的位置，调整布局避免重叠。

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