Scipy 的solve_bvp 和耦合微分方程的性能问题

Question

我在尝试在 Python 3.8.3 中实现下面的耦合微分方程（也称为单模耦合方程）时遇到问题。至于求解器，我使用的是 Scipy 的函数scipy.integrate.solve_bvp，其文档可以阅读here。我想求解复域中的方程，针对不同的传播轴值 (z) 和不同的 beta 值 (beta_analysis)。

问题在于，与在 Matlab 中使用函数 bvp4c、bvpinit 和 bvpset 的等效实现相比，它非常慢（无法管理）。评估两次执行的前几次迭代，它们返回相同的结果，除了生成的网格在 Scipy 的情况下要大得多。网格有时甚至会饱和到最大值。

下面显示了要求解的方程以及边界条件函数。

import h5py
import numpy as np
from scipy import integrate
    
def coupling_equation(z_mesh, a):
    ka_z = k    # Global
    z_a = z     # Global
    a_p = np.empty_like(a).astype(complex)

    for idx, z_i in enumerate(z_mesh): 
        beta_zf_i = np.interp(z_i, z_a, beta_zf)    # Get beta at the desired point of the mesh
        ka_z_i = np.interp(z_i, z_a, ka_z)          # Get ka at the desired point of the mesh

        coupling_matrix = np.empty((2, 2), complex)
        coupling_matrix[0] = [-1j * beta_zf_i, ka_z_i]
        coupling_matrix[1] = [ka_z_i, 1j * beta_zf_i]

        a_p[:, idx] = np.matmul(coupling_matrix, a[:, idx])    # Solve the coupling matrix

    return a_p

def boundary_conditions(a_a, a_b):
    return np.hstack(((a_a[0]-1), a_b[1]))

此外，鉴于solve_bpv 函数的fun 参数必须是一个可使用参数(x, y) 调用。我的方法是定义一些全局变量，但如果有更好的解决方案，我也将不胜感激。

我要编写的分析函数是：

def analysis(k, z, beta_analysis, max_mesh):
    s11_analysis = np.empty_like(beta_analysis, dtype=complex)
    s21_analysis = np.empty_like(beta_analysis, dtype=complex)
    
    initial_mesh = np.linspace(z[0], z[-1], 10)    # Initial mesh of 10 samples along L
    mesh = initial_mesh
    
    # a_init must be complex in order to solve the problem in a complex domain
    a_init = np.vstack((np.ones(np.size(initial_mesh)).astype(complex), 
                        np.zeros(np.size(initial_mesh)).astype(complex)))
    
    for idx, beta in enumerate(beta_analysis):
        print(f"Iteration {idx}: beta_analysis = {beta}")
        global beta_zf 
        beta_zf = beta * np.ones(len(z))    # Global variable so as to use it in coupling_equation(x, y)
        
        a = integrate.solve_bvp(fun=coupling_equation, 
                                bc=boundary_conditions, 
                                x=mesh, 
                                y=a_init, 
                                max_nodes=max_mesh,
                                verbose=1)
#         mesh = a.x      # Mesh for the next iteration
#         a_init = a.y    # Initial guess for the next iteration, corresponding to the current solution
        s11_analysis[idx] = a.y[1][0]
        s21_analysis[idx] = a.y[0][-1]
    return s11_analysis, s21_analysis

我怀疑问题与传递给不同迭代的初始猜测有关（请参阅analysis 函数中循环内的注释行）。我尝试将迭代的解决方案设置为以下的初始猜测（这必须减少求解器所需的时间），但它甚至更慢，我不明白。也许我错过了什么，因为这是我第一次尝试求解微分方程。

用于执行的参数如下：

f2 = h5py.File(r'path/to/file', 'r')
k = np.array(f2['k']).squeeze()
z = np.array(f2['z']).squeeze()
f2.close()

analysis_points = 501
max_mesh = 1e6 

beta_0 = 3e2; 
beta_low = 0;       # Lower value of the frequency for the analysis
beta_up = beta_0;   # Upper value of the frequency for the analysis
beta_analysis = np.linspace(beta_low, beta_up, analysis_points);

s11_analysis, s21_analysis = analysis(k, z, beta_analysis, max_mesh)

关于如何提高这些功能的性能的任何想法？提前谢谢大家，如果问题没有很好地表达，我很抱歉，我接受任何关于此的建议。

编辑：添加了一些关于性能和问题大小的信息。

• 实际上，我找不到确定调用coupling_equation 次数的关系。这一定是求解器内部操作的问题。我通过打印一行检查了一次迭代中的调用次数，它发生在 133 次（这是最快的一次）。这必须乘以 beta 的迭代次数。对于已分析的，求解器返回以下内容：

11 次迭代求解，节点数 529。 最大相对残差：9.99e-04 最大边界残差：0.00e+00

• a 和 z_mesh 的形状是相关的，因为 z_mesh 是一个向量，其长度与网格大小相对应，求解器每次调用 coupling_equation 时都会重新计算。鉴于a包含z_mesh各点的前进波和后退波的幅度，a的形状为(2, len(z_mesh))。
• 在计算时间方面，我只用 Python 在大约 2 小时内完成了 19 次迭代。在这种情况下，初始迭代速度更快，但随着网格的增长，它们开始花费更多时间，直到网格饱和到最大允许值为止。我认为这是因为输入耦合系数在该点的值，因为当beta_analysis 中没有执行循环时也会发生这种情况（只是solve_bvp 函数用于beta 的中间值）。相反，Matlab 在大约 6 分钟内设法返回了整个问题的解决方案。如果我将最后一次迭代的结果作为initial_guess 传递（analysis 函数中的注释行，则网格溢出的速度会更快，并且不可能进行多次迭代。

Answer 1

基于半随机输入，我们可以看到有时会达到max_mesh。这意味着coupling_equation 可以用相当大的z_mesh 和a 数组调用。问题在于coupling_equation 包含一个慢速纯Python 循环，它在数组的每一列上进行迭代。您可以使用 Numpy 矢量化 大大加快计算速度。这是一个实现：

def coupling_equation_fast(z_mesh, a):
    ka_z = k    # Global
    z_a = z     # Global
    a_p = np.empty(a.shape, dtype=np.complex128)
    beta_zf_i = np.interp(z_mesh, z_a, beta_zf)    # Get beta at the desired point of the mesh
    ka_z_i = np.interp(z_mesh, z_a, ka_z)          # Get ka at the desired point of the mesh
    # Fast manual matrix multiplication
    a_p[0] = (-1j * beta_zf_i) * a[0] + ka_z_i * a[1]
    a_p[1] = ka_z_i * a[0] + (1j * beta_zf_i) * a[1]
    return a_p

与原始实现相比，此代码提供了类似的半随机输入输出，但在我的机器上快了大约 20 倍。

此外，我不知道max_mesh 您的输入是否也很大，即使这是正常/有意的。减少max_mesh 的值以进一步减少执行时间可能是有意义的。