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

时间:2014-01-20 05:05:13

标签: python python-2.7 numpy

假设我有2个一维(1D)numpy数组ab,长度分别为n1n2。我还有一个函数F(x,y),它有两个值。现在我想将该函数应用于我的两个1D数组中的每对值,因此结果将是具有形状n1, n2的2D numpy数组。二维数组的i, j元素为F(a[i], b[j])

在没有大量for循环的情况下,我无法找到这样做的方法,而且我确信在numpy中这样做的方法要简单得多(而且速度更快!)。

提前致谢!

6 个答案:

答案 0 :(得分:13)

您可以使用numpy broadcasting对两个数组进行计算,使用anewaxis转换为垂直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

答案 1 :(得分:6)

如果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

答案 2 :(得分:1)

您可以使用列表推导来创建数组数组:

import numpy as np

# Arrays
a = np.array([1, 2, 3]) # n1 = 3
b = np.array([4, 5]) # n2 = 2

# Your function (just an example)
def f(i, j):
    return i + j

result = np.array([[f(i, j)for j in b ]for i in a])
print result

<强>输出:

[[5 6]
 [6 7]
 [7 8]]

答案 3 :(得分:1)

我建议,如果您的用例更受限于产品,that you use the outer-product?

e.g:

import numpy

a = array([0, 1, 2])
b = array([0, 1, 2, 3])

numpy.outer(a,b)

返回

array([[0, 0, 0, 0],
       [0, 1, 2, 3],
       [0, 2, 4, 6]])

然后您可以应用其他转换:

numpy.outer(a,b) + 1

返回

array([[1, 1, 1, 1],
       [1, 2, 3, 4],
       [1, 3, 5, 7]])

这要快得多:

>>> import timeit
>>> timeit.timeit('numpy.array([[i*j for i in a] for j in b])', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
31.79583477973938

>>> timeit.timeit('numpy.outer(a,b)', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
9.351550102233887
>>> timeit.timeit('numpy.outer(a,b)+1', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
12.308301210403442

答案 4 :(得分:1)

作为另一种比点积更具可扩展性的替代方案,在嵌套列表推导时间的不到1/5 - 1/9时,使用numpy.newaxistook a bit more digging to find):

>>> import numpy
>>> a = numpy.array([0,1,2])
>>> b = numpy.array([0,1,2,3])

这次,使用电源功能:

>>> pow(a[:,numpy.newaxis], b)
array([[1, 0, 0, 0],
       [1, 1, 1, 1],
       [1, 2, 4, 8]])

与替代方案相比:

>>> numpy.array([[pow(i,j) for j in b] for i in a])
array([[1, 0, 0, 0],
       [1, 1, 1, 1],
       [1, 2, 4, 8]])

比较时间:

>>> import timeit
>>> timeit.timeit('numpy.array([[pow(i,j) for i in a] for j in b])', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
31.943181037902832
>>> timeit.timeit('pow(a[:, numpy.newaxis], b)', 'import numpy; a=numpy.arange(3); b=numpy.arange(4)')
5.985810041427612

>>> timeit.timeit('numpy.array([[pow(i,j) for i in a] for j in b])', 'import numpy; a=numpy.arange(10); b=numpy.arange(10)')
109.74687385559082
>>> timeit.timeit('pow(a[:, numpy.newaxis], b)', 'import numpy; a=numpy.arange(10); b=numpy.arange(10)')
11.989138126373291

答案 5 :(得分:1)

如果F()适用于广播参数,请务必使用,正如其他人所描述的那样 另一种方法是使用 np.fromfunctionfunction_on_an_int_grid会是一个更好的名字。) 以下只是将int网格映射到a-b网格,然后映射到F()

import numpy as np

def func_allpairs( F, a, b ):
    """ -> array len(a) x len(b):
        [[ F( a0 b0 )  F( a0 b1 ) ... ]
         [ F( a1 b0 )  F( a1 b1 ) ... ]
         ...
        ]
    """
    def fab( i, j ):
        return F( a[i], b[j] )  # F scalar or vec, e.g. gradient

    return np.fromfunction( fab, (len(a), len(b)), dtype=int )  # -> fab( all pairs )


#...............................................................................
def F( x, y ):
    return x + 10*y

a = np.arange( 100 )
b = np.arange( 222 )
A = func_allpairs( F, a, b )
# %timeit: 1000 loops, best of 3: 241 µs per loop -- imac i5, np 1.9.3