提高 Velocity-Verlet 积分算法的性能

首先要做的是使用探查器来分析代码中的热点。你可以用它来做到这一点。基本报告显示，几乎所有时间都花在函数上： 和 （称为 200_000 次）：cProfile.runx_funcv_funcapply_forces

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   200000    0.496    0.000    0.496    0.000 apply_forces
   199999    0.279    0.000    2.833    0.000 func_verlet
   199999    0.713    0.000    0.713    0.000 x_func
   199999    0.780    0.000    0.780    0.000 v_func
        1    0.000    0.000    0.000    0.000 SLP
   199999    0.222    0.000    0.718    0.000 a_func
        1    0.446    0.446    3.279    3.279 verlet_integration
[...]

快速分析并显示您正在微小数组上调用 Numpy 函数。问题是 Numpy 没有针对处理如此小的数组进行优化（Numpy 对输入进行许多检查，与微小数组的计算时间相比，这些检查成本很高）。x_funcv_func

解决此问题的主要方法是使用在更大的数组上运行的矢量化 Numpy 调用。问题是这个解决方案需要完全重新设计你的代码，以便获得循环，这是问题的根源。for i in range(steps - 1)

另一种解决方案是使用 Numba 以避免 Numpy 开销。这个解决方案要简单得多（尽管如果你的目标是学习 Numpy，你的老师可能不会预料到这一点）。下面是一个示例：

import numpy as np
import numba as nb


@nb.njit
def SLP(x, y, v0, q):
    return 0.5 * v0 ** 2 * np.log(x ** 2 + (y / q) ** 2)


@nb.njit
def calc_vx(x, y, vy, q):
    """
    Calculate starting value of x velocity
    """
    vx2 = -2 * SLP(x, y, 1, q) - vy ** 2
    return np.sqrt(vx2)


@nb.njit
def apply_forces(x, y, v0, q):
    """
    Apply forces to determine the accelerations
    """
    Fx = -(v0 ** 2 * x) / (y ** 2 / q ** 2 + x ** 2)
    Fy = -(v0 ** 2 * y) / (q ** 2 * x ** 2 + y ** 2)
    return np.array([Fx, Fy])


@nb.njit
def x_func(x_vec, v_vec, a_vec, dt):
        return x_vec + dt * (v_vec + 0.5 * a_vec * dt)

@nb.njit
def v_func(v_vec, a_vec, new_a_vec, dt):
    return v_vec + 0.5 * (a_vec + new_a_vec) * dt

@nb.njit
def a_func(x_vec, dt, q):
    x, y = x_vec
    return apply_forces(x, y, 1, q)


# The parameter is a signature of the function that provides the input type to 
# Numba so it can eagerly compile the function and all the dependent functions.
# Please read the Numba documentation for more information about this.
@nb.njit('(float64[:], float64[:], float64[:], float64, float64)')
def func_verlet(x_vec, v_vec, a_vec, dt, q):
    # calculate the new position, velocity and acceleration
    new_x_vec = x_func(x_vec, v_vec, a_vec, dt)
    new_a_vec = a_func(new_x_vec, dt, q)
    new_v_vec = v_func(v_vec, a_vec, new_a_vec, dt)
    out = np.empty((len(new_x_vec), 3))
    out[:,0] = new_x_vec
    out[:,1] = new_v_vec
    out[:,2] = new_a_vec
    return out


def verlet_integration(start, f, dt, steps, q):
    # initialise an array and set the first value to the starting value
    vals = np.zeros((steps, *np.shape(start)))
    vals[0] = start

    # run through all elements and apply the integrator to each value
    for i in range(steps - 1):
        vals[i + 1] = f(*vals[i].T, dt, q)

    # I return vals.T so i can use the following to get arrays for the position, velocity and acceleration
    # ((x, y), (vx, vy), (ax, ay)) = verlet_integration_vec( ... )
    return vals.T


def integration(y, vy, dt, t0, t1, q):
    # calculate the starting values
    x = 0
    vx = calc_vx(x, y, vy, q)
    ax, ay = apply_forces(x, y, 1, q)
    start = [(x, vx, ax), (y, vy, ay)]
    steps = round((t1 - t0) / dt)  # calculate the number of necessary steps

    return verlet_integration(start, func_verlet, dt, steps, q)


((x_, y_), (vx_, vy_), (ax_, ay_)) = integration(0.1, 0.2, 1E-4, 0, 20, 0.7)

请注意，某些函数已被修改，以避免使用 Numba 不支持的函数/运算符。例如，不支持用于解包 Numpy 数组或列表的一元运算符，但无论如何它通常效率很低。*

生成的代码速度提高了 5 倍。

探查器现在显示，由于 for 循环（包括运算符和函数调用）的开销，该函数负责执行的主要部分。由于 lambda，这部分不容易移植到 Numba。我认为，如果您成功地重新设计了这部分以避免 lambda 和解包，它可以提高两倍的速度。事实上，即使使用 Numba，对 2 个项目的数组进行操作也是非常低效的。在标量上操作会使代码的可读性降低一些，但速度要快得多（当然，无论有没有 Numba）。我想代码可以再次提高几倍。verlet_integration*

更新：通过更新的代码，Numba 可以更好地提供帮助，因为主要的性能瓶颈现已修复。这是新的Numba版本：

import numpy as np
import numba as nb


@nb.njit
def calc_vx(y, vy, q):
    vx2 = -np.log((y / q) ** 2) - vy ** 2
    return np.sqrt(vx2)


@nb.njit
def apply_forces(x, y, q):
    Fx = -x / (y ** 2 / q ** 2 + x ** 2)
    Fy = -y / (q ** 2 * x ** 2 + y ** 2)
    return np.array([Fx, Fy])


@nb.njit('(float64[:,:], float64, int_, float64)')
def verlet_integration(start, dt, steps, q):
    vals = np.zeros((steps, 3, 2))
    vals[0] = start

    for i in range(steps - 1):
        x_vec, v_vec, a_vec = vals[i]
        new_x_vec = x_vec + dt * (v_vec + 0.5 * a_vec * dt)
        x, y = new_x_vec
        new_a_vec = apply_forces(x, y, q)
        new_v_vec = v_vec + 0.5 * (a_vec + new_a_vec) * dt
        vals[i + 1, 0] = new_x_vec
        vals[i + 1, 1] = new_v_vec
        vals[i + 1, 2] = new_a_vec

    return vals.T


def integration(y, vy, dt, t0, t1, q):
    vx = calc_vx(y, vy, q)
    ax, ay = apply_forces(0, y, q)
    start = [(0, y), (vx, vy), (ax, ay)]
    steps = round((t1 - t0) / dt)

    e = verlet_integration(np.array(start), dt, steps, q)
    return e


((x, vx, ax), (y, vy, ay)) = integration(0.1, 0.2, 1E-4, 0, 100, 0.7)  # 9.7

这比问题的更新代码快 36 倍。在我的机器上只需要 0.27 秒，而初始代码需要 9.7 秒。

我看到有两件事可以在这里提供帮助。

（1） 考虑使用库中的 ODE 求解器，而不是您手动编写的 Verlet 方法。Runge-Kutta 方法及其变体（例如 Runge-Kutta-Fehlberg）广泛适用，我会先尝试一下。库方法很可能是用 C 代码编写的，因此速度要快得多，而且有人已经解决了错误。这两个方面对于这个问题都很有用。

（2） 如果您需要自己编写 ODE 求解器，那么在 Python 中加快实现速度的一个想法是在轨迹上矢量化，以便您可以使用 Numpy 数组操作。也就是说，创建一个表示所有 1000 或 10,000 条轨迹的数组，然后一步一步地推进所有轨迹。我假设你可以用矩阵形式重写运动方程。

顺便说一句，我的建议是保留你已经制定的解决方案，因为它似乎正在工作，并使用想法（1）或（2）开始一个单独的实现，并使用你的原始实现来验证结果。

在我的基准测试中，内联apply_forces将 @JérômeRichard 的解决方案速度提高了 ~2 倍。

import numpy as np
import numba as nb   #tested with numba 0.55.2

@nb.njit
def calc_vx(y, vy, q):
    vx2 = -np.log((y / q) ** 2) - vy ** 2
    return np.sqrt(vx2)

@nb.njit
def verlet_integration(start, dt, steps, q):
    vals = np.zeros((steps, 3, 2))
    vals[0] = start

    for i in range(steps - 1):
        x_vec, v_vec, a_vec = vals[i]
        x, y = x_vec + dt * (v_vec + 0.5 * a_vec * dt)
        ax, ay =  -x / (y ** 2 / q ** 2 + x ** 2), -y / (q ** 2 * x ** 2 + y ** 2)
        vx, vy = v_vec[0] + 0.5 * (a_vec[0] + ax) * dt, v_vec[1] + 0.5 * (a_vec[1] + ay) * dt

        vals[i + 1, 0] = x, y
        vals[i + 1, 1] = vx, vy
        vals[i + 1, 2] = ax, ay
    return vals.T

@nb.njit
def integration(y, vy, dt, t0, t1, q):
    vx = calc_vx(y, vy, q)
    ax, ay = 0, -y / y**2
    start = np.array([[0, y], [vx, vy], [ax, ay]])
    steps = round((t1 - t0) / dt) 

    return verlet_integration(start, dt, steps, q)

((x, vx, ax), (y, vy, ay)) = integration(0.1, 0.2, 1E-4, 0, 100, 0.7)