python - 在 NumPy 中高效计算给定向量元素的所有成对乘积

Question

我正在寻找一种“最佳”方法来计算给定向量元素的所有成对乘积。如果向量的大小为 N ，输出将是大小为 N * (N + 1) // 2 的向量并包含x[i] * x[j]所有人的值(value)(i, j)与 i <= j 配对。计算这个的简单方法如下:

import numpy as np

def get_pairwise_products_naive(vec: np.ndarray):
    k, size = 0, vec.size
    output = np.empty(size * (size + 1) // 2)
    for i in range(size):
        for j in range(i, size):
            output[k] = vec[i] * vec[j]
            k += 1
    return output

需求:

• 最大限度地减少额外的内存分配/使用:如果可能，直接写入输出缓冲区。
• 使用矢量化 NumPy 例程而不是显式循环。
• 避免额外(不必要的)计算。

我一直在玩像outer这样的例程。 , triu_indices和einsum以及一些索引/ View 技巧，但尚未找到适合上述需求的解决方案。

Answer 1

方法#1

对于使用 NumPy 进行矢量化，您可以在使用外乘法获得所有成对乘法后使用掩码，如下所示 -

def pairwise_multiply_masking(a):
    return (a[:,None]*a)[~np.tri(len(a),k=-1,dtype=bool)]

方法#2

对于真正大的输入一维数组，我们可能需要诉诸使用单循环的迭代切片方法 -

def pairwise_multiply_iterative_slicing(a):
    n = len(a)
    N = (n*(n+1))//2
    out = np.empty(N, dtype=a.dtype)
    c = np.r_[0,np.arange(n,0,-1)].cumsum()
    for ii,(i,j) in enumerate(zip(c[:-1],c[1:])):
        out[i:j] = a[ii:]*a[ii]
    return out

基准测试

我们将包括pairwise_products and pairwise_products_numba from @orlp's solution在设置中。

使用benchit包(很少有基准测试工具打包在一起；免责声明:我是它的作者)来对建议的解决方案进行基准测试。

import benchit
funcs = [pairwise_multiply_masking, pairwise_multiply_iterative_slicing, pairwise_products_numba, pairwise_products]
in_ = [np.random.rand(n) for n in [10,50,100,200,500,1000,5000]]
t = benchit.timings(funcs, in_)
t.plot(logx=True, save='timings.png')
t.speedups(-1).plot(logx=True, logy=False, save='speedups.png')

结果(pairwise_products 的时间和加速)-

从绘图趋势可以看出，对于非常大的数组，基于切片的数组将开始获胜，否则矢量化的数组会做得很好。

建议

• 我们还可以研究 numexpr 来更有效地对大型数组执行外部乘法。