我的函数使用标量，但不是数组。我能做些什么来克服这个问题？

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数组时，它不起作用。我认为逻辑运算符不适用于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])))