Моя функция работает со скаляром, но не с массивом. Что я могу сделать, чтобы преодолеть эту проблему?

Question

def diffusitivity(z):
   a=1
   if (0<=z and z<1.0/2.0):
        temp=2*(1+a)*(1+2*a)/(a^2)*z*((1-2*z)**(1/a))
        return temp
   elif (1/2<=z and z<=1):
        temp=2*(1+a)*(1+2*a)/(a^2)*(1-z)*((2*z-1)**(1/a))
        return temp
   else:
        return 0

Работает нормально, когда на входе z одно число. Но это не работает, когда z - это пустой массив. Я думаю, что это логический оператор, который не работает с массивом NumPy. Что я могу сделать, чтобы решить эту проблему?

Answer 1

Вы должны немного изменить свой образ мышления:

d = numpy.zeros_like(z)

m = z >= 0
m &= z < 0.5
d[m] = -24*z[m]**2 + 12*z[m]

m = z >= 0.5
m &= z <= 1
d[m] = -24*z[m]**2 + 36*z[m] - 12

Уравнения немного упрощаются с a = 1 в вашем примере, но идея состоит в том, чтобы использовать маску для выборочного применения условий (затем повторно использовать для экономии памяти).

Answer 2

Вам необходимо векторизовать функцию перед применением к массиву или вектору.

import numpy as np
vectorized_func = np.vectorize(diffusitivity)
print(vectorized_func(np.array([1,2,3,4])))