Sympy: Lambdify Such That Operations On Arrays Always Result In Arrays, Also For Constants?

I need to evaluate the derivative of functions (f') given by the user in many points. The points are in a list (or numpy.array, pandas.Series...). I obtain the expected value when

Solution 1:

lamb1 is a function that returns the constant 1: def lamb1(x): return 1.

lamb2 is a function that returns its argument: def lamb2(x): return x.

So, the output is very well the expected one.

Here is an approach that might work. I changed the test function for f2 to t*t as that was more annoying in my tests (dealing with Pow(t,2)).

import sympy as sp
import numpy as np

f1 = sp.sympify('1')
f2 = sp.sympify('t*t')

defnp_lambdify(varname, func):
    lamb = sp.lambdify(varname, func, modules=['numpy'])
    if func.is_constant():
        returnlambda t: np.full_like(t, lamb(t))
    else:
        returnlambda t: lamb(np.array(t))

lamb1 = np_lambdify('t', f1)
lamb2 = np_lambdify('t', f2)

print(lamb1(1))
print(lamb1([1, 2, 3]))
print(lamb2(2))
print(lamb2([1, 2, 3]))

Outputs:

1
[1 1 1]
4
[1 4 9]

Solution 2:

With isympy/ipython introspection:

In [28]: lamb2??                                                                                 
Signature: lamb2(t)
Docstring:
Created with lambdify. Signature:

func(arg_0)

Expression:

t

Source code:

def _lambdifygenerated(t):
    return (t)

and for the first:

In [29]: lamb1??                                                                                 
Signature: lamb1(t)
Docstring:
Created with lambdify. Signature:

func(arg_0)

Expression:1

Source code:

def _lambdifygenerated(t):
    return (1)

So one returns the input argument; the other returns just the constant, regardless of the input. lambdify does a rather simple lexical translation from sympy to numpy Python.

edit

Putting your functions in a sp.Matrix:

In [55]: lamb3 = lambdify('t',Matrix([f1,f2]))                                                   

In [56]: lamb3??                                                                                 
...
def _lambdifygenerated(t):
    return (array([[1], [t]]))
...

In [57]: lamb3(np.arange(3))                                                                     
Out[57]: 
array([[1],
       [array([0, 1, 2])]], dtype=object)

So this returns a numpy array; but because of the mix of shapes the result is object dtype, not 2d.

We can see this with a direct array generation:

In [53]: np.array([[1],[1,2,3]])                                                                 
Out[53]: array([list([1]), list([1, 2, 3])], dtype=object)

In [54]: np.array([np.ones(3,int),[1,2,3]])                                                      
Out[54]: 
array([[1, 1, 1],
       [1, 2, 3]])

Neither sympy nor the np.array attempts to 'broadcast' that constant. There are numpy constructs that will do that, such as multiplication and addition, but this simple sympy function and lambdify don't.

edit

frompyfunc is a way of passing an array (or arrays) to a function that only works with scalar inputs. While lamb2 works with an array input, you aren't happy with the lamb1 case, or presumably lamb3.

In [60]: np.frompyfunc(lamb1,1,1)([1,2,3])                                                       
Out[60]: array([1, 1, 1], dtype=object)

In [61]: np.frompyfunc(lamb2,1,1)([1,2,3])                                                       
Out[61]: array([1, 2, 3], dtype=object)

This [61] is slower than simply lamb2([1,2,3]) since it effectively iterates.

In [62]: np.frompyfunc(lamb3,1,1)([1,2,3])                                                       
Out[62]: 
array([array([[1],
       [1]]), array([[1],
       [2]]),
       array([[1],
       [3]])], dtype=object)

In this Matrix case the result is an array of arrays. But since shapes match they can be combined into one array (in various ways):

In [66]: np.concatenate(_62, axis=1)                                                             
Out[66]: 
array([[1, 1, 1],
       [1, 2, 3]])

Solution 3:

Usually it isn't actually a problem for lambdify to return a constant, because NumPy's broadcasting semantics will automatically treat a constant as an array of that constant of the appropriate shape.

If it is a problem, you can use a wrapper like

def broadcast(fun):return lambda *x: numpy.broadcast_arrays(fun(*x), *x)[0]

(this is taken from https://github.com/sympy/sympy/issues/5642, which has more discussion on this issue).

Note that using broadcast is better than full_like as in JohanC's answer, because broadcasted constant arrays do not actually take up more memory, whereas full_like will copy the constant in memory to make the array.

Solution 4:

I often use the trick t * 0 + 1 to create a zero-vector the same length as my input, but then add 1 to each of its elements. It works with NumPy; check if it works with Sympy!

Solution 5:

I never use lambdify so I can't be too critical of how it is working. But it appears that you will need to fool it by giving it an expression that doesn't simplify to a scalar which, when evaluated with numbers will reduce to the desired value:

>>>import numpy as np>>>lambdify('t','(1+t)*t-t**2-t+42','numpy')(np.array([1,2,3]))
array([42, 42, 42])