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

Question

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

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

我需要加快速度，因为我希望能够实现至少 30 次迭代，并且我希望在我退休之前让它完成。

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

这是包含计算导数的函数的代码(inf_integrate 函数只是复合 Simpson 方法，因为它比使用 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()

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