Convert sympy expressions to function of numpy arrays

Question

I have a system of ODEs written in sympy:

from sympy.parsing.sympy_parser import parse_expr

xs = symbols('x1 x2')
ks = symbols('k1 k2')
strs = ['-k1 * x1**2 + k2 * x2', 'k1 * x1**2 - k2 * x2']
syms = [parse_expr(item) for item in strs]

I would like to convert this into a vector valued function, accepting a 1D numpy array of the x value, a 1D numpy array of the k values, returning a 1D numpy array of the equations evaluated at those points. The signature would look something like this:

import numpy as np
x = np.array([3.5, 1.5])
k = np.array([4, 2])
xdot = my_odes(x, k)

The reason I want something like this is to give this function to scipy.integrate.odeint, so it needs to be fast.

Attempt 1: subs

Of course, I can write a wrapper around subs:

def my_odes(x, k):
    all_dict = dict(zip(xs, x))
    all_dict.update(dict(zip(ks, k)))
    return np.array([sym.subs(all_dict) for sym in syms])

But this is super slow, especially for my real system which is much bigger and is run many times. I need to compile this operation to C code.

Attempt 2: theano

I can get close with sympy's integration with theano:

from sympy.printing.theanocode import theano_function

f = theano_function(xs + ks, syms)

def my_odes(x, k):
    return np.array(f(*np.concatenate([x,k]))))

This compiles each expression, but all this packing and unpacking of the inputs and outputs slows it back down. The function returned by theano_function accepts numpy arrays as arguments, but it needs one array for each symbol rather than one element for each symbol. This is the same behavior for functify and ufunctify as well. I don't need the broadcast behavior; I need it to interpret each element of the array as a different symbol.

Attempt 3: DeferredVector

If I use DeferredVector I can make a function that accepts numpy arrays, but I can't compile it to C code or return a numpy array without packaging it myself.

import numpy as np
import sympy as sp
from sympy import DeferredVector

x = sp.DeferredVector('x')
k =  sp.DeferredVector('k')
deferred_syms = [s.subs({'x1':x[0], 'x2':x[1], 'k1':k[0], 'k2':k[1]}) for s in syms]
f = [lambdify([x,k], s) for s in deferred_syms]

def my_odes(x, k):
    return np.array([f_i(x, k) for f_i in f])

Using DeferredVector I do not need to unpack the inputs, but I still need to pack the outputs. Also, I can use lambdify, but the ufuncify and theano_function versions perish, so no fast C code is generated.

from sympy.utilities.autowrap import ufuncify
f = [ufuncify([x,k], s) for s in deferred_syms] # error

from sympy.printing.theanocode import theano_function
f = theano_function([x,k], deferred_syms) # error

Answer 1

You can use the sympy function lambdify. For example,

from sympy import symbols, lambdify
from sympy.parsing.sympy_parser import parse_expr
import numpy as np

xs = symbols('x1 x2')
ks = symbols('k1 k2')
strs = ['-k1 * x1**2 + k2 * x2', 'k1 * x1**2 - k2 * x2']
syms = [parse_expr(item) for item in strs]

# Convert each expression in syms to a function with signature f(x1, x2, k1, k2):
funcs = [lambdify(xs + ks, f) for f in syms]


# This is not exactly the same as the `my_odes` in the question.
# `t` is included so this can be used with `scipy.integrate.odeint`.
# The value returned by `sym.subs` is wrapped in a call to `float`
# to ensure that the function returns python floats and not sympy Floats.
def my_odes(x, t, k):
    all_dict = dict(zip(xs, x))
    all_dict.update(dict(zip(ks, k)))
    return np.array([float(sym.subs(all_dict)) for sym in syms])

def lambdified_odes(x, t, k):
    x1, x2 = x
    k1, k2 = k
    xdot = [f(x1, x2, k1, k2) for f in funcs]
    return xdot


if __name__ == "__main__":
    from scipy.integrate import odeint

    k1 = 0.5
    k2 = 1.0
    init = [1.0, 0.0]
    t = np.linspace(0, 1, 6)
    sola = odeint(lambdified_odes, init, t, args=((k1, k2),))
    solb = odeint(my_odes, init, t, args=((k1, k2),))
    print(np.allclose(sola, solb))

True is printed when the script is run.

It is much faster (note the change in units of the timing results):

In [79]: t = np.linspace(0, 10, 1001)

In [80]: %timeit sol = odeint(my_odes, init, t, args=((k1, k2),))
1 loops, best of 3: 239 ms per loop

In [81]: %timeit sol = odeint(lambdified_odes, init, t, args=((k1, k2),))
1000 loops, best of 3: 610 µs per loop

Answer 2

I wrote a module named JiTCODE, which is tailored to problems such as yours. It takes symbolic expressions, converts them to C code, wraps a Python extension around it, compiles this, and loads it for use with scipy.integrate.ode or scipy.integrate.solve_ivp.

Your example would look like this:

from jitcode import y, jitcode
from sympy.parsing.sympy_parser import parse_expr
from sympy import symbols

xs = symbols('x1 x2')
ks = symbols('k1 k2')
strs = ['-k1 * x1**2 + k2 * x2', 'k1 * x1**2 - k2 * x2']
syms = [parse_expr(item) for item in strs]

substitutions = {x_i:y(i) for i,x_i in enumerate(xs)}
f = [sym.subs(substitutions) for sym in syms]

ODE = jitcode(f,control_pars=ks)

You can then use ODE pretty much like an instance of scipy.integrate.ode.

While you would not need this for your application, you can also extract and use the compiled function:

ODE.compile_C()
import numpy as np
x = np.array([3.5, 1.5])
k = np.array([4, 2])
print(ODE.f(0.0,x,*k))

Note that in contrast to your specifications, k is not passed as a NumPy array. For most ODE applications, this should not be relevant, because it is more efficient to hardcode the control parameters.

Finally, note that for this small example, you may not get the best performance due to the overheads of scipy.integrate.ode or scipy.integrate.solve_ivp (also see SciPy Issue #8257 or this answer of mine). For large differential equations (as you have), this overhead becomes irrelevant.