加速sympy-lamodified和向量化函数

Question

我正在使用sympy来生成一些用于数值计算的函数。因此，我把一个表达式转化为可以和numpy数组一起使用的形式。这里有一个例子：

import numpy as np
import sympy as sp

def numpy_function():
    x, y, z = np.mgrid[0:1:40*1j, 0:1:40*1j, 0:1:40*1j]
    T   = (1 - np.cos(2*np.pi*x))*(1 - np.cos(2*np.pi*y))*np.sin(np.pi*z)*0.1
    return T

def sympy_function():
    x, y, z = sp.Symbol("x"), sp.Symbol("y"), sp.Symbol("z")
    T   = (1 - sp.cos(2*sp.pi*x))*(1 - sp.cos(2*sp.pi*y))*sp.sin(sp.pi*z)*0.1
    lambda_function = np.vectorize(sp.lambdify((x, y, z), T, "numpy"))
    x, y, z = np.mgrid[0:1:40*1j, 0:1:40*1j, 0:1:40*1j]
    T = lambda_function(x,y,z)
    return T

在sympy版本和纯numpy版本之间，主要的问题是速度，也就是：

In [3]: timeit test.numpy_function()  
100 loops, best of 3: 11.9 ms per loop

与

In [4]: timeit test.sympy_function()
1 loops, best of 3: 634 ms per loop

相比。

那么，有没有办法让sympy的速度更接近numpy的版本呢？我觉得np.vectorize的速度挺慢的，但我代码中的某些部分没有它就不工作。谢谢你的建议。

编辑：
我找到了为什么需要vectorize函数的原因，也就是：

In [35]: y = np.arange(10)

In [36]: f = sp.lambdify(x,sin(x),"numpy")

In [37]: f(y)
Out[37]: 
array([ 0.        ,  0.84147098,  0.90929743,  0.14112001, -0.7568025 ,
       -0.95892427, -0.2794155 ,  0.6569866 ,  0.98935825,  0.41211849])

这个似乎工作得很好：

In [38]: y = np.arange(10)

In [39]: f = sp.lambdify(x,1,"numpy")

In [40]: f(y)
Out[40]: 1

所以对于像1这样的简单表达式，这个函数不会返回一个数组。有没有办法解决这个问题？这难道不是某种bug，或者至少是设计不一致吗？

Answer 1

在这个情况下，使用 np.vectorize() 就像是在对 x、y 和 z 的第一维进行循环，所以这会导致速度变慢。如果你告诉 lambdify() 使用 NumPy 的函数，那么就不需要 np.vectorize() 了，而这正是你现在所做的。这样使用：

def sympy_function():
    x, y, z = sp.Symbol("x"), sp.Symbol("y"), sp.Symbol("z")
    T   = (1 - sp.cos(2*sp.pi*x))*(1 - sp.cos(2*sp.pi*y))*sp.sin(sp.pi*z)*0.1
    lambda_function = sp.lambdify((x, y, z), T, "numpy")
    x, y, z = np.mgrid[0:1:40*1j, 0:1:40*1j, 0:1:40*1j]
    T = lambda_function(x,y,z)
    return T

性能就会变得相当不错：

In [26]: np.allclose(numpy_function(), sympy_function())
Out[26]: True

In [27]: timeit numpy_function()
100 loops, best of 3: 4.08 ms per loop

In [28]: timeit sympy_function()
100 loops, best of 3: 5.52 ms per loop

Answer 2

lambdify 函数在处理常量时只返回一个值，这是因为它没有使用到 numpy 的函数。这是 lambdify 的工作方式所决定的（具体可以参考 这个链接）。

不过，这通常不是个问题，因为常量在与数组进行任何操作时，会自动调整为合适的形状。换句话说，你可以直接使用常量，它会自动适应。而如果你特意用一个包含相同常量的数组，那样效率就会低很多，因为你会重复计算同样的操作多次。