Numpy: манипулировать элементами в зависимости от значения без цикла по всему массиву

Question

У меня есть numpy массив x, и я хотел бы выполнить действие с элементами x в зависимости от его значения. Например, предположим, что я хочу взять квадрат для всех отрицательных элементов и взять четвертую степень для всех остальных элементов. Следующий код делает свое дело

import numpy as np
x = np.array([-2,-1,0,1,2,3])
y = np.zeros(len(x))

for i in range(len(x)):
    if x[i]<0:
        y[i] = np.square(x[i])
    else:
        y[i] = np.power(x[i],4)

print(y)

Поскольку я обрабатываю большие массивы и выполняю это действие много раз, я ищу способ сделать это без зацикливания всех элементов.

Answer 1

обычно

mask = x<0
y[mask] = fn1(x[mask])
mask =~mask
y[mask] = fn2(x[mask])

Тест маски и fn написаны для работы с массивами.

Answer 2

Попробуйте это

import numpy as np
x = np.array([-2,-1,0,1,2,3])
mask1 = np.clip(x,0,None)
mask1 = np.power(mask1,4)
mask2 = np.clip(x,None,0)
mask2 = np.power(mask2,2)
res = mask1+mask2

вывод

array([ 4,  1,  0,  1, 16, 81], dtype=int32)