将返回一维数组的函数应用于numpy数组的所有元素

时间:2018-07-26 14:00:51

标签: python arrays numpy

说我有一个d维数组,我想应用一个函数,该函数将一个1D数组返回到每个元素,产生一个d+1维数组。非常适合嵌入查找。例如:

def f(x):
    return np.array([x * i for i in range(6)])

m = np.random.randint(0,10, (2,3,4,5))

g = np.vectorize(f)
h = np.vectorize(f, otypes=[np.ndarray])

我希望n=g(m)具有(2,3,4,5,6)的形状,就像我编写了4个for循环一样,将f应用于m中的每个元素。

但是,g引发错误:ValueError: setting an array element with a sequence.h返回一个奇数数组,其结尾如下:

...
[array([0, 9, 18, 27, 36, 45, 54, 63, 72]),
array([0, 5, 10, 15, 20, 25, 30, 35, 40]),
array([0, 7, 14, 21, 28, 35, 42, 49, 56]), ...,
array([0, 1, 2, 3, 4, 5, 6, 7, 8]),
array([0, 4, 8, 12, 16, 20, 24, 28, 32]),
array([0, 3, 6, 9, 12, 15, 18, 21, 24])]]]], dtype=object)

我无法改变自己的期望。

我得到@hpaulj的评论,我只是找不到其他方法。

1 个答案:

答案 0 :(得分:1)

vectorize添加了一个signature参数,该参数可以满足您的需求。但是,它甚至比otypes方法要慢。为了完整起见,我将举例说明:

In [198]: def f(x):
     ...:     return np.array([x * i for i in range(6)])
In [199]: g = np.vectorize(f, signature='()->(6)')

它采用任何形状数组,并返回一个尺寸为6的新数组:

In [202]: g(1)
Out[202]: array([0, 1, 2, 3, 4, 5])
In [203]: g(np.array([1,2]))
Out[203]: 
array([[ 0,  1,  2,  3,  4,  5],
       [ 0,  2,  4,  6,  8, 10]])
In [204]: g(np.array([[1],[2]]))
Out[204]: 
array([[[ 0,  1,  2,  3,  4,  5]],

       [[ 0,  2,  4,  6,  8, 10]]])

要更快地执行此操作,请使用已编译的f方法编写numpy。没有一个函数会采用Python函数并将其重写为更快。

为此f,我们可以利用numpy广播并编写:

In [205]: def fn(x):
     ...:     return x[...,None] * np.arange(6)
     ...: 

In [207]: fn(np.array(1))
Out[207]: array([0, 1, 2, 3, 4, 5])
In [208]: fn(np.array([1,2]))
Out[208]: 
array([[ 0,  1,  2,  3,  4,  5],
       [ 0,  2,  4,  6,  8, 10]])
In [209]: fn(np.array([[1],[2]]))
Out[209]: 
array([[[ 0,  1,  2,  3,  4,  5]],

       [[ 0,  2,  4,  6,  8, 10]]])

这里矢量化很容易。有时它需要一些聪明的思考,而有时则是不可能的(特别是如果问题本质上是连续的)。有一些工具可以从迭代解决方案中生成更快的编译代码,例如numbacython