如何在特定索引处将两个数组相乘？

Question

玩具示例

我有两个具有不同形状的数组，例如：

import numpy as np
matrix = np.arange(5*6*7*8).reshape(5, 6, 7, 8)
vector = np.arange(1, 20, 2)

我想要做的是将矩阵的每个元素乘以“向量”的元素之一，然后在最后两个轴上求和。为此，我有一个与“矩阵”形状相同的数组，它告诉我要使用哪个索引，例如：

Idx = (matrix+np.random.randint(0, vector.size, size=matrix.shape))%vector.size

我知道解决方案之一是：

matVec = vector[Idx]
res = np.sum(matrix*matVec, axis=(2, 3))

甚至：

res = np.einsum('ijkl, ijkl -> ij', matrix, matVec)

问题

然而，我的问题是我的数组很大，并且 matVec 的构建既需要时间又需要内存。那么有没有办法绕过这个并仍然达到相同的结果呢？

更现实的例子

这是我实际做的一个更现实的例子：

import numpy as np

order = 20
dim = 23

listOrder = np.arange(-order, order+1, 1)
N, P = np.meshgrid(listOrder, listOrder)
K = np.arange(-2*dim+1, 2*dim+1, 1)
X = np.arange(-2*dim, 2*dim, 1)

tN = np.einsum('..., p, x -> ...px', N, np.ones(K.shape, dtype=int), np.ones(X.shape, dtype=int))#, optimize=pathInt)
tP = np.einsum('..., p, x -> ...px', P, np.ones(K.shape, dtype=int), np.ones(X.shape, dtype=int))#, optimize=pathInt)
tK = np.einsum('..., p, x -> ...px', np.ones(P.shape, dtype=int), K, np.ones(X.shape, dtype=int))#, optimize=pathInt)
tX = np.einsum('..., p, x -> ...px', np.ones(P.shape, dtype=int), np.ones(K.shape, dtype=int), X)#, optimize=pathInt)
tL = tK + tX

mini, maxi = -4*dim+1, 4*dim-1
NmPp2L = np.arange(2*mini-2*order, 2*maxi+2*order+1)
Idx = (2*tL+tN-tP) - NmPp2L[0]

np.random.seed(0)
matrix = (np.random.rand(Idx.size) + 1j*np.random.rand(Idx.size)).reshape(Idx.shape)
vector = np.random.rand(np.max(Idx)+1) + 1j*np.random.rand(np.max(Idx)+1)

res = np.sum(matrix*vector[Idx], axis=(2, 3))

Answer 1

对于更大的数据数组

\n

import numpy as np\n\nmatrix = np.arange(50*60*70*80).reshape(50, 60, 70, 80)\nvector = np.arange(1, 50, 2)\nIdx = (matrix+np.random.randint(0, vector.size, size=matrix.shape))%vector.size\n

\n

并行numba加速计算并避免创建matVec。

\n

在 4 核 Intel Xeon Platinum 8259CL 上

\n

matVec = vector[Idx]\n# %timeit 48.4 ms \xc2\xb1 167 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n\nres = np.einsum(\'ijkl, ijkl -> ij\', matrix, matVec)\n# %timeit 26.9 ms \xc2\xb1 40.7 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n

\n

针对未优化的numba实施

\n

import numba as nb\n\n@nb.njit(parallel=True)\ndef func(matrix, idx, vector):\n    res = np.zeros((matrix.shape[0],matrix.shape[1]), dtype=matrix.dtype)\n    for i in nb.prange(matrix.shape[0]):\n        for j in range(matrix.shape[1]):\n            for k in range(matrix.shape[2]):\n                for l in range(matrix.shape[3]):\n                    res[i,j] += matrix[i,j,k,l] * vector[idx[i,j,k,l]]\n    return res\n\nfunc(matrix, Idx, vector)  # compile run\nfunc(matrix, Idx, vector)\n# %timeit 21.7 ms \xc2\xb1 486 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10 loops each)\n# (48.4 + 26.9) / 21.7 = ~3.47x speed up\n\nnp.testing.assert_equal(func(matrix, Idx, vector), np.einsum(\'ijkl, ijkl -> ij\', matrix, matVec))\n

\n\n

性能和进一步优化

\n

处理复数时，上述 Numba 代码应该受内存限制。事实上，matrix和Idx很大，必须从相对较慢的 RAM 中完全读取。matrix大小为或41*41*92*92*8*2 = 217.10 MiB或关于目标平台（在大多数 x86-64 Windows 平台上应为 int32 类型，在大多数 Linux x86-64 平台上应为 int64 类型）。这也是速度慢的原因：Numpy 需要在内存中写入一个大数组（更不用说在这种情况下，在 x86-64 平台上写入数据应该比读取数据慢两倍）。Idx41*41*92*92*8 = 108.55 MiB41*41*92*92*4 = 54.28 MiBvector[Idx]

\n

假设代码受内存限制，加速的唯一方法是减少从 RAM 读取的数据量。Idx这可以通过存储在uint16数组中而不是默认np.int_数据类型（大 2~4）来实现。这是可能的，因为vector.size它很小。在配备 i5-9600KF 处理器和 38-40 GiB/s RAM 的 Linux 上，这种优化带来了约 29% 的速度提升，而代码仍然主要受内存限制。该平台上的实现几乎是最佳的。

\n