Numpy quirk:将函数应用于两个1D数组的所有对,以获得一个2D数组

Question

比方说,我有2个一维(1D)numpy的阵列,a并b与长度n1和n2分别.我还有一个函数F(x,y),它有两个值.现在我想将该函数应用于我的两个1D数组中的每对值,因此结果将是具有形状的2D numpy数组n1, n2.i, j二维数组的元素是F(a[i], b[j]).

如果没有大量的for循环,我无法找到这样做的方法,而且我确信在numpy中有一个更简单(更快!)的方法.

提前致谢!

Answer 1

您可以使用numpy广播对两个阵列进行计算,a使用newaxis以下方法转换为垂直2D阵列:

In [11]: a = np.array([1, 2, 3]) # n1 = 3
    ...: b = np.array([4, 5]) # n2 = 2
    ...: #if function is c(i, j) = a(i) + b(j)*2:
    ...: c = a[:, None] + b*2

In [12]: c
Out[12]: 
array([[ 9, 11],
       [10, 12],
       [11, 13]])

基准测试:

In [28]: a = arange(100)

In [29]: b = arange(222)

In [30]: timeit r = np.array([[f(i, j) for j in b] for i in a])
10 loops, best of 3: 29.9 ms per loop

In [31]: timeit c = a[:, None] + b*2
10000 loops, best of 3: 71.6 us per loop

Answer 2

如果F超出您的控制范围,您可以使用自动将其包装为"矢量感知" numpy.vectorize.我在下面提供一个工作示例,我将其定义F为完整性.这种方法具有简单的优点,但是如果你有控制权F,那么用一点小心重写它来正确地进行矢量化可以获得巨大的速度优势

import numpy

n1 = 100
n2 = 200

a = numpy.arange(n1)
b = numpy.arange(n2)

def F(x, y):
    return x + y

# Everything above this is setup, the answer to your question lies here:
fv = numpy.vectorize(F)
r = fv(a[:, numpy.newaxis], b)

在我的计算机上,找到以下时间,显示您为"自动"矢量化支付的价格:

%timeit fv(a[:, numpy.newaxis], b)
100 loops, best of 3: 3.58 ms per loop

%timeit F(a[:, numpy.newaxis], b)
10000 loops, best of 3: 38.3 µs per loop