Scipy Optimize约束函数提速：单函数为何慢10倍？

问题背景 #

我有一个包含22个约束函数的优化任务，将其改为单约束函数通过数据框动态获取参数索引和阈值后，优化器运行时间变为原来的10倍。需要在保持灵活性（支持通过字典/CSV更新阈值、无需硬编码索引）的前提下，将速度提升至接近原方案水平。

原方案（高性能但硬编码） #

原约束函数通过硬编码参数索引范围和阈值实现，每个函数逻辑简单，执行效率高：

def MaxConstraint000(parameters): 
    _count = np.sum(parameters[0:2])
    return (129 - _count)

def MaxConstraint001(parameters): 
    _count = np.sum(parameters[2:5])
    return (2571 - _count)

# 共22个类似约束函数，最终约束列表如下
_Constraints = ({'type': 'ineq', 'fun': MaxConstraint000},
                {'type': 'ineq', 'fun': MaxConstraint001},
                # ... 其余20个约束
               )

改造后的慢方案（灵活但低效） #

为避免硬编码，改用单约束函数通过keyValue从数据框查询索引和阈值，但每次调用都执行数据框查询和数组转换，导致性能大幅下降：

def MaxConstraint(parameters, keyValue):
    _parameters= np.array(parameters)
    _index = np.where(DFData['KeyValues'].values == keyValue)[0] 
    
    _count = np.sum(_parameters[_index])
    _target = DFMaxValues.loc[[keyValue], ['MaxValue']].values[0][0]
    return (_target - _count )

_Constraints = ({'type': 'ineq', 'fun': MaxConstraint, 'args': ('keyValue1', )},
                {'type': 'ineq', 'fun': MaxConstraint, 'args': ('keyValue2', )},
                # ... 其余20个约束
               )

提速方案 #

1. 预计算约束映射表，避免重复查询 #

在初始化阶段提前构建每个keyValue对应的索引数组和阈值，存储在字典中，约束函数直接从字典读取，消除每次调用的查询开销：

# 预构建约束映射字典（仅初始化时执行一次）
constraint_map = {}
for key in DFMaxValues.index:
    constraint_map[key] = {
        'indices': np.where(DFData['KeyValues'].values == key)[0],
        'max_value': DFMaxValues.loc[key, 'MaxValue']
    }

def MaxConstraint(parameters, keyValue):
    params = np.array(parameters)
    constraint = constraint_map[keyValue]
    return constraint['max_value'] - np.sum(params[constraint['indices']])

2. 合并约束为单函数批量计算 #

将所有约束合并为一个函数，一次性计算所有约束结果，减少多次函数调用的开销：

# 预构建约束索引和阈值列表（初始化时执行）
constraint_indices = []
constraint_max_values = []
for key in DFMaxValues.index:
    constraint_indices.append(np.where(DFData['KeyValues'].values == key)[0])
    constraint_max_values.append(DFMaxValues.loc[key, 'MaxValue'])

# 转为numpy数组提升计算效率
constraint_indices = [np.array(inds) for inds in constraint_indices]
constraint_max_values = np.array(constraint_max_values)

def AllConstraints(parameters):
    params = np.array(parameters)
    results = []
    for inds, max_val in zip(constraint_indices, constraint_max_values):
        results.append(max_val - np.sum(params[inds]))
    return np.array(results)

# 约束定义改为单个函数
_Constraints = ({'type': 'ineq', 'fun': AllConstraints},)

3. 向量化计算（最优性能） #

构建权重矩阵，利用numpy向量化运算一次性计算所有约束结果，完全避免循环：

import numpy as np

# 构建权重矩阵：行数=约束数，列数=参数数（初始化时执行）
num_params = len(_Bounds)  # 参数数量与边界一致
num_constraints = len(DFMaxValues)
weight_matrix = np.zeros((num_constraints, num_params), dtype=np.int32)

for i, key in enumerate(DFMaxValues.index):
    inds = np.where(DFData['KeyValues'].values == key)[0]
    weight_matrix[i, inds] = 1

# 预提取阈值数组
constraint_max_values = DFMaxValues['MaxValue'].values

def VectorizedConstraints(parameters):
    params = np.array(parameters)
    # 矩阵乘法计算所有约束的参数和，再与阈值计算差值
    sums = weight_matrix @ params
    return constraint_max_values - sums

_Constraints = ({'type': 'ineq', 'fun': VectorizedConstraints},)

4. 优化动态查询（如需保留动态更新） #

如果需要动态修改约束，无法预计算，可将数据框转为numpy数组或字典提升查询速度：

# 预转换为numpy数组和字典（初始化时执行）
key_values_array = DFData['KeyValues'].values
max_values_dict = DFMaxValues['MaxValue'].to_dict()

def FastMaxConstraint(parameters, keyValue):
    params = np.array(parameters)
    inds = np.where(key_values_array == keyValue)[0]
    return max_values_dict[keyValue] - np.sum(params[inds])

方案效果 #

• 预计算映射表可将性能提升至原方案的80%-90%
• 批量计算或向量化计算可将性能提升至原方案水平，甚至更优（减少函数调用开销）
• 所有方案均保留灵活性：只需更新CSV/字典后重新生成约束映射表/权重矩阵即可，无需修改约束函数代码

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