如何使用Numpy批量调用函数计算字典内所有参数组合的结果

核心思路 #

• 用np.meshgrid生成所有参数的笛卡尔积组合，不需要引入其他第三方库
• 按参数名构造关键字参数传入目标函数，自动兼容函数默认参数
• 可选择普通遍历或numpy向量化两种实现，适配不同性能需求

代码实现 #

基础遍历版本 #

import numpy as np

def apply_all_combinations(func, args_dict):
    # 提取参数名和对应取值列表
    param_names = list(args_dict.keys())
    param_values = list(args_dict.values())
    # 生成笛卡尔积网格，indexing='ij'保证参数顺序和输入字典顺序一致
    grids = np.meshgrid(*param_values, indexing='ij')
    # 展平每个参数的网格为一维数组
    flat_grids = [grid.flatten() for grid in grids]
    # 遍历所有组合调用函数
    return [func(**dict(zip(param_names, vals))) for vals in zip(*flat_grids)]

高效向量化版本（适合大量参数组合场景） #

import numpy as np

def apply_all_combinations_vectorized(func, args_dict):
    param_names = list(args_dict.keys())
    param_values = list(args_dict.values())
    grids = np.meshgrid(*param_values, indexing='ij')
    flat_grids = [grid.flatten() for grid in grids]
    # 将自定义函数转为numpy支持的ufunc，适配数组输入
    ufunc = np.frompyfunc(
        lambda *args: func(**dict(zip(param_names, args))),
        len(param_names),
        1
    )
    # 批量运算返回numpy数组
    return ufunc(*flat_grids)

测试示例 #

# 测试用参数
args1 = {'a' : [1, 2, 3], 'b' : [6, 7]}
args2 = {'a' : [1, 2, 3], 'b' : [6, 7], 'c' : [5, 4]}

# 测试用目标函数
def fun(a, b, c = 1):
    return a + b + c

# 调用示例
print(apply_all_combinations(fun, args1))
# 输出：[8, 9, 9, 10, 10, 11]
print(apply_all_combinations_vectorized(fun, args2).tolist())
# 输出：[12, 11, 13, 12, 13, 12, 14, 13, 14, 13, 15, 14]

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