将python函数广播到numpy数组

Question

假设我们有一个特别简单的功能

import scipy as sp
def func(x, y):
   return x + y

这个功能显然适用于几种内置Python数据类型x和y像串,列表,整数,浮点,数组,等等.因为我们是在阵列特别感兴趣,我们考虑两个数组:

x = sp.array([-2, -1, 0, 1, 2])
y = sp.array([-2, -1, 0, 1, 2])

xx = x[:, sp.newaxis]
yy = y[sp.newaxis, :]

>>> func(xx, yy)

这回来了

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

正如我们所期望的那样.

现在,如果想要将数组作为以下函数的输入？

def func2(x, y):
  if x > y:
     return x + y
  else:
     return x - y

做>>>func(xx, yy)会引起错误.

人们想出的第一个明显的方法是sp.vectorizescipy/numpy中的函数.尽管如此,这种方法已被证明效率不高.任何人都可以想到一种更强大的方式来广播任何功能到numpy数组吗？

如果以阵列友好的方式重写代码是唯一的方法,如果你也可以在这里提一下它会有所帮助.

Answer 1

np.vectorize 是一种将操作数字的Python函数转换为在ndarrays上运行的numpy函数的通用方法.

但是,正如你所指出的那样,它并不是很快,因为它是在"引擎盖下"使用Python循环.

为了获得更好的速度,你必须手工制作一个将numpy数组作为输入的函数,并利用这个numpy-ness:

import numpy as np

def func2(x, y):
    return np.where(x>y,x+y,x-y)      

x = np.array([-2, -1, 0, 1, 2])
y = np.array([-2, -1, 0, 1, 2])

xx = x[:, np.newaxis]
yy = y[np.newaxis, :]

print(func2(xx, yy))
# [[ 0 -1 -2 -3 -4]
#  [-3  0 -1 -2 -3]
#  [-2 -1  0 -1 -2]
#  [-1  0  1  0 -1]
#  [ 0  1  2  3  0]]

关于表现:

test.py:

import numpy as np

def func2a(x, y):
    return np.where(x>y,x+y,x-y)      

def func2b(x, y):
    ind=x>y
    z=np.empty(ind.shape,dtype=x.dtype)
    z[ind]=(x+y)[ind]
    z[~ind]=(x-y)[~ind]
    return z

def func2c(x, y):
    # x, y= x[:, None], y[None, :]
    A, L= x+ y, x<= y
    A[L]= (x- y)[L]
    return A

N=40
x = np.random.random(N)
y = np.random.random(N)

xx = x[:, np.newaxis]
yy = y[np.newaxis, :]

运行:

N = 30时:

% python -mtimeit -s'import test' 'test.func2a(test.xx,test.yy)'
1000 loops, best of 3: 219 usec per loop

% python -mtimeit -s'import test' 'test.func2b(test.xx,test.yy)'
1000 loops, best of 3: 488 usec per loop

% python -mtimeit -s'import test' 'test.func2c(test.xx,test.yy)'
1000 loops, best of 3: 248 usec per loop

N = 1000时:

% python -mtimeit -s'import test' 'test.func2a(test.xx,test.yy)'
10 loops, best of 3: 93.7 msec per loop

% python -mtimeit -s'import test' 'test.func2b(test.xx,test.yy)'
10 loops, best of 3: 367 msec per loop

% python -mtimeit -s'import test' 'test.func2c(test.xx,test.yy)'
10 loops, best of 3: 186 msec per loop

这似乎表明它func2a略快func2c(并且func2b非常慢).

Answer 2

对于这种特殊情况,您还可以编写一个在NumPy数组和普通Python浮点数上运行的函数:

def func2d(x, y):
    z = 2.0 * (x > y) - 1.0
    z *= y
    return x + z

这个版本的速度也是unutbu的四倍func2a()(经过测试N = 100).