Scipy 的solve_bvp 和耦合微分方程的性能問題

Question

我在嘗試在 Python 3.8.3 中實現下面的耦合微分方程（也稱為單模耦合方程）時遇到問題。 至於求解器，我使用的是 Scipy 的 function scipy.integrate.solve_bvp ，可以在此處閱讀其文檔。 我想求解復域中的方程，對於不同的傳播軸值（ z ）和不同的 beta 值（ beta_analysis ）。

問題在於，與使用函數bvp4c 、 bvpinit和bvpset的 Matlab 中的等效實現相比，它非常慢（無法管理）。 評估兩次執行的前幾次迭代，它們返回相同的結果，除了生成的網格在 Scipy 的情況下要大得多。 網格有時甚至會飽和到最大值。

下面顯示了要求解的方程以及邊界條件 function。

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

Moreover, I couldn't find a way to pass k , z and beta_zf as arguments of the function coupling_equation , given that the fun argument of the solve_bpv function must be a callable with the parameters (x, y) . 我的方法是定義一些全局變量，但如果有更好的解決方案，我也將不勝感激。

我試圖編碼的分析 function 是：

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中沒有執行循環時發生（只是 beta 的中間值的solve_bvp function）。 相反，Matlab 設法在大約 6 分鍾內返回了整個問題的解決方案。 如果我將最后一次迭代的結果作為initial_guess傳遞（ analysis function 中的注釋行，則網格溢出的速度會更快，並且不可能進行多次迭代。

Answer 1

基於半隨機輸入，我們可以看到有時會達到max_mesh 。 這意味着可以使用相當大a z_mesh和 arrays 調用coupling_equation方程。 問題在於， coupling_equation包含一個緩慢的純 Python 循環，該循環在 arrays 的每一列上進行迭代。 您可以使用Numpy vectorization 大大加快計算速度。 這是一個實現：

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

與原始實現相比，此代碼提供了類似的 output 與半隨機輸入，但在我的機器上大約快 20 倍。

此外，我不知道max_mesh是否會因您的輸入而變大，即使這是正常/有意的。 為了進一步減少執行時間，減小max_mesh的值可能是有意義的。