如何绘制两个随机变量联合概率的2D热力图

实现步骤 #

• 首先计算两个离散随机变量各自的边缘概率：权重除以总权重即可得到对应取值的发生概率
• 生成联合概率矩阵：两个变量相互独立，联合概率直接等于各自边缘概率的乘积
• 调用可视化库绘制热力图，手动设置坐标轴范围为0-1即可

完整可运行代码 #

import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

# 定义x的取值、权重，计算边缘概率
x_values = [0.25, 0.75]
x_weights = [25, 75]
x_probs = [w / sum(x_weights) for w in x_weights]

# 定义y的取值、权重，计算边缘概率
y_values = [0.4, 0.6]
y_weights = [4, 6]
y_probs = [w / sum(y_weights) for w in y_weights]

# 计算联合概率矩阵
joint_prob_matrix = np.array(x_probs).reshape(-1, 1) @ np.array(y_probs).reshape(1, -1)

# 绘制热力图
plt.figure(figsize=(6, 5))
ax = sns.heatmap(
    joint_prob_matrix,
    xticklabels=y_values,
    yticklabels=x_values,
    cmap="viridis",
    annot=True, # 开启色块内显示概率数值
    fmt=".3f"
)

# 设置x、y轴取值范围为0到1
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# 设置图表标注
ax.set_xlabel("Y 取值")
ax.set_ylabel("X 取值")
ax.set_title("X与Y联合概率分布热力图")

plt.show()

补充说明 #

如果需要得到覆盖0-1整个区间的平滑热力图，可以对离散的联合概率做插值扩展，比如用scipy.interpolate模块将离散点拟合到0-1的网格上再绘制，不过原变量本身只有两组离散取值，默认的离散热力图已经能准确反映联合概率分布情况。

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