如何高效计算大型稀疏概率矩阵的逐列信息熵

熵的计算公式为 $H = -\sum p_i \log(p_i)$，其中概率为0的项对结果无贡献，因此无需处理稀疏矩阵中的0值，直接基于非零元素做矢量化计算即可，完全避免逐列循环和稠密转换的开销。

实现步骤 #

1. 将稀疏矩阵转为CSC格式，优化列相关操作的性能
2. 生成所有非零元素对应的列索引
3. 计算每个非零概率值的 $p\log(p)$ 贡献
4. 按列对贡献值求和，取负数得到每列的熵值

代码实现 #

import numpy as np
from scipy.sparse import csc_matrix

# 示例概率矩阵，实际使用时替换为你的稀疏矩阵，先转为CSC格式
P = [[0.2, 0.0, 0.5, 0.3, 0.0, 0.0],
     [0.5, 0.5, 0.5, 0.6, 1.0, 0.0],
     [0.0, 0.0, 0.0, 0.1, 0.0, 0.5],
     [0.3, 0.5, 0.0, 0.0, 0.0, 0.5]]
P_csc = csc_matrix(P)

# 生成每个非零元素对应的列索引
col_counts = np.diff(P_csc.indptr)
col_indices = np.repeat(np.arange(P_csc.shape[1]), col_counts)

# 计算非零元素的p*log(p)贡献
non_zero_p = P_csc.data
contrib = non_zero_p * np.log(non_zero_p)

# 按列求和后取负得到熵，结果顺序和列顺序完全对应
col_entropy = -np.bincount(col_indices, weights=contrib, minlength=P_csc.shape[1])

补充说明 #

• 上述代码计算结果和原有循环调用scipy.stats.entropy的结果完全一致，但性能提升可达数百倍。对于200万×17万的高稀疏词袋矩阵，原本需要7小时的计算可以压缩到数分钟内完成。
• 如果你使用的是原始词频计数矩阵，还可以省去提前转换概率矩阵的步骤，直接在计算阶段合并归一化操作，进一步节省内存开销：

import numpy as np
from scipy.sparse import csc_matrix

# X为词频计数CSC矩阵
X_csc = csc_matrix(your_count_matrix)

# 计算每列总词频
col_total = X_csc.sum(axis=0).A1
# 生成每个非零元素对应的列索引
col_counts = np.diff(X_csc.indptr)
col_indices = np.repeat(np.arange(X_csc.shape[1]), col_counts)

# 计算非零元素对应的概率及贡献
non_zero_x = X_csc.data
non_zero_p = non_zero_x / col_total[col_indices]
contrib = non_zero_p * np.log(non_zero_p)

col_entropy = -np.bincount(col_indices, weights=contrib, minlength=X_csc.shape[1])

• 如果需要计算以2为底的熵，将np.log替换为np.log2即可。

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