如何向量化包含if语句的函数？

Question

假设我们有以下功能:

def f(x, y):
    if y == 0:
        return 0
    return x/y

这与标量值一起工作正常.不幸的是,当我尝试使用numpy的阵列用于x与y该比较y == 0被视为这导致错误的数组运算:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-13-9884e2c3d1cd> in <module>()
----> 1 f(np.arange(1,10), np.arange(10,20))

<ipython-input-10-fbd24f17ea07> in f(x, y)
      1 def f(x, y):
----> 2     if y == 0:
      3         return 0
      4     return x/y

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

我尝试使用,np.vectorize但它没有什么区别,代码仍然失败,同样的错误. np.vectorize是一个选项,它给出了我期望的结果.

我能想到的唯一解决方案是np.where在y数组上使用类似的东西:

def f(x, y):
    np.where(y == 0, 0, x/y)

这对标量不起作用.

有没有更好的方法来编写包含if语句的函数？它应该适用于标量和数组.

Answer 1

一种方法是转换x,并y以numpy的自己的函数中的数组:

def f(x, y):
    x = np.array(x)
    y = np.array(y)
    return np.where(y == 0, 0, x/y)

当其中一个x或者y是标量而另一个是numpy数组时,这将起作用.如果它们都是可以广播的阵列,它也将起作用.如果它们是不兼容形状的阵列(例如,不同长度的1D阵列),它将不起作用,但是不清楚在这种情况下期望的行为是什么.

Answer 2

我想知道你面临的问题是什么np.vectorize.它在我的系统上工作正常:

In [145]: def f(x, y):
     ...:     if y == 0:
     ...:         return 0
     ...:     return x/y

In [146]: vf = np.vectorize(f)

In [147]: vf([[3],[10]], [0,1,2,0])
Out[147]: 
array([[ 0,  3,  1,  0],
       [ 0, 10,  5,  0]])

请注意,结果dtype由第一个元素的结果决定.您也可以自己设置所需的输出:

In [148]: vf = np.vectorize(f, otypes=[np.float])

In [149]: vf([[3],[10]], [0,1,2,0])
Out[149]: 
array([[  0. ,   3. ,   1.5,   0. ],
       [  0. ,  10. ,   5. ,   0. ]])

文档中有更多示例.

Answer 3

您可以使用仅在以下位置执行除法的蒙版数组y!=0:

def f(x, y):
    x = np.atleast_1d(np.array(x))
    y = np.atleast_1d(np.ma.array(y, mask=(y==0)))
    ans = x/y
    ans[ans.mask]=0
    return np.asarray(ans)

Answer 4

一种笨重但有效的方法是基本上预处理数据:

def f(x, y):
    if type(x) == int and type(y) == int: return x/y # Will it ever be used for this?

    # Change scalars to arrays
    if type(x) == int: x = np.full(y.shape, x, dtype=y.dtype)
    if type(y) == int: y = np.full(x.shape, y, dtype=x.dtype)

    # Change all divide by zero operations to 0/1
    div_zero_idx = (y==0)
    x[div_zero_idx] = 0
    y[div_zero_idx] = 1

    return x/y

我计时了所有不同的方法:

def f_mask(x, y):
    x = np.ma.array(x, mask=(y==0))
    y = np.array(y)
    ans = x/y
    ans[ans.mask]=0
    return np.asarray(ans)

def f_where(x, y):
    x = np.array(x)
    y = np.array(y)
    return np.where(y == 0, 0, x/y)

def f_vect(x, y):
    if y == 0:
        return 0
    return x/y

vf = np.vectorize(f_vect)

print timeit.timeit('f(np.random.randint(10, size=array_length), np.random.randint(10, size=array_length))', number=10000, setup="from __main__ import f; import numpy as np; array_length=1000")
print timeit.timeit('f_mask(np.random.randint(10, size=array_length), np.random.randint(10, size=array_length))', number=10000, setup="from __main__ import f_mask; import numpy as np; array_length=1000")
print timeit.timeit('f_where(np.random.randint(10, size=array_length), np.random.randint(10, size=array_length))', number=10000, setup="from __main__ import f_where; import numpy as np; array_length=1000")
print timeit.timeit('vf(np.random.randint(10, size=array_length), np.random.randint(10, size=array_length))', number=10000, setup="from __main__ import vf; import numpy as np; array_length=(1000)")

# f
# 0.760189056396

# f_mask
# 2.24414896965

# f_where
# RuntimeWarning: divide by zero encountered in divide return np.where(y == 0, 0, x/y)
# 1.08176398277

# f_vect
# 3.45374488831

第一个功能是最快的,没有警告.如果x或y是标量,则时间比率相似.对于更高维数组,掩蔽数组方法变得相对更快(尽管它仍然是最慢的).