Overload division operator in sympy

42 Views Asked by At

I would like to overload the division operator in sympy to handle division by zero.

Here is the context: I have expressions, that are simplified through sp.simplify. As a consequence, sympy may simplify an expression into one that makes the division operator / appear. Consider for instance those lines of code:

import sympy as sp
x = sp.Symbol('x', real=True)
expr = 5*sp.exp(-1*sp.log(x))
print(expr)

Which prints 5/x. I managed to "override" functions like log, for that I just used the solution in this post and it works:

import math

# protected log function for evaluation
def pl(x):
    if x == 0:
        return x
    if x < 0:
        return math.log(-x)
    else:
        return math.log(x)

expr = sp.log(x)
f = sp.lambdify(x, expr, modules={'log': pl})
print(f(2) == f(-2))

Which prints True as expected. I assume this works because log appears in the expression. For pow, it is a bit more complicated. I had to use the ccode function:

# protected pow function
def pp(x, p):
    if not isinstance(p, int) and x < 0:
        return math.pow(-x, p)
    if p < 0 and x == 0:
        return 0
    return math.pow(x, p)

expr = sp.Pow(x, 5.5)
print(expr, sp.ccode(expr))
f = sp.lambdify(x, expr, modules={'pow': pp})
g = sp.lambdify(x, sp.ccode(expr), modules={'pow': pp})
print(g(-2) == g(2), f(-2), g(-2))

Which prints x**5.5 pow(x, 5.5) and then True (-1.1087038706409865e-13-45.254833995939045j) 45.254833995939045. So for log, exp or pow functions, I managed to override at evaluation time. But for the division operator I didn't. sp.ccode(5/x) and 5/x print the same. I would like 5/x to appear like div(5, x) then I could use the module keyword argument to override at evaluation time. I understand sympy transforms / with a combination of Mul and Pow functions. But then how to make those operators (/, +, -, *) appear like in the ccode of x**5 which gives pow(x, 5)?

0

There are 0 best solutions below