如何高效存储复用SymPy Lambdified函数？求优化及替代方案

我正在使用SymPy逐步构建方程组，随后通过scipy.optimize.least_squares进行求解。添加方程的代码如下：

def get_variables_in_equation(equation):
    free_vars = list(equation.free_symbols)
    return sorted(free_vars, key=lambda v: variables.index(v) if v in variables else float('inf'))

I_current = sp.sympify("my equation")
equations.append(I_current)
eq_lambdas.append(sp.lambdify(get_variables_in_equation(equations[-1]), equations[-1], modules='numpy'))

整体功能正常，但存在两个问题：

1. 性能问题：处理大量方程时速度较慢。
2. 持久化问题：无法直接存储lambdified函数，只能先保存SymPy方程，重新加载后再重新lambdify，当方程数量达数千个时非常耗时。

请问是否有直接存储并重新加载lambdified函数的方法，或是能优化该场景性能的方案？

一、性能优化方案 #

1. 批量处理方程与变量，减少函数调用开销 #

逐个生成lambda函数会带来大量重复的变量解析和函数创建开销，建议将所有方程合并为一个向量值函数，统一进行lambdify：

# 假设variables是全局变量列表，包含所有可能的变量
all_vars = variables
# 将所有方程组合成一个SymPy矩阵/列表
eq_system = sp.Matrix(equations)
# 一次性生成返回向量的lambda函数
system_lambda = sp.lambdify(all_vars, eq_system, modules='numpy')

# 使用时，传入变量数组即可得到所有方程的结果向量
def residuals(x):
    return system_lambda(*x).flatten()

# 传入least_squares
scipy.optimize.least_squares(residuals, x0)

这种方式利用numpy的向量化运算，比调用数千个单独的lambda函数效率高得多。

2. 优化变量排序的查找效率 #

原代码中variables.index(v)是O(n)操作，当变量数量较多时会累积大量耗时，改用字典映射优化：

# 提前构建变量到索引的映射字典
var_index_map = {var: idx for idx, var in enumerate(variables)}

def get_variables_in_equation(equation):
    free_vars = list(equation.free_symbols)
    return sorted(free_vars, key=lambda v: var_index_map.get(v, float('inf')))

字典查找是O(1)操作，能显著提升排序速度。

3. 选择更高效的lambdify后端 #

• Numba后端：安装numba后，使用modules='numba'可以将函数编译为机器码，大幅提升计算速度，尤其适合复杂方程：

  system_lambda = sp.lambdify(all_vars, eq_system, modules='numba')
  

• C代码生成：使用SymPy的代码生成功能，将方程编译为C扩展，性能最优：

  from sympy.utilities.codegen import codegen
  # 生成C代码和头文件
  [(c_file, c_code), (h_file, h_code)] = codegen(('residuals', eq_system), 'C', 'my_system', header=True)
  # 手动编译为共享库（或用setuptools自动编译），之后通过ctypes调用
  

4. 避免重复sympify解析 #

如果方程来自字符串，尽量一次性解析所有字符串，或缓存sympify的结果，避免重复解析开销：

# 假设所有方程字符串存在eq_strings列表中
equations = [sp.sympify(s) for s in eq_strings]

二、持久化方案 #

1. 使用CloudPickle序列化lambdified函数 #

标准pickle对动态生成的lambda函数支持有限，cloudpickle可以更好地处理这类对象：

import cloudpickle

# 保存函数
with open('system_lambda.pkl', 'wb') as f:
    cloudpickle.dump(system_lambda, f)

# 加载函数
with open('system_lambda.pkl', 'rb') as f:
    loaded_lambda = cloudpickle.load(f)

注意：如果使用了numba或C编译的后端，序列化可能会有兼容性问题，建议测试后使用。

2. 生成Python源代码保存 #

将lambdified函数对应的Python代码生成并保存为文件，下次直接import使用，比重新lambdify快很多：

from sympy.printing.pycode import pycode

# 生成方程系统的Python代码
code = pycode(eq_system, assign_to='residuals')
# 添加变量参数和函数定义
full_code = f"""
import numpy as np

def residuals({', '.join([str(v) for v in all_vars])}):
    {code}
    return residuals.flatten()
"""

# 保存为py文件
with open('residuals_func.py', 'w') as f:
    f.write(full_code)

# 加载时直接导入
from residuals_func import residuals

3. 优化SymPy方程的持久化 #

如果必须保存SymPy表达式，直接序列化表达式对象而非字符串，避免重新sympify的开销：

import pickle

# 保存SymPy方程列表
with open('equations.pkl', 'wb') as f:
    pickle.dump(equations, f)

# 加载方程
with open('equations.pkl', 'rb') as f:
    loaded_equations = pickle.load(f)

# 再批量lambdify（比从字符串sympify后再lambdify快很多）

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