Numpy高效计算：对两个数组对应元素应用函数后的求和

问题根源 #

你之前用zip加列表推导的方式速度慢，本质是因为Python级别的显式循环没有利用numpy的向量化运算能力——numpy的优势在于把运算推到底层C实现，而手动循环会带来大量Python调用开销，当矩阵数量很大时，这种开销会被放大。

最优方案：向量化改造函数 #

这是效率最高的解决方式，核心是把your_function的逻辑改成能直接作用于批量维度的numpy运算，完全避免循环。

假设你的your_function是某个基于3×3矩阵的标量运算（比如矩阵乘积的迹、对应元素乘积和等），我们可以利用numpy的axis参数指定对每个3×3子矩阵做运算。

举个例子，如果your_function是计算两个矩阵的对应元素乘积之和（即内积）：

# 原来的逐矩阵函数
def your_function(m1, m2):
    return np.sum(m1 * m2)

直接改造成批量运算：

# 对整个数组的最后两个轴（3×3维度）求和，再对结果总和
result = np.sum(array_1 * array_2)

再比如，如果函数是计算两个矩阵的迹的乘积：

# 原来的函数
def your_function(m1, m2):
    return np.trace(m1) * np.trace(m2)

批量版本：

# 对每个3×3矩阵计算迹，得到形状为(N,)的数组，对应相乘后求和
traces1 = np.trace(array_1, axis1=1, axis2=2)
traces2 = np.trace(array_2, axis1=1, axis2=2)
result = np.sum(traces1 * traces2)

只要能把函数逻辑拆解为numpy的内置向量化操作，就能获得最优性能。

备选方案1：用Numba加速循环 #

如果your_function的逻辑非常复杂，无法直接向量化，可以用Numba的JIT编译把Python循环转换成机器码，大幅提升速度：

from numba import jit

# 先确保your_function可以被Numba编译（如果有复杂逻辑可能需要调整）
@jit(nopython=True)
def batch_sum(array1, array2):
    total = 0.0
    count = array1.shape[0]
    for i in range(count):
        total += your_function(array1[i], array2[i])
    return total

result = batch_sum(array_1, array_2)

nopython=True会强制编译成纯机器码，避免Python对象的开销，性能接近C语言循环。

备选方案2：numpy.vectorize（仅作最后选择） #

np.vectorize是语法糖，底层依然是Python循环，但比手动列表推导略简洁，且能自动处理批量输入：

# 指定signature：输入两个(n,m)数组，输出标量
vectorized_func = np.vectorize(your_function, signature='(n,m),(n,m)->()')
# 直接对整个数组运算，得到每个矩阵对的结果，再求和
result = np.sum(vectorized_func(array_1, array_2))

注意：这个方法的性能提升有限，仅当无法改造函数且不想用Numba时使用。

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