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

Question

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

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

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

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

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

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

您可以在这里做几件事：

如果分析您的代码，您会发现大部分时间花在集成函数 inf_integrate 上，主要是因为您使用的是手动 Python 循环。这可以通过将参数转换为矢量化函数并使用 SciPy 的集成例程（已编译因此速度很快）来修改。
当您使用嵌套的符号表达式时，可能值得检查一下偶尔的显式简化是否有助于控制爆炸式的复杂性。这似乎是这里的情况。
您定义的所有Lamda函数都不需要。您可以使用表达式简化工作。我还没有检查这是否真的会影响运行时，但它肯定有助于下一步（因为 SymEngine 还没有 Lambda）。
使用 SymEngine 而不是 SymPy。 SymPy（截至目前）纯粹是 Python-based，因此很慢。 SymEngine 是其正在制作的编译核心，并且可以相当快。它几乎具有您需要的所有功能。
每走一步，您都会求解一个性质不变的方程：它始终是相同的二次方程，只是系数发生变化。一般来说，通过一次解决这个问题，您可以节省很多时间，特别是 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()

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

Answer 2

Wrzlprmft 的回答很棒。我已经进行了清理，并将笨重的集成功能与 SymPy 的集成功能进行了交换。这不适用于我的原始代码，但在 Wrzlprmft 的 corrections/additions 之后可以完美运行。该程序有点慢（仍然比我原来的程序快几个数量级），但它不再有限制精度的错误。这是最终代码：

from symengine import *
from sympy import *
import sympy

r, E1 = symbols('r, E1')
H11, H22, H12, H21 = symbols("H11, H22, H12, H21")
S11, S22, S12, S21 = symbols("S11, S22, S12, S21")
low = 0
high = oo
n = 100000

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)


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


psi0 = exp(-3*r/2)
I1 = N(integrate(4*pi*(r**2)*psi0*H(psi0), (r, low, high)))
I2 = N(integrate(4*pi*(r**2)*psi0**2, (r, low, high)))
E0 = I1/I2
print(E0)

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

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

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

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

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

How to speed up symbolic derivatives of long functions using SymPy?

python

optimization

sympy

integral

derivative