如何将SymPy分段复函数lambdify后用numba.njit加速？

问题场景 #

在复平面特定区域定义SymPy分段函数，经lambdify转换后尝试用@njit装饰加速时触发报错。根源是SymPy生成的代码依赖NumPy的select和logical_or.reduce，而这些特性不被Numba nopython模式支持。

原代码示例：

import sympy as smp
from numba import njit
import numpy as np

x = smp.symbols('x')
g = smp.sin(x)
f = smp.Piecewise((0, smp.re(x) > 2*smp.pi), (0, smp.re(x) < 0), (g, True))

# 非加速版本正常运行
f_num = smp.lambdify(x, f)
print(f_num(1+2j))

# njit装饰后报错
f_num_njit = njit(f_num)
f_num_njit(1+2j)

报错信息：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'reduce' of type Function(<ufunc 'logical_or'>)

File "<lambdifygenerated-128>", line 2:
def _lambdifygenerated(x):
    return select([logical_or.reduce((less(real(x), 0),greater(real(x), 2*pi))),True], [0,sin(x)], default=nan)
    ^

解决方案 #

核心思路是自定义lambdify的转换规则，将SymPy的Piecewise结构转换为Numba完全支持的原生Pythonif-elif-else条件判断，同时保留NumPy的复数兼容函数以维持计算效率。

完整实现代码 #

import sympy as smp
from numba import njit
import numpy as np

# 自定义Piecewise到Python原生条件判断的转换逻辑
def piecewise_to_numba_compatible(expr, args):
    pieces = []
    # 处理除默认分支外的所有条件
    for cond, expr_val in expr.args[:-1]:
        cond_func = smp.lambdify(args, cond, modules=['numpy'])
        expr_func = smp.lambdify(args, expr_val, modules=['numpy'])
        pieces.append((cond_func, expr_func))
    # 处理默认分支
    default_expr_func = smp.lambdify(args, expr.args[-1].expr, modules=['numpy'])
    
    def wrapped_func(x):
        for cond, expr_val in pieces:
            if cond(x):
                return expr_val(x)
        return default_expr_func(x)
    return wrapped_func

# 定义SymPy分段函数
x = smp.symbols('x')
# 实际场景中替换为复杂计算的g(x)
g = smp.sin(x)
f = smp.Piecewise((0, smp.re(x) > 2*smp.pi), (0, smp.re(x) < 0), (g, True))

# 生成兼容Numba的函数
f_num = piecewise_to_numba_compatible(f, [x])

# 应用njit加速
f_num_njit = njit(f_num)

# 测试用例
print(f_num_njit(1+2j))          # 实部在[0,2π]，返回sin(1+2j)
print(f_num_njit(3*np.pi + 1j))  # 实部超过2π，返回0
print(f_num_njit(-1+2j))         # 实部小于0，返回0

方案说明 #

• 规避了NumPyselect和reduce的依赖，通过原生Python条件判断适配Numba nopython模式。
• 分段内的表达式仍使用NumPy模块转换，完全保留NumPy的高效向量计算能力，适合复杂函数场景。
• 转换逻辑可扩展，支持多条件、多变量的分段函数定义。

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