Python版Velocity-Verlet轨道积分算法性能优化求助

问题背景 #

当前在简单对数势场下使用Velocity-Verlet积分器开展恒星轨道积分计算，初始版本单条轨道计算耗时约25秒，完成1000条轨道积分总耗时约7小时；经初步优化后单条轨道耗时降至9秒，仍无法满足千级、万级轨道的批量计算需求，现征集纯Python环境下的性能优化方案。

现有实现基础信息 #

• 核心算法：采用Velocity-Verlet积分器，核心逻辑封装在代码的verlet_integration()函数中。
• 物理模型：使用简单对数势（Simple Logarithmic Potential, SLP），固定参数v_0=1，参数q仅取0.7、0.9两个值，势场对应的加速度计算逻辑封装在apply_forces()函数中。
• 计算参数规则：
  • 所有轨道初始位置固定x=0、总能量E=0，初始x方向速度vx通过calc_vx()函数求解得到
  • 为保证计算精度，积分时间步长不得大于1E-4
  • 单条轨道积分时间范围为t=0至t=200
• 计算目标：覆盖(y, vy)构成的全部允许相空间（即calc_vx()计算过程中不会出现负数开平方的参数区间），最低要求完成1000条轨道计算，理想状态支持10000条轨道批量计算。
• 约束条件：不得更换编程语言，仅可基于现有Python实现做优化。

现有可运行完整代码 #

import numpy as np


def calc_vx(y, vy, q):
    """
    Calculate starting value of x velocity
    """
    vx2 = -np.log((y / q) ** 2) - vy ** 2
    return np.sqrt(vx2)


def apply_forces(x, y, q):
    """
    Apply forces to determine the accelerations
    """
    Fx = -x / (y ** 2 / q ** 2 + x ** 2)
    Fy = -y / (q ** 2 * x ** 2 + y ** 2)
    return Fx, Fy


def verlet_integration(start, 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):
        x_vec, v_vec, a_vec = vals[i]
        new_x_vec = x_vec + dt * (v_vec + 0.5 * a_vec * dt)
        new_a_vec = apply_forces(*new_x_vec, q)
        new_v_vec = v_vec + 0.5 * (a_vec + new_a_vec) * dt
        vals[i + 1] = new_x_vec, new_v_vec, new_a_vec

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


def integration(y, vy, dt, t0, t1, q):
    # calculate the starting values
    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)  # bereken het aantal benodigde stappen

    e = verlet_integration(start, dt, steps, q)  # integreer
    return e

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

内容的提问来源于stack exchange，提问作者Viktor VN