为何处理分段函数时使用numpy.vectorize而非numpy.piecewise？

一、直接调用自定义函数报错的原因 #

你写的普通分段函数（比如用if-else判断）是针对单个数值设计的，一旦传入numpy数组，像if x < 0这种判断会返回一个布尔数组（比如输入[-1,0,1]的话，返回[True, False, False]）。但Python的if语句只能接受单个布尔值当判断条件，没法处理数组形式的结果，这就导致了那个“真值歧义”的报错。

举个直观的例子：

import numpy as np
def fun(x):
    if x < 0:
        return x**2
    else:
        return x+1

x = np.array([-1, 0, 1])
fun(x)  # 这里直接触发报错

二、numpy.vectorize的具体作用 #

numpy.vectorize其实就是个包装工具，它能把只能处理单个数值的函数，改成可以批量处理numpy数组的版本：

• 它会自动遍历输入数组的每一个元素，把单个数值传入你的自定义函数里执行判断和计算
• 最后把所有元素的计算结果重新拼成一个numpy数组返回

比如你用它包装后的代码：

vec_fun = np.vectorize(fun)
result = vec_fun(x)  # 正常得到array([1, 1, 2])

这里vec_fun会挨个取出x里的-1、0、1，分别丢进fun里跑if-else，再把结果合起来。

三、关于numpy.piecewise报错的补充 #

你用numpy.piecewise报错，大概率是条件参数写得不对。piecewise要求传入数组形式的布尔条件，不能用标量逻辑。正确的写法应该是这样：

def piecewise_fun(x):
    return np.piecewise(x, 
                       [x < 0, x >= 0],  # 每个条件都是布尔数组，对应不同分支
                       [lambda x: x**2, lambda x: x+1])  # 每个分支的处理函数

result = piecewise_fun(np.array([-1,0,1]))  # 正常返回结果

要是你写的条件没适配数组，比如用了针对单个值的判断逻辑，也会触发同样的歧义错误。

四、为什么需要vectorize？ #

核心原因就是不用改你原来的标量函数逻辑，就能直接处理数组输入。虽然vectorize本质上还是在循环遍历数组元素（性能不如原生numpy数组操作），胜在代码简洁，能直接复用你写好的单个值判断逻辑。

如果追求性能，更推荐用numpy原生的数组掩码操作，比如：

x = np.array([-1,0,1])
result = np.where(x < 0, x**2, x+1)

这种是真正的向量化操作，速度比vectorize快很多，但需要把代码逻辑改成数组操作的形式。

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