Преобразование выражений sympy в функции массивов numpy

У меня есть система ODE, написанная в 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]

Я хотел бы преобразовать это в векторнозначную функцию, приняв массив 1D numpy значения x, массив 1D numpy значений k, возвращая массив 1D numpy уравнений, оцененных в этих точках. Подпись будет выглядеть примерно так:

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

Причина, по которой мне нужно что-то вроде этого, - дать этой функции scipy.integrate.odeint, поэтому она должна быть быстрой.

Попытка 1: subs

Конечно, я могу написать обертку вокруг 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])

Но это очень медленно, особенно для моей реальной системы, которая намного больше и выполняется много раз. Мне нужно скомпилировать эту операцию с кодом C.

Попытка 2: theano

Я могу приблизиться к интеграции sympy с 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]))))

Это компилирует каждое выражение, но вся эта упаковка и распаковка входов и выходов замедляет его. Функция, возвращаемая theano_function, принимает массивы numpy в качестве аргументов, но для каждого символа требуется один массив, а не один элемент для каждого символа. Это то же поведение для functify и ufunctify. Мне не нужно поведение в эфире; Мне нужно, чтобы он интерпретировал каждый элемент массива как другой символ.

Попытка 3: Отложенный вектор

Если я использую DeferredVector, я могу сделать функцию, которая принимает массивы numpy, но я не могу ее скомпилировать с кодом C или вернуть массив numpy без его упаковки.

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])

Использование DeferredVector Мне не нужно распаковывать входы, но мне все равно нужно упаковывать выходы. Кроме того, я могу использовать lambdify, но версии ufuncify и theano_function погибают, поэтому не создается быстрый код C.

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

Ответ 1

Вы можете использовать функцию sympy lambdify. Например,

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 печатается при запуске script.

Это намного быстрее (обратите внимание на изменение в единицах результатов синхронизации):

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