如何用numpy.apply_along_axis优化共现矩阵关联强度计算？

关于apply_along_axis的可行性 #

用apply_along_axis实现需求可行但不推荐：它本质是封装了底层循环，速度提升有限，远不如Numpy向量化操作高效，且实现时需要额外处理行索引，代码冗余。

如果一定要用apply_along_axis，可以参考以下方式（仅作演示，不推荐）：

import numpy as np

# 将单词总频次转换为numpy数组
word_occur_np = np.array(word_occurrences)

# 定义行处理函数，需手动传递行索引
def calc_row(row, idx, word_occur):
    return row / (word_occur[idx] * word_occur)

# 逐行应用函数生成结果矩阵
association_matrix = np.zeros_like(cooc, dtype=np.float64)
for i in range(cooc.shape[0]):
    association_matrix[i] = np.apply_along_axis(calc_row, 1, cooc[i:i+1], i, word_occur_np)[0]

推荐：用Numpy广播实现高效计算 #

最快捷高效的方式是利用Numpy的广播机制，直接完成向量化运算，无需循环：

import numpy as np

# 将单词总频次转换为numpy数组
word_occur = np.array(word_occurrences)

# 计算所有(i,j)对的频次乘积矩阵（外积）
occur_product = word_occur[:, np.newaxis] * word_occur[np.newaxis, :]

# 直接生成关联强度矩阵
association_matrix = cooc / occur_product

原理说明 #

• word_occur[:, np.newaxis]将一维数组转为列向量，word_occur[np.newaxis, :]转为行向量
• 两者相乘触发广播机制，自动生成每个位置的word_occur[i] * word_occur[j]
• 最后用共现矩阵直接除以该乘积矩阵，一次性完成所有元素的计算，速度比循环快几个数量级。

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