How do I use sympy.lambdify with Max function to substitute numpy.maximum instead of numpy.amax?

Question

I'm trying to lambdify big analytic expression with sp.Max(x, 0) inside. I want to use numpy to vectorize my calculations, so x is going to be an array. I need element-wise maximum values of x and 0. Still, sympy changes sp.Max to np.amax by default. It finds maximum along the axis, that's not what I need. "modules" keyword in lambdify doesn't work as I expect. I've tried:

import numpy as np
import sympy as sp

arr = np.array([1, 2, 3])
expr = sp.sin(x) + sp.Max(x, 0)
f = sp.lambdify(x, expr, modules=[{'Max': np.maximum}, 'numpy'])  # docs say, priority of modules matters
help(f)

It gives:

Help on function _lambdifygenerated:
_lambdifygenerated(x)
    Created with lambdify. Signature:

    func(x)

    Expression:

    sin(x) + Max(0, x)

    Source code:

    def _lambdifygenerated(x):
        return (sin(x) + amax((0,x)))


    Imported modules:

sp.Max changed to amax for some reason.

If 'numpy' is not included into 'modules' list, it simply skips all other functions. I've also tried to swap dict and 'numpy' in list, but it haven't helped. Please, clarify, what's wrong? Is it a bug in sympy?

Answer 1

When using lambdify to create numpy functions intended to work vectorized, there often are subtle problems, especially when variables (x) and constants (0) are mixed.

In this case, sp.max supposes all of its possible many parameters being single values. np.amax gets the maximum of one flattened array. np.maximum gets the element-wise maximum of two arrays. The problem here is that the constant 0 doesn't automatically get expanded to a numpy array.

My workaround would be to replace sp.max with a custom function based on sp.Piecewise. Note that you would need a separate function if there would be more than 2 arguments to sp.max.

import numpy as np
import sympy as sp
from sympy.abc import x

def sympy_max2(a, b):
    return sp.Piecewise((b, a < b), (a, True))

arr = np.array([11, 22, 33, -1, -2])
expr = sp.sin(x) + sympy_max2(0, x)
f = sp.lambdify(x, expr, modules=['numpy'])

print(f(arr)) # [10.00000979 21.99114869 33.99991186 -0.84147098 -0.90929743]

Answer 2

To use the np.maximum function instead of np.amax, I found that specifying the amax method instead of the Max works. The np.maximum function also requires some tweaking, to receive the arguments that are used for the amax function.

import numpy as np
import sympy as sp

arr = np.array([11, 22, 33, -1, -2])
expr = sp.sin(x) + sp.Max(x, 0)
def custom_amax(x,**kwargs): 
    return np.maximum(x[0],x[1])
f = sp.lambdify(x, expr, modules=[{'amax': custom_amax}, 'numpy']) 
f(arr) # [10.00000979, 21.99114869, 33.99991186, -0.84147098, -0.90929743]

Answer 3

In the current version of SymPy I get return (sin(x) + amax((0,x), axis=0)) in the signature. Is this what you want?