使用Numpy vectorize含if语句时元素顺序致结果不一致问题咨询

先定义两个简单函数：

def x(t):
    return 0 if t < 0 else 1

def h(t):
    return 0 if t < 0 else np.exp(-t)

对两个函数分别执行x = np.vectorize(x)和h = np.vectorize(h)后，第一个函数表现符合预期，输出结果和输入数组元素顺序无关：

x(np.array([-1,0,1,2])) -> array([0, 1, 1, 1])
x(np.array([0,1,2,-1])) -> array([1, 1, 1, 0])

但第二个函数的结果会随输入中负数元素的位置变化：

h(np.array([-1,0,1,2])) -> array([0, 1, 0, 0])
h(np.array([0,1,2,-1])) -> array([1., 0.3678, 0.1353, 0.])

这是Numpy的Bug吗？

这不是Numpy的Bug，是np.vectorize的返回类型推断机制导致的：

• np.vectorize会根据函数对第一个输入元素的返回值类型，推断整个输出数组的类型。
• 对于h函数，当输入数组第一个元素是-1时，函数返回整数0，np.vectorize就会把输出数组类型设为整数。后续元素计算出的浮点数（比如np.exp(-0)=1.0、np.exp(-1)≈0.3678）会被强制转换为整数，所以变成1、0，最终得到array([0, 1, 0, 0])。
• 当输入数组第一个元素是0时，函数返回浮点数1.0，np.vectorize就会把输出数组类型设为浮点数，后续所有结果都保留小数，得到正常的浮点数组。

解决方法 #

1. 直接用Numpy原生向量化操作替代np.vectorize（推荐，因为np.vectorize本质是Python循环，效率远低于原生向量化）：

import numpy as np

t1 = np.array([-1,0,1,2])
t2 = np.array([0,1,2,-1])

def x(t):
    return np.where(t < 0, 0, 1)

def h(t):
    return np.where(t < 0, 0, np.exp(-t))

# 测试
print(h(t1))  # 输出 array([0.        , 1.        , 0.36787944, 0.13533528])
print(h(t2))  # 输出 array([1.        , 0.36787944, 0.13533528, 0.        ])

2. 如果一定要用np.vectorize，显式指定输出类型：
   通过otypes参数强制指定输出为浮点类型，避免自动推断出错：

import numpy as np

def h(t):
    return 0 if t < 0 else np.exp(-t)

h_vec = np.vectorize(h, otypes=[np.float64])

t1 = np.array([-1,0,1,2])
t2 = np.array([0,1,2,-1])

print(h_vec(t1))  # 输出 array([0.        , 1.        , 0.36787944, 0.13533528])
print(h_vec(t2))  # 输出 array([1.        , 0.36787944, 0.13533528, 0.        ])

内容的提问来源于stack exchange，提问作者sboeser