为numpy数组分配2D数组函数返回值

时间:2017-08-02 09:27:27

标签: python arrays numpy

我有一个完美的功能。它需要一个numpy 2D数组并通过数组执行某些操作然后将其返回。 我尝试使用返回的值按条件填充数组。看下面的代码:

>>> import numpy as np
>>> x=np.random.randint(20, size=(4,5))
>>> y=np.zeros_like(x) * np.nan
>>> x
array([[19,  0,  6, 17,  5],
       [18, 18, 10, 19,  9],
       [ 2,  5, 10,  5, 15],
       [ 9,  3,  0,  6,  9]])
>>> y
array([[ nan,  nan,  nan,  nan,  nan],
       [ nan,  nan,  nan,  nan,  nan],
       [ nan,  nan,  nan,  nan,  nan],
       [ nan,  nan,  nan,  nan,  nan]])
>>> y[ x>15 ] = 1000
>>> y
array([[ 1000.,    nan,    nan,  1000.,    nan],
       [ 1000.,  1000.,    nan,  1000.,    nan],
       [   nan,    nan,    nan,    nan,    nan],
       [   nan,    nan,    nan,    nan,    nan]])

问题是添加类似的函数时。

>>> def foo(x):
    return x*2

>>> y[ x>15 ] = foo(x)

Warning (from warnings module):
  File "__main__", line 1
FutureWarning: assignment exception type will change in the future
Traceback (most recent call last):
  File "<pyshell#59>", line 1, in <module>
    y[ x>15 ] = foo(x)
ValueError: boolean index array should have 1 dimension

或类似的东西:

>>> _=foo(x)
>>> y[ x>15 ]=_
Traceback (most recent call last):
  File "<pyshell#64>", line 1, in <module>
    y[ x>15 ]=_
ValueError: boolean index array should have 1 dimension

为什么它不再起作用了??

1 个答案:

答案 0 :(得分:2)

你的主要问题是布尔索引数组返回1d数组,所以:

y[x > 15]
[Out]: array([nan, nan, nan, nan, nan])

因此,如果你想分配它,你需要一个相同大小的1-d数组(或者广播到1d的东西,就像你在第一个例子中那样的0d标量)。

所以要么布尔切片你的输入

y[x > 15] = foo(x[x > 15])  # or
y[x > 15] = foo(x)[x > 15]

或者使用保留形状的np.where

y = np.where(x > 15, foo(x), y)

第一个选项更快,但np.where通常更清晰,更具扩展性。