Передача функции Python на массивы

Question

Допустим, у нас есть особенно простая функция, такая как

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

Эта функция, очевидно, работает для нескольких встроенных типов данных Python x и y, таких как string, list, int, float, array и т. Д. Поскольку мы особенно заинтересованы в массивах, мы рассмотрим два массива:

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.vectorize в scipy / numpy. Этот метод, тем не менее, оказался не очень эффективным. Может ли кто-нибудь придумать более надежный способ передачи какой-либо функции в целом на массивы numpy?

Если перезапись кода в виде массива является единственным способом, было бы полезно, если бы вы упомянули об этом и здесь.

Answer 1

np.vectorize - это общий способ преобразования функций Python, работающих с числами, в простые функции, работающие с ndarrays.

Однако, как вы указываете, это не очень быстро, так как он использует цикл Python "под капотом".

Чтобы добиться лучшей скорости, вы должны вручную изготовить функцию, которая ожидает вводные массивы в качестве входных данных и использует эту уникальную комбинацию:

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'sfunc2a() (протестировано с N = 100).

Answer 3

Просто, чтобы получить основную идею, вы можете изменить свою функцию, например, таким образом:

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

Таким образом, в вашем случае, что-то подобное должно быть очень разумной отправной точкой:

In []: def func(x, y):
   ..:     x, y= x[:, None], y[None, :]
   ..:     return x+ y
   ..:
In []: def func2(x, y):
   ..:     x, y= x[:, None], y[None, :]
   ..:     A, L= x+ y, x<= y
   ..:     A[L]= (x- y)[L]
   ..:     return A
   ..:
In []: x, y= arange(-2, 3), arange(-2, 3)
In []: func(x, y)
Out[]:
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]])
In []: func2(x, y)
Out[]:
array([[ 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]])

Хотя такая обработка может показаться бесполезной, это не всегда так. Всегда измеряйте фактическую производительность ваших программ и вносите затем (не ранее) необходимые изменения.

ИМХО для дополнительного преимущества: этот тип «векторизации» делает ваш код действительно непротиворечивым и читаемым.