如何让SymPy的lambdify在分段函数条件中使用np.isclose替代np.equal？

问题背景 #

用SymPy定义了如下分段表达式：

from sympy import Piecewise, Eq
from sympy.abc import x

expr = Piecewise(
    (1, Eq(x, 0.0)), 
    (x, True)
)

通过lambdify转成NumPy可调用函数后，生成的代码会用np.equal做条件判断：

import inspect
from sympy import lambdify

f = lambdify(x, expr, modules="numpy")
print(''.join(inspect.findsource(f)[0]))
# 输出：
# def _lambdifygenerated(x):
#     return select([equal(x, 0),True], [1,x], default=nan)

但数值求解场景中，浮点运算的微小偏差会让equal的判断失效，想换成np.isclose，之前试了直接在modules加{"==": np.isclose}或用implemented_function自定义函数都没成——后者还触发了TypeError: Second argument must be a Boolean, not my_eq，因为Piecewise要求条件是布尔关系。

方法一：自定义Lambdify转换字典 #

直接扩展lambdify的numpy转换规则，把SymPy的Eq映射到np.isclose：

1. 导入模块并复制原有转换规则：

import numpy as np
from sympy import Piecewise, Eq, lambdify
from sympy.abc import x
from sympy.utilities.lambdify import NUMPY_TRANSLATIONS

# 复制原有规则，避免修改全局字典
custom_trans = dict(NUMPY_TRANSLATIONS)
# 添加Eq到np.isclose的映射
custom_trans['Eq'] = np.isclose

2. 用自定义转换字典执行lambdify：

expr = Piecewise(
    (1, Eq(x, 0.0)), 
    (x, True)
)

f = lambdify(x, expr, modules=[custom_trans, "numpy"])
# 查看生成的代码
import inspect
print(''.join(inspect.findsource(f)[0]))

生成的代码会变成：

def _lambdifygenerated(x):
    return select([isclose(x, 0),True], [1,x], default=nan)

完美替换成了np.isclose。

方法二：自定义SymPy函数并指定映射 #

如果需要更灵活的控制，自定义一个SymPy函数，关联np.isclose：

1. 定义自定义函数：

from sympy import Function, Piecewise, lambdify
from sympy.abc import x
import numpy as np

class CloseEq(Function):
    @classmethod
    def eval(cls, a, b):
        # 可选：常量直接判断，做符号层面简化
        if isinstance(a, (float, int)) and isinstance(b, (float, int)):
            return np.isclose(a, b)

# 定义lambdify时的映射规则
custom_modules = {
    "CloseEq": np.isclose,
    "numpy": np
}

2. 用自定义函数构造Piecewise：

expr = Piecewise(
    (1, CloseEq(x, 0.0)),
    (x, True)
)

f = lambdify(x, expr, modules=custom_modules)
# 查看生成的代码
import inspect
print(''.join(inspect.findsource(f)[0]))

生成的代码同样会使用isclose，同时满足Piecewise对布尔条件的要求。

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