Search code examples
pythonsympypiecewiselambdify

lambdification of sympy piecewise function evaluates every expression


I'm lambdifying a sympy piecewise function trying to do something like this:

f = Piecewise((1,(p > -1e-10) & (p < 1e-10)), (1/p, True))
g = lambdify(p,f,"numpy")

While

>>> f.subs(p,0)
1 

I get

>>> g(0)
/usr/lib/python2.7/dist-packages/numpy/__init__.py:1: RuntimeWarning: divide by zero encountered in true_divide
  """ 
array(1.0)

It seems, that (the lambdified ?)-Piecewise evaluates all expressions before returning the one with the true condition. Is there a way around this?


Solution

  • The NumPy code printer used by lambdify translates Piecewise to

    numpy.select(conditions, expressions, default=numpy.nan)
    

    This means that the array expressions is computed in its entirety before numpy.select selects one element of that array. Some ways to get around it are:

    1) Change the backend to math (or mpmath, or anything other than numpy), which will cause Piecewise to be translated as a nested if statement.

    g = lambdify(p, f, "math")
    g(0)  # 1, no warnings
    

    2) Rewrite the formula in terms of Max/Min/Abs/sign, which can express some piecewise functions and lambdify easily. This isn't always possible but in your case,

    f = 0.5 * (sign(p + 1e-10) + sign(p - 1e-10)) / Max(1e-10, Abs(p)) + 0.5 * (sign(p + 1e-10) - sign(p - 1e-10))
    

    does the job. The trick is that 0.5 * (sign(p + 1e-10) + sign(p - 1e-10)) is sign(p) when p is not too close to 0, and is 0 when it is. Similarly, 0.5 * (sign(p + 1e-10) - sign(p - 1e-10)) is 1 if p is not too close to 0 and is 0 when it is. These factors cause the formula to switch from one mode to the other, and Max in the denominator avoids the division by zero error in any case.

    3) Suppress Runtime Warnings