使用numpy.frompyfunc创建的ufunc无法省略默认参数的问题与解决

我在Python中编写了一个带默认参数的简单函数，希望将其用于numpy数组。该函数的第二个参数默认值为1，通常可省略，但使用frompyfunc创建ufunc后，省略默认参数调用会报错。请解释原因并说明如何让原函数适配numpy数组使用。

代码示例 #

import numpy as np

# 定义带默认参数的简单函数
def my_func(x, y=1):
    return x + y

# 创建适配数组的ufunc，声明2个输入、1个输出
my_ufunc = np.frompyfunc(my_func, 2, 1)

# 传入2个参数的调用都正常工作
print(my_ufunc([1, 2, 3], 1))
print(my_ufunc([1, 2, 3], [1]))
print(my_ufunc([1, 2, 3], [1, 2, 3]))

# 省略第二个默认参数调用，报错
print(my_ufunc([1, 2, 3]))

报错信息（翻译后） #

TypeError: my_func (vectorized)() takes from 2 to 3 positional arguments but 1 were given
类型错误：向量化后的my_func()需要2到3个位置参数，但只传入了1个

原因解释 #

np.frompyfunc创建ufunc时，会严格按照你指定的输入参数数量（第二个参数2）生成向量化函数，它完全不会识别原函数的默认参数逻辑。也就是说，生成的my_ufunc被固定为必须接收2个输入，原函数里y=1的默认值在向量化过程中被直接忽略了——因为frompyfunc只关注你明确声明的参数个数，不会去解析原函数的参数默认值信息。

解决方法 #

方法1：包装函数保留默认参数逻辑 #

自己写一个包装函数，手动处理默认参数，再对包装函数做向量化：

import numpy as np

def my_func(x, y=1):
    return x + y

# 包装函数，处理y的默认值逻辑
def wrapped_func(x, y=None):
    y = 1 if y is None else y
    return my_func(x, y)

# 声明2个输入参数，由包装函数内部处理默认值
my_ufunc = np.frompyfunc(wrapped_func, 2, 1)

# 两种调用方式都正常工作
print(my_ufunc([1,2,3]))  # 输出 array([2, 3, 4], dtype=object)
print(my_ufunc([1,2,3], 2))  # 输出 array([3, 4, 5], dtype=object)

方法2：用np.vectorize替代frompyfunc #

np.vectorize能识别原函数的参数签名，包括默认参数，使用起来更灵活：

import numpy as np

def my_func(x, y=1):
    return x + y

# 直接对原函数做向量化，自动识别默认参数
my_ufunc = np.vectorize(my_func)

# 支持省略默认参数调用
print(my_ufunc([1,2,3]))  # 输出 [2 3 4]
print(my_ufunc([1,2,3], [1,2,3]))  # 输出 [2 4 6]

注意：np.vectorize本质是循环处理数组元素，和frompyfunc一样不是真正的C级向量化，适合简单场景。

方法3：直接用numpy原生运算（最优解） #

如果你的函数逻辑像示例这么简单，直接用numpy原生数组运算效率最高，完全不需要自定义向量化：

import numpy as np

arr = np.array([1,2,3])
print(arr + 1)  # 输出 [2 3 4]
print(arr + [1,2,3])  # 输出 [2 4 6]

numpy会自动做广播运算，速度比自定义向量化函数快很多。

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