Python 循环中不同大小向量的Numpy乘法

我有一个矩阵，比如说，

p

的大小

（X，Y）

。同样，我有两个矩阵，比如，大小为

（M，N）

的

Kx和
Ky
，大小为
（M，N）
的矩阵
pk
，以及分别为
X
和
Y
的两个向量
u
和
v
。例如，它们可以定义如下：

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
可以是相同的，如果解决方案不存在的话。
在NumPy计算中消除
for loop
s的一个常见策略是使用高维数组

例如，考虑行

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

该行取决于索引
m
、
n
、
i
和
j
。因此
Arg
取决于
m
、
n
、
i
和
j
。这意味着
Arg
可以被认为是一个由
m
、
n
、
i
和
j
索引的四维数组。因此，我们可以通过计算来消除循环

（就

Arg而言）的4个

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，）
。将它们相乘用于创建形状数组
（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
最终成为仅由
i
和
j
索引的二维数组

通过使用这些4维数组，我们可以在整个NumPy数组上应用NumPy操作。相反，当使用4
进行循环时，我们必须对标量进行逐值计算。例如，当<代码> ARG是一个4维数组时，考虑<代码> NP.COS（ARG）< /代码>。这将在一个NumPy函数调用中卸载所有余弦的计算，该函数在编译的C代码中执行底层循环。这比为每个标量调用一次
np.cos
要快得多。这就是为什么使用高维数组比基于循环的代码快得多的原因



健全性检查（显示alt和orig返回相同结果）：


显示
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

“X和Y的两个向量v和u”：是否在那里交换了
u
和
v
？当然，现在将编辑它，谢谢@EugeneB以下是如何让你的问题更清楚：用几行代码来描述数量
P
，
Kx
，
Ky
和
pk
比用一段文字要好得多。如果您可以编辑问题中的代码，使其在不需要任何编辑或猜测的情况下实际运行，那么对试图帮助您的人来说会更好。有关示例和一些建议，请参阅。谢谢，为方便起见进行了编辑。@EugeneB您是否尝试运行此代码？非常感谢您快速而详细的回复！看来这就是我一直在寻找的！
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

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