如何使用SymPy加速长函数的符号导数？

Question

我正在用 Python 编写一个程序来使用自由 ICI 方法求解薛定谔方程（好吧，现在是 SICI 方法......但它会变成自由 ICI）。如果这听起来不熟悉，那是因为关于该主题的信息很少，而且绝对没有示例代码可供使用。

该过程涉及迭代地得出偏微分方程的解。在这样做时，需要执行许多符号导数。问题是，随着程序运行，需要微分的函数会越来越大，以至于到第五次迭代时，计算符号导数需要大量时间。

我需要加快速度，因为我希望能够实现至少 30 次迭代，而且我希望在我退休之前做到这一点。

我已经完成并删除了不必要的重复计算（或者至少是我所知道的那些），这有很大帮助。除此之外，我完全不知道如何加快速度。

这是包含计算导数的inf_integrate函数的代码（该函数只是复合辛普森的方法，因为它比使用 SymPy 的方法快得多integrate，并且不会因振荡函数而引发错误）：

from sympy import *


def inf_integrate(fun, n, a, b):
    f = lambdify(r, fun)
    h = (b-a)/n
    XI0 = f(a) + f(b)
    XI1 = 0
    XI2 = 0

    for i in range(1, n):
        X = a + i*h

        if i % 2 == 0:
            XI2 = XI2 + f(X)
        else:
            XI1 = XI1 + f(X)

    XI = h*(XI0 + 2*XI2 + 4*XI1)/3

    return XI


r = symbols('r')

def H(fun):
    return (-1/2)*diff(fun, r, 2) - (1/r)*diff(fun, r) - (1/r)*fun

E1 = symbols('E1')
low = 10**(-5)
high = 40
n = 5000

g = Lambda(r, r)


psi0 = Lambda(r, exp(-1.5*r))

I1 = inf_integrate(4*pi*(r**2)*psi0(r)*H(psi0(r)), n, low, high)
I2 = inf_integrate(4*pi*(r**2)*psi0(r)*psi0(r), n, low, high)

E0 = I1/I2
print(E0)

for x in range(10):

    f1 = Lambda(r, psi0(r))
    f2 = Lambda(r, g(r)*(H(psi0(r)) - E0*psi0(r)))
    Hf1 = Lambda(r, H(f1(r)))
    Hf2 = Lambda(r, H(f2(r)))

    H11 = inf_integrate(4*pi*(r**2)*f1(r)*Hf1(r), n, low, high)
    H12 = inf_integrate(4*pi*(r**2)*f1(r)*Hf2(r), n, low, high)
    H21 = inf_integrate(4*pi*(r**2)*f2(r)*Hf1(r), n, low, high)
    H22 = inf_integrate(4*pi*(r**2)*f2(r)*Hf2(r), n, low, high)

    S11 = inf_integrate(4*pi*(r**2)*f1(r)*f1(r), n, low, high)
    S12 = inf_integrate(4*pi*(r**2)*f1(r)*f2(r), n, low, high)
    S21 = S12
    S22 = inf_integrate(4*pi*(r**2)*f2(r)*f2(r), n, low, high)

    eqn = Lambda(E1, (H11 - E1*S11)*(H22 - E1*S22) - (H12 - E1*S12)*(H21 - E1*S21))

    roots = solve(eqn(E1), E1)

    E0 = roots[0]

    C = -(H11 - E0*S11)/(H12 - E0*S12)

    psi0 = Lambda(r, f1(r) + C*f2(r))

    print(E0)

该程序正在运行并收敛到预期的结果，但速度太慢了。非常感谢任何有关加快速度的帮助。

Answer 1

您可以在这里做几件事：

1. 如果您分析您的代码，您会注意到您将大部分时间花在了集成函数上inf_integrate，主要是因为您使用了手动 Python 循环。这可以通过将参数转换为向量化函数并使用 SciPy 的集成例程（编译后速度很快）来修改。

2. 当您使用嵌套的符号表达式时，检查偶尔的显式简化是否有助于控制爆炸的复杂性可能是值得的。这似乎是这里的情况。

3. Lamda不需要您定义的所有函数。您可以使用表达式简化工作。我还没有检查这是否真的影响了运行时，但它肯定有助于下一步（因为 SymEngine 还没有Lambda）。

4. 使用 SymEngine 而不是 SymPy。SymPy（截至目前）纯粹是基于 Python 的，因此速度很慢。SymEngine 是其正在编译的核心，可以显着更快。它几乎具有您需要的所有功能。

5. 每一步，您都在求解一个不会改变其性质的方程：它始终是相同的二次方程，只是系数发生了变化。一般来说，通过一次性解决这个问题，您可以节省大量时间，尤其是 SymPy 不必处理复杂的系数。

综上所述，我得出以下结论：

from symengine import *
import sympy
from scipy.integrate import trapz
import numpy as np

r, E1 = symbols('r, E1')
H11, H22, H12, H21 = symbols("H11, H22, H12, H21")
S11, S22, S12, S21 = symbols("S11, S22, S12, S21")
low = 1e-5
high = 40
n = 5000

quadratic_expression = (H11-E1*S11)*(H22-E1*S22)-(H12-E1*S12)*(H21-E1*S21)
general_solution = sympify( sympy.solve(quadratic_expression,E1)[0] )
def solve_quadratic(**kwargs):
    return general_solution.subs(kwargs)

sampling_points = np.linspace(low,high,n)
def inf_integrate(fun):
    f = lambdify([r],[fun])
    values = f(sampling_points)
    return trapz(values,sampling_points)

def H(fun):
    return -fun.diff(r,2)/2 - fun.diff(r)/r - fun/r

psi0 = exp(-3*r/2)
I1 = inf_integrate(4*pi*(r**2)*psi0*H(psi0))
I2 = inf_integrate(4*pi*(r**2)*psi0**2)
E0 = I1/I2
print(E0)

for x in range(30):
    f1 = psi0
    f2 = r * (H(psi0)-E0*psi0)
    Hf1 = H(f1)
    Hf2 = H(f2)

    H11 = inf_integrate( 4*pi*(r**2)*f1*Hf1 )
    H12 = inf_integrate( 4*pi*(r**2)*f1*Hf2 )
    H21 = inf_integrate( 4*pi*(r**2)*f2*Hf1 )
    H22 = inf_integrate( 4*pi*(r**2)*f2*Hf2 )

    S11 = inf_integrate( 4*pi*(r**2)*f1**2 )
    S12 = inf_integrate( 4*pi*(r**2)*f1*f2 )
    S21 = S12
    S22 = inf_integrate( 4*pi*(r**2)*f2**2 )

    E0 = solve_quadratic(
            H11=H11, H22=H22, H12=H12, H21=H21,
            S11=S11, S22=S22, S12=S12, S21=S21,
        )
    print(E0)

    C = -( H11 - E0*S11 )/( H12 - E0*S12 )
    psi0 = (f1 + C*f2).simplify()

在我的机器上，这会在几秒钟内收敛到 ?½。