python - 不同大小向量的 Numpy 乘法避免循环

Question

我有一个矩阵，比方说，P尺寸(X,Y) .另外，我有两个矩阵，比如说 Kx和 Ky尺寸(M,N)两者都是一个矩阵 pk尺寸(M,N)和两个向量 u和 v的 X和 Y分别。例如，它们可以定义如下:

import numpy as np
P = np.zeros((X,Y));
pk = np.random.rand(M,N);
Kx = np.random.rand(M,N);
Ky = np.random.rand(M,N);
u = np.random.rand(X);
v = np.random.rand(Y);

当然，在实际代码中它们不是随机的，但这对于本例来说无关紧要。问题是，是否存在与以下等价的纯 numpy:

for m in range(0, M):
    for n in range(0, N):
        for i in range(0,X):
            for j in range(0,Y):
               Arg = Kx[m,n]*u[i] + Ky[m,n]*v[j];
               P[i,j] += pk[m,n]*np.cos(Arg);

全部M , N , X , Y是不同的，但是 X和 Y如果解决方案不存在，则可以相同。

Answer 1

消除 for-loop 的常用策略NumPy 计算中的 s 用于处理高维数组。

例如，考虑这条线

Arg = Kx[m,n]*u[i] + Ky[m,n]*v[j]

此行取决于索引 m , n , i和 j . 所以 Arg取决于 m , n , i和 j 。这意味着 Arg可以被认为是一个由 m 索引的 4 维数组, n , i和 j .所以我们可以消除 4 for-loops -- 到 Arg是关心——通过计算

Kxu = Kx[:,:,np.newaxis]*u
Kyv = Ky[:,:,np.newaxis]*v
Arg = Kxu[:,:,:,np.newaxis] + Kyv[:,:,np.newaxis,:]

Kx[:,:,np.newaxis]有形状 (M, N, 1) , 和 u有形状 (X,) .将它们相乘使用 NumPy broadcasting创建形状数组 (M, N, X) .因此，上面，new axes are used somewhat like placeholders ，所以 Arg最终得到由 m 索引的 4 个轴, n , i , j按照这个顺序。

同样，P可以定义为

P = (pk[:,:,np.newaxis,np.newaxis]*np.cos(Arg)).sum(axis=0).sum(axis=0)

sum(axis=0) (调用两次)沿着 m 求和和 n轴，所以 P最终成为由 i 索引的二维数组和 j仅。

通过使用这些 4 维数组，我们可以对整个 NumPy 数组应用 NumPy 操作。相反，当使用 4 for-loops 时，我们必须对标量进行逐值计算。例如考虑什么 np.cos(Arg) Arg 时正在做什么是一个 4 维数组。这在一个 NumPy 函数调用中减轻了 所有 余弦的计算，该函数调用在已编译的 C 代码中执行底层循环。这比调用 np.cos 快得多每个标量一次。这就是为什么使用高维数组最终比 for-loop 快得多的原因。基于代码。

---

import numpy as np

def orig(Kx, Ky, u, v, pk):
    M, N = Kx.shape
    X = u.size
    Y = v.size
    P = np.empty((X, Y), dtype=pk.dtype)
    for m in range(0, M):
        for n in range(0, N):
            for i in range(0,X):
                for j in range(0,Y):
                   Arg = Kx[m,n]*u[i] + Ky[m,n]*v[j]
                   P[i,j] += pk[m,n]*np.cos(Arg)
    return P

def alt(Kx, Ky, u, v, pk):
    Kxu = Kx[:,:,np.newaxis]*u
    Kyv = Ky[:,:,np.newaxis]*v
    Arg = Kxu[:,:,:,np.newaxis] + Kyv[:,:,np.newaxis,:]
    P = (pk[:,:,np.newaxis,np.newaxis]*np.cos(Arg)).sum(axis=0).sum(axis=0)
    return P

M, N = 10, 20
X, Y = 5, 15
Kx = np.random.random((M, N))
Ky = np.random.random((M, N))
u = np.random.random(X)
v = np.random.random(Y)
pk = np.random.random((M, N))

---

健全性检查，(显示 alt 和 orig 返回相同的结果):

In [57]: P2 = alt(Kx, Ky, u, v, pk)

In [58]: P1 = orig(Kx, Ky, u, v, pk)

In [59]: np.allclose(P1, P2)
Out[59]: True

---

一个基准，显示alt明显快于 orig :

In [60]: %timeit orig(Kx, Ky, u, v, pk)
10 loops, best of 3: 33.6 ms per loop

In [61]: %timeit alt(Kx, Ky, u, v, pk)
1000 loops, best of 3: 349 µs per loop