Performance issue with Scipy's solve_bvp and coupled differential equations

I'm facing a problem while trying to implement the coupled differential equation below (also known as single-mode coupling equation) in Python 3.8.3. As for the solver, I am using Scipy's function scipy.integrate.solve_bvp, whose documentation can be read here. I want to solve the equations in the complex domain, for different values of the propagation axis (z) and different values of beta (beta_analysis).

The problem is that it is extremely slow (not manageable) compared with an equivalent implementation in Matlab using the functions bvp4c, bvpinit and bvpset. Evaluating the first few iterations of both executions, they return the same result, except for the resulting mesh which is a lot greater in the case of Scipy. The mesh sometimes even saturates to the maximum value.

The equation to be solved is shown here below, along with the boundary conditions 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). My approach is to define some global variables, but I would appreciate any help on this too if there is a better solution.

The analysis function which I am trying to code is:

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

I suspect that the problem has something to do with the initial guess that is being passed to the different iterations (see commented lines inside the loop in the analysis function). I try to set the solution of an iteration as the initial guess for the following (which must reduce the time needed for the solver), but it is even slower, which I don't understand. Maybe I missed something, because it is my first time trying to solve differential equations.

The parameters used for the execution are the following:

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)

Any ideas on how to improve the performance of these functions? Thank you all in advance, and sorry if the question is not well-formulated, I accept any suggestions about this.

Edit: Added some information about performance and sizing of the problem.

• In practice, I can't find a relation that determines de number of times coupling_equation is called. It must be a matter of the internal operation of the solver. I checked the number of callings in one iteration by printing a line, and it happened in 133 ocasions (this was one of the fastests). This must be multiplied by the number of iterations of beta. For the analyzed one, the solver returned this:

Solved in 11 iterations, number of nodes 529. Maximum relative residual: 9.99e-04 Maximum boundary residual: 0.00e+00

• The shapes of a and z_mesh are correlated, since z_mesh is a vector whose length corresponds with the size of the mesh, recalculated by the solver each time it calls coupling_equation. Given that a contains the amplitudes of the progressive and regressive waves at each point of z_mesh, the shape of a is (2, len(z_mesh)).
• In terms of computation times, I only managed to achieve 19 iterations in about 2 hours with Python. In this case, the initial iterations were faster, but they start to take more time as their mesh grows, until the point that the mesh saturates to the maximum allowed value. I think this is because of the value of the input coupling coefficients in that point, because it also happens when no loop in beta_analysisis executed (just the solve_bvp function for the intermediate value of beta). Instead, Matlab managed to return a solution for the entire problem in just 6 minutes, aproximately. If I pass the result of the last iteration as initial_guess (commented lines in the analysis function, the mesh overflows even faster and it is impossible to get more than a couple iterations.