np.vectorize处理自定义函数返回全0值的问题及手动向量化方案

我想将以下自定义函数快速应用到numpy数组：

def f(x):
    return max(0, 1 - abs(x))

测试用的绘图代码如下：

import numpy as np
import matplotlib.pyplot as plt
x = np.linspace(start=-4, stop=4, num=100)
plt.plot(x, list(map(f, x)))

但使用np.vectorize(f)(x)时，却返回了全0向量。明明np.vectorize本质是循环实现，至少应该返回正确结果才对？

另外，我已经手动实现了该函数的向量化版本：

def f(x):
    z = np.zeros(x.shape)
    return np.fmax(z, 1 - np.abs(x), out=z)

问题原因 #

问题出在np.vectorize的输出类型自动推断上：

• 调用np.vectorize(f)(x)时，它会先传入x的第一个元素（此处为-4）推断输出类型：max(0, 1 - abs(-4)) = max(0, -3) = 0，这是一个整数。
• 由于x是linspace(-4, 4, 100)生成的数组，其中没有恰好等于0的元素，所有元素经过f计算后得到的都是0到1之间的浮点数。
• 但np.vectorize已推断输出类型为整数，这些浮点数会被强制转换为整数0，最终返回全0向量。

解决方法 #

1. 显式指定输出类型 #

使用np.vectorize时，通过otypes参数指定输出为浮点数，避免类型推断错误：

vec_f = np.vectorize(f, otypes=[np.float64])
result = vec_f(x)

2. 使用numpy原生向量化操作（推荐） #

直接用numpy的向量化函数替换Python内置函数，性能和正确性更有保障，代码也更简洁：

def f(x):
    return np.fmax(0, 1 - np.abs(x))

3. 修改原函数返回浮点数 #

将原函数中的0改为0.0，让第一次调用返回浮点数，使np.vectorize正确推断输出类型：

def f(x):
    return max(0.0, 1 - abs(x))

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