在分段中对 sympy 相等条件进行羔羊化时使用 np.isclose

Question

考虑以下分段表达式：

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

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

用 numpy 对该表达式进行 Lambdizing 会得到以下结果：

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)

但是，此函数设计为与数值求解器一起使用，因此我有小的浮点差异，可能会在分段条件下给出错误的结果。 是否可以让

lambdify

使用

np.isclose

代替

np.equal

？

我尝试将其添加到

lambdify

的

modules

参数中，但我不知道该将什么作为键，因为我不是替换函数而是替换关系。例如，这不起作用：

modules={"==": np.isclose}

。

我也尝试过创建一个

implemented_function

：

my_eq = implemented_function('my_eq', lambda x, y: np.isclose(x, y))
expr = Piecewise(
    (1, my_eq(x, 0.0)), 
    (x, True)
)

但是在羔羊化时会给出

TypeError: Second argument must be a Boolean, not my_eq

，如果 Piecewise 期望布尔关系，这是可以理解的。

Answer 1

实现目标的一种方法是创建自定义 numpy 打印机，重写

_print_Relational

方法以通知打印机使用

isequal

，并将打印机交给

lambdify

:

class MyNumPyPrinter(NumPyPrinter):
    def _print_Relational(self, expr):
        "Relational printer for Equality and Unequality"
        op = {
            '==' :'isclose',
            '!=' :'not_equal',
            '<'  :'less',
            '<=' :'less_equal',
            '>'  :'greater',
            '>=' :'greater_equal',
        }
        if expr.rel_op in op:
            lhs = self._print(expr.lhs)
            rhs = self._print(expr.rhs)
            return '{op}({lhs}, {rhs})'.format(op=self._module_format(self._module + '.'+op[expr.rel_op]),
                                               lhs=lhs, rhs=rhs)
        return super()._print_Relational(expr)
    
f = lambdify(x, expr, printer=MyNumPyPrinter)
f(1e-12)
# out: array(1.)