将sympy表达式转换为numpy数组的函数

Question

我有一个用sympy编写的ODE系统：

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]

我想将这转换为一个向量值函数，接受一个一维numpy数组的x值，一个一维numpy数组的k值，返回一个在这些点上计算的方程的一维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

如果我使用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

- drhagen

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

- Warren Weckesser

---

这确实比我任何一个版本都要快得多。我得到了subs 120毫秒，theano_function 8.7毫秒和lambdify 0.6毫秒。 - drhagen

如果我们能够想出如何在C中完成整个评估，而不是每次迭代都使用Python列表进行打包和解包，那么我们应该能够节省更多。 - drhagen

如果我将lambdafied_odes中的调用转换为使用splatting [f(*np.concatenate([x,k])) for f in funcs]，当状态数量变化时这是必要的，时间会稍微增加到1.4毫秒 - 仍然是最好的。 - drhagen

如果你真正想要速度，你可以尝试使用autowrap或ufuncify。 - asmeurer

据我所知，ufuncify和autowrap都需要对输入参数进行解包（如果您得到一个numpy数组）并重新打包输出值（如果您想要一个numpy数组）。 - drhagen

SymPy中的DeferredVector可以帮助在lambdify中打包和解包数组参数。 - moorepants

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)

您可以像使用 scipy.integrate.ode 的实例一样使用 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 或 this answer of mine），您可能无法获得最佳性能。对于大型微分方程（如您的情况），这种开销变得不相关。

- Wrzlprmft