不规则矢量的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-loopNumPy计算中的s的常用策略是使用更高维数组.

例如,考虑一下这条线

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

该行依赖于指数m,n,i和j.所以Arg要看m,n,i和j.这个装置Arg可以被认为是作为一个4维阵列索引通过m,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广播来创建形状数组(M, N, X).因此,以上,新轴用来有点像占位符,以使得Arg与由索引4轴结束m,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最终被一个2维阵列由索引i和j只.

通过使用这些4维数组,我们可以在整个NumPy数组上应用NumPy操作.相反,当使用4时for-loops,我们必须在标量上按值计算.例如,考虑什么np.cos(Arg)是Arg4维阵列时正在做什么.这会在一个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