Преобразуйте выражения sympy в функцию numpy

Question

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

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

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:

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: Теано

Я могу приблизиться к интеграции 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 принимает числовые массивы в качестве аргументов, но ему нужен один массив для каждого символа, а не один элемент для каждого символа. Это то же самое поведение для functify а также ufunctify также. Мне не нужно поведение вещания; Мне нужно, чтобы каждый элемент массива интерпретировался как отдельный символ.

Попытка 3: DeferredVector

Если я использую DeferredVector Я могу создать функцию, которая принимает числовые массивы, но я не могу скомпилировать ее в код на Си или вернуть массив пустых, не упаковав его сам.

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

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

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

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

Я написал модуль с именем JiTCODE, который предназначен для таких проблем, как ваша. Он принимает символические выражения, преобразует их в код C, оборачивает вокруг него расширение Python, компилирует его и загружает для использования с scipy.integrate.ode или же scipy.integrate.solve_ivp,

Ваш пример будет выглядеть так:

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)

Вы можете использовать ODE очень похоже на пример scipy.integrate.ode,

Хотя это не требуется для вашего приложения, вы также можете извлечь и использовать скомпилированную функцию:

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

Обратите внимание, что в отличие от ваших спецификаций, k не передается как массив NumPy. Для большинства приложений ODE это не должно быть актуально, поскольку эффективнее жестко задавать параметры управления.

Наконец, обратите внимание, что для этого небольшого примера вы можете не получить наилучшую производительность из-за накладных расходов scipy.integrate.ode или же scipy.integrate.solve_ivp (также см. SciPy Issue #8257 или мой ответ). Для больших дифференциальных уравнений (как у вас) эти издержки становятся неактуальными.