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

Question

我正在寻找一种“最佳”方法来计算给定向量元素的所有两两乘积。如果向量的大小为N，则输出将是一个大小为N * (N + 1) // 2的向量，并包含所有(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等例程以及一些索引/视图技巧，但未能找到符合上述要求的解决方案

Answer 1

您还可以将此算法并行化。如果可以只分配一次足够大的阵列（该阵列上较小的视图几乎不需要任何成本），然后覆盖它，则可以实现更大的加速

示例

@numba.njit(parallel=True)
def pairwise_products_numba_2_with_allocation(vec):
    k, size = 0, vec.size
    k_vec=np.empty(vec.size,dtype=np.int64)
    output = np.empty(size * (size + 1) // 2)

    #precalculate the indices
    for i in range(size):
        k_vec[i] = k
        k+=(size-i)

    for i in numba.prange(size):
        k=k_vec[i]
        for j in range(size-i):
            output[k+j] = vec[i] * vec[j+i]

    return output

@numba.njit(parallel=True)
def pairwise_products_numba_2_without_allocation(vec,output):
    k, size = 0, vec.size
    k_vec=np.empty(vec.size,dtype=np.int64)

    #precalculate the indices
    for i in range(size):
        k_vec[i] = k
        k+=(size-i)

    for i in numba.prange(size):
        k=k_vec[i]
        for j in range(size-i):
            output[k+j] = vec[i] * vec[j+i]

    return output

计时

A=np.arange(5000)
k, size = 0, A.size
output = np.empty(size * (size + 1) // 2)

%timeit res_1=pairwise_products_numba_2_without_allocation(A,output)
#7.84 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_2=pairwise_products_numba_2_with_allocation(A)
#16.9 ms ± 325 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit res_3=pairwise_products_numba(A) #@orlp
#43.3 ms ± 134 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)