两个numpy数组的每对列之间的差异（如何更有效地执行）？

I am trying to optimize some code, and by profiling i noticed that this particular loop takes a lot of time. Can you help me write it faster?

我正在尝试优化一些代码,通过分析我注意到这个特定的循环需要花费很多时间。你能帮我写得更快吗?

import numpy as np

rows_a, rows_v, cols = (10, 15, 3)
a = np.arange(rows_a*cols).reshape(rows_a,cols)
v = np.arange(rows_v*cols).reshape(rows_v,cols)

c = 0
for i in range(len(v)):
    D = ((a-v[i])**2).sum(axis=-1)
    c += D.min()

print(c)

Is there any numpy function that can do this efficiently?

是否有任何numpy功能可以有效地做到这一点?

1 个解决方案

#1

import numpy as np

rows_a, rows_v, cols = (10, 15, 3)
a = np.arange(rows_a*cols).reshape(rows_a,cols)
v = np.arange(rows_v*cols).reshape(rows_v,cols)

def using_loop():
    c = 0
    for i in range(len(v)):
        D = ((a-v[i])**2).sum(axis=-1)
        c += D.min()
    return c

def using_broadcasting():
    return ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).min(axis=0).sum()

In [106]: %timeit using_loop()
1000 loops, best of 3: 233 µs per loop

In [107]: %timeit using_broadcasting()
10000 loops, best of 3: 29.1 µs per loop

In [108]: assert using_loop() == using_broadcasting()

When using NumPy it usually helps to eliminate for-loops (if possible) and express the calculation with operations done on entire arrays -- or at least on arrays that are as large as possible. By doing so, you off-load more of the work to fast algorithms written in C or Fortran without intermediate Python code.

使用NumPy时,它通常有助于消除for循环(如果可能)并通过在整个数组上完成的操作来表达计算 - 或者至少在尽可能大的数组上进行计算。通过这样做,您可以将更多的工作卸载到用C或Fortran编写的快速算法,而无需中间Python代码。

In the original code, D has shape (10,) for each iteration of the loop. Since there are 15 iterations of the loop, if we could express all the values for D from all 15 iterations at once as one big array, then D would have shape (10, 15). In fact, we can do that:

在原始代码中,D对于循环的每次迭代具有形状(10,)。由于循环有15次迭代,如果我们可以将所有15次迭代中的D的所有值一次表示为一个大数组,那么D将具有形状(10,15)。事实上,我们可以这样做:

Since a has shape (10,3), a[:, np.newaxis, :] has shape (10,1,3).

由于a具有形状(10,3),[:,np.newaxis,:]具有形状(10,1,3)。

Using NumPy broadcasting, since v has shape (15,3),

使用NumPy广播,因为v具有形状(15,3),

a[:,np.newaxis,:]-v

has shape (10,15,3). Squaring, then summing on the last axis gives an array of shape (10, 15). This is the new D:

有形状(10,15,3)。平方,然后在最后一个轴上求和得到一个形状的数组(10,15)。这是新的D:

In [109]: ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).shape
Out[109]: (10, 15)

Once you have D, the rest of the calculation follows naturally.

一旦你有了D,剩下的计算就会自然而然地进行。

#1

import numpy as np

rows_a, rows_v, cols = (10, 15, 3)
a = np.arange(rows_a*cols).reshape(rows_a,cols)
v = np.arange(rows_v*cols).reshape(rows_v,cols)

def using_loop():
    c = 0
    for i in range(len(v)):
        D = ((a-v[i])**2).sum(axis=-1)
        c += D.min()
    return c

def using_broadcasting():
    return ((a[:,np.newaxis,:]-v)**2).sum(axis=-1).min(axis=0).sum()

In [106]: %timeit using_loop()
1000 loops, best of 3: 233 µs per loop

In [107]: %timeit using_broadcasting()
10000 loops, best of 3: 29.1 µs per loop

In [108]: assert using_loop() == using_broadcasting()