Как векторизовать эту простую NumPy функцию?

Question

Учитывая функцию:

def f(x, c=0.7):
    if x >= 0:
        if x <= c:
            return 0.0
        if x <= 2*c:
            return x-c
        else:
            return c
    else:
        return -f(-x, c=c)

Я хотел бы применить ее к NumPy массивам. Раньше я делал это с np.vectorize, но у меня не получается. В чем здесь идея?

Answer 1

Я просто хотел указать на следующее из документации по np.vectorize:

Функция vectorize предназначена в первую очередь для удобства, а не для производительности. Реализация, по сути, предназначена для l oop.

Так что, на самом деле, вы здесь НЕ используете возможности векторизации NumPy. Используя NumPy индексацию логического массива и np.where, вы можете переписать свою функцию так, чтобы у вас была «реальная» векторизация.

Вот идея с моей стороны. Фактический код выглядит довольно некрасиво, я должен признать, но, предварительно рассчитав логические массивы, мы минимизируем время обработки и использование памяти.

def f_vec(x, c=0.7):

    # Initialize output array of same size and type as input array
    out = np.zeros_like(x)

    # Pre-calculate boolean arrays to prevent multiple calculation in following steps
    x_gtq_0 = (x >= 0)
    x_lt_0 = (x < 0)
    x_gt_c = (x > c)
    x_ltq_2c = (x <= 2 * c)
    x_gt_2c = (x > 2 * c)
    abs_x = np.abs(x)
    abs_x_gt_c = abs_x > c
    abs_x_ltq_2c = abs_x <= 2 * c
    abs_x_gt_2c = (abs_x > 2 * c)

    # Re-writing if-else blocks as operations on before calculated boolean arrays
    out[np.where(x_gtq_0 & x_gt_c & x_ltq_2c)] = x[np.where(x_gtq_0 & x_gt_c & x_ltq_2c)] - c
    out[np.where(x_gtq_0 & x_gt_2c)] = c
    out[np.where(x_lt_0 & abs_x_gt_c & abs_x_ltq_2c)] = c - abs_x[np.where(x_lt_0 & abs_x_gt_c & abs_x_ltq_2c)]
    out[np.where(x_lt_0 & abs_x_gt_2c)] = -c

    return out

Я добавил следующую небольшую тестовую функцию, чтобы выполнить некоторые сравнения :

def test(x):

    print(x.shape)
    vfunc = np.vectorize(f)

    tic = time.perf_counter()
    res_func = vfunc(x, c=0.7)
    print(time.perf_counter() - tic)

    tic = time.perf_counter()
    res_vec = f_vec(x, c=0.7)
    print(time.perf_counter() - tic)

    print('Differences: ', np.count_nonzero(np.abs(res_func - res_vec) > 10e-9), '\n')


test((np.random.rand(10) - 0.5) * 4)
test((np.random.rand(1000, 1000) - 0.5) * 4)
test((np.random.rand(1920, 1280, 3) - 0.5) * 4)

Вот результаты:

(10,)
0.0001590869999999467
7.954300000001524e-05
Differences:  0 

(1000, 1000)
1.53853834
0.0843256779999999
Differences:  0 

(1920, 1280, 3)
10.974010127
0.7489308680000004
Differences:  0 

Таким образом, с точки зрения производительности разница между np.vectorize и фактическим векторизованным подходом огромна для больших входов. Тем не менее, если для ваших входных данных достаточно решения np.vectorize, и вы не хотите прилагать слишком много усилий для переписывания кода, придерживайтесь этого! Как я уже сказал, я просто хотел показать, что векторизация - это нечто большее.

Надеюсь, это поможет!

----------------------------------------
System information
----------------------------------------
Platform:    Windows-10-10.0.16299-SP0
Python:      3.8.1
NumPy:       1.18.1
----------------------------------------

Answer 2

Эта функция отлично работает с вашей. Попробуйте:

vfunc = np.vectorize(f)
vfunc(a, c=0.7)

Если вы все еще получаете ошибки - пожалуйста, опубликуйте их с примером входных данных