如何在np.vectorized向量化函数中传递*args参数？

常规Python中*args的用法可以正常展开参数，但当函数被np.vectorize向量化后，如果args是二维数组，函数会将数组的所有行拆分成独立参数传入，而非每行对应一组参数。如果函数有大量输入参数，不想逐个定义如my_function(a,b,c,d,e)的入参，可以通过以下方式解决：

原代码及错误 #

import numpy as np

a = np.random.rand(50,1)
b = np.random.rand(50,1)

args = np.random.rand(50,3)

def my_function(a,b,c,d,e):
    
    result = a * b * c * d * e
    

    return result


my_func_vec = np.vectorize(my_function)
res = my_func_vec(a,b,*args)

# 类型错误：my_function()需要5个位置参数，但传入了52个

错误原因 #

args是形状为(50,3)的数组，使用*args时会将其拆分为50个长度为3的一维数组，加上a和b，总共传入了52个参数，而函数仅需5个位置参数，因此报错。

解决方案 #

方法1：拆分args的列作为独立参数 #

将args的每一列单独取出，作为函数的对应参数传入：

import numpy as np

a = np.random.rand(50,1)
b = np.random.rand(50,1)
args = np.random.rand(50,3)

def my_function(a,b,c,d,e):
    return a * b * c * d * e

my_func_vec = np.vectorize(my_function)
# 拆分args的三列作为c、d、e参数
res = my_func_vec(a, b, args[:,0], args[:,1], args[:,2])

方法2：利用signature参数适配批量参数 #

通过np.vectorize的signature参数指定输入参数的形状，让函数能正确识别每组参数：

import numpy as np

a = np.random.rand(50,1)
b = np.random.rand(50,1)
args = np.random.rand(50,3)

def my_function(a,b,args_arr):
    # 从数组中拆分出c、d、e
    c, d, e = args_arr
    return a * b * c * d * e

# signature指定：a和b是单元素，args_arr是长度为3的数组，输出为单元素
my_func_vec = np.vectorize(my_function, signature='(),(),(3)->()')
res = my_func_vec(a, b, args)

方法3：修改函数接收*args并适配向量化 #

调整函数定义，直接接收*args，并在函数内部处理参数：

import numpy as np

a = np.random.rand(50,1)
b = np.random.rand(50,1)
args = np.random.rand(50,3)

def my_function(a,b,*args):
    # args会被拆分为三个对应列的数组
    c, d, e = args
    return a * b * c * d * e

my_func_vec = np.vectorize(my_function)
res = my_func_vec(a, b, args[:,0], args[:,1], args[:,2])

内容的提问来源于stack exchange，提问作者THATS MY QUANT MY QUANTITATIVE