Matplotlib提取corner plot对角子图并设置带色左右对齐统计标题

方法1：直接生成并排对角子图（推荐） #

无需从已生成的figure中提取，直接利用corner库的绘图逻辑，基于原始数据在新figure中绘制对角分布，同时保留原corner风格的彩色统计标题。

示例代码：

import corner
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.offsetbox import AnchoredText

# 生成测试数据
ndim = 3
np.random.seed(42)
data1 = np.random.multivariate_normal(np.zeros(ndim), np.eye(ndim), 1000)
data2 = np.random.multivariate_normal(np.ones(ndim), np.eye(ndim), 1000)

# 计算每组数据的均值、标准差（用于生成标题）
def get_dist_stats(data):
    stats = []
    for i in range(data.shape[1]):
        stats.append((np.mean(data[:, i]), np.std(data[:, i])))
    return stats

stats1 = get_dist_stats(data1)
stats2 = get_dist_stats(data2)

# 创建新figure，布局为2行ndim列
fig, axes = plt.subplots(2, ndim, figsize=(12, 6))

# 绘制第一组corner的对角子图
for i in range(ndim):
    ax = axes[0, i]
    # 绘制直方图+密度曲线，和corner默认样式一致
    corner.hist2d(data1[:, i], np.zeros_like(data1[:, i]), ax=ax, bins=20, color="#1f77b4")
    # 添加corner风格的彩色统计标题
    mean, std = stats1[i]
    at = AnchoredText(
        f"$\\mu = {mean:.2f}$\n$\\sigma = {std:.2f}$",
        loc="upper right", frameon=True,
        prop={"color": "#1f77b4", "size": 10}
    )
    ax.add_artist(at)
    ax.set_xlabel(f"Param {i}")
    ax.set_yticks([])

# 绘制第二组corner的对角子图
for i in range(ndim):
    ax = axes[1, i]
    corner.hist2d(data2[:, i], np.zeros_like(data2[:, i]), ax=ax, bins=20, color="#ff7f0e")
    mean, std = stats2[i]
    at = AnchoredText(
        f"$\\mu = {mean:.2f}$\n$\\sigma = {std:.2f}$",
        loc="upper right", frameon=True,
        prop={"color": "#ff7f0e", "size": 10}
    )
    ax.add_artist(at)
    ax.set_xlabel(f"Param {i}")
    ax.set_yticks([])

plt.tight_layout()
plt.show()

方法2：从已生成的Corner Plot中提取内容 #

如果已经有绘制好的corner figure，可以遍历其轴网格，提取对角位置的绘图元素，复制到新figure的轴中：

import corner
import numpy as np
import matplotlib.pyplot as plt

# 模拟你的场景：在同一个figure中绘制两个corner plot
ndim = 3
data1 = np.random.multivariate_normal(np.zeros(ndim), np.eye(ndim), 1000)
data2 = np.random.multivariate_normal(np.ones(ndim), np.eye(ndim), 1000)

fig, (ax_container1, ax_container2) = plt.subplots(1, 2, figsize=(12, 6))
# 获取corner生成的轴网格
corner_axes1 = np.array(corner.corner(data1, fig=fig, ax=ax_container1).axes).reshape(ndim, ndim)
corner_axes2 = np.array(corner.corner(data2, fig=fig, ax=ax_container2).axes).reshape(ndim, ndim)
# 提取对角轴
diag_axes1 = [corner_axes1[i,i] for i in range(ndim)]
diag_axes2 = [corner_axes2[i,i] for i in range(ndim)]

# 创建新figure复制内容
new_fig, new_axes = plt.subplots(2, ndim, figsize=(12,6))

# 复制第一组对角子图
for idx, orig_ax in enumerate(diag_axes1):
    dest_ax = new_axes[0, idx]
    # 复制所有绘图元素（线条、直方图补丁）
    for artist in orig_ax.get_children():
        if isinstance(artist, (plt.Line2D, plt.Patch)):
            dest_ax.add_artist(artist.copy())
    # 同步坐标轴范围、标签
    dest_ax.set_xlim(orig_ax.get_xlim())
    dest_ax.set_ylim(orig_ax.get_ylim())
    dest_ax.set_xlabel(orig_ax.get_xlabel())
    dest_ax.set_yticks([])
    # 复制统计文本
    for text in orig_ax.texts:
        dest_ax.text(text.get_x(), text.get_y(), text.get_text(), 
                     fontproperties=text.get_fontproperties(),
                     color=text.get_color())

# 复制第二组对角子图
for idx, orig_ax in enumerate(diag_axes2):
    dest_ax = new_axes[1, idx]
    for artist in orig_ax.get_children():
        if isinstance(artist, (plt.Line2D, plt.Patch)):
            dest_ax.add_artist(artist.copy())
    dest_ax.set_xlim(orig_ax.get_xlim())
    dest_ax.set_ylim(orig_ax.get_ylim())
    dest_ax.set_xlabel(orig_ax.get_xlabel())
    dest_ax.set_yticks([])
    for text in orig_ax.texts:
        dest_ax.text(text.get_x(), text.get_y(), text.get_text(), 
                     fontproperties=text.get_fontproperties(),
                     color=text.get_color())

plt.tight_layout()
plt.show()

关键说明 #

• 直接提取AxesSubplot对象无法复用，因为每个轴绑定到特定figure，必须复制其内部的绘图元素到新轴。
• 使用AnchoredText可以完美复刻corner中右上角对齐的彩色统计标题，解决普通绘图无法复用该格式的问题。
• 方法1更稳定，避免了复制元素可能出现的布局错乱问题。

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