Implementing a joint differential equation and eigenvalue solver

Question

I've written a Python script to solve a system of coupled differential equations and then use the results to solve an eigenvalue problem. The full problem is:

\$\frac{d^2}{dr^2} K(r) = \frac{K(r)}{r^2}(K^2(r) -1) -4 g^2v^2h^2(r)K(r)\$

\$\frac{d^2}{dr^2} h(r) + \frac{2}{r} \frac{d}{dr} h(r) = \frac{2}{r^2}K^2(r)h(r) -4 v^2λ(1-h^2(r))h(r)\$

\$h(0) = 0; K(0) = 1; h(∞) = 1; K(∞) = 0\$

---

\$\frac{d}{dr}F_L(r) + \frac{1+K(r)}{r} F_L(r) +f_g\frac{v}{\sqrt{2}} h(r)F_R(r) = ε G_L(r)\$

\$\frac{d}{dr}G_L(r) + \frac{1-h(r)}{r} G_L(r) +f_g\frac{v}{\sqrt{2}} h(r)G_R(r) = -ε F_L(r)\$

\$\frac{d}{dr}F_(r) + \frac{2}{r} F_R(r) +f_g\frac{v}{\sqrt{2}} h(r)F_L(r) = -ε G_R(r)\$

\$\frac{d}{dr}G_R(r) + f_g\frac{v}{\sqrt{2}} h(r)G_L(r) = ε F_R(r)\$

\$F_L(∞) = G_L(∞) = F_R(∞) = G_R(∞) = 0; F_L(0) = F_R(0); G_L(0), G_r(0) {finite}\$

In the first set of equations, \$r\$ is a variable, \$K(r)\$ and \$h(r)\$, are functions of \$r\$ and \$v\$, \$g\$ and \$λ\$ are parameters. After solving the first set, I want to put \$K(r)\$ and \$h(r)\$ in the second set which is an eigensystem with the eigenvalue \$ε\$, where \$f_g\$ is a parameter.

The variable \$r\$ can vary from zero to infinity.

The parameters \$\{ g, v, λ, f_g\} \$ are going to be set by hand.

The eigenvalue \$ε\$ along with the eigenfunction $$\begin{pmatrix}G_L \\\ F_L \\\ G_R \\\ F_R\end{pmatrix}$$ should be found by solving the second set of equations, which can be written as follows:

$$\begin{bmatrix} \begin{array}{cc|cc|cc|cc} 0 && \frac{d}{dr} + \frac{1 + k}{r} && 0 && f_g\frac{v}{\sqrt{2}}h \\ \hline -\frac{d}{dr} + \frac{k-1}{r} && 0 && -f_g\frac{v}{\sqrt{2}}h && 0 \\ \hline 0 && -f_g\frac{v}{\sqrt{2}}h && 0 && -\frac{d}{dr} - \frac{2}{r} \\ \hline f_g \frac{v}{\sqrt{2}}h && 0 && \frac{d}{dr} && 0 \end{array} \end{bmatrix} \begin{bmatrix}G_L \\\ F_L \\\ G_R \\\ F_R\end{bmatrix} = ε \begin{bmatrix}G_L \\\ F_L \\\ G_R \\\ F_R\end{bmatrix} $$

And This is my approach for solving this problem:

 import numpy as np from scipy.integrate import solve_bvp from scipy.interpolate import interp1d from scipy.linalg import eig import matplotlib.pyplot as plt # Define the parameters (these will be set by hand) g = 1.0 v = 1.0 lambda_val = 1.0 fq = 1.0 # ----------------------------------------------------------------------- # Part 1: Solve the first set of equations for K(r) and h(r) # ----------------------------------------------------------------------- def first_set_ode(r, y): """Defines the system of ODEs for the first set.""" K = y[0] dK_dr = y[1] h = y[2] dh_dr = y[3] ddK_ddr = (K / r**2) * (K**2 - 1) - 4 * g**2 * v**2 * h**2 * K ddh_ddr = (2 / r**2) * K**2 * h - (2 / r) * dh_dr - 4 * v**2 * lambda_val * (1 - h**2) * h return np.array([dK_dr, ddK_ddr, dh_dr, ddh_ddr]) def first_set_bc(ya, yb): """Defines the boundary conditions for the first set.""" return np.array([ya[2], ya[0] - 1, yb[2] - 1, yb[0]]) # h(0)=0, K(0)=1, h(inf)=1, K(inf)=0 # Define the spatial grid for the first set (adjust as needed) r_first = np.linspace(1e-3, 10, 100) # Start slightly above 0 to avoid division by zero # Initial guess for the solution (can be improved) y_guess = np.zeros((4, r_first.size)) y_guess[0, :] = 1 # Guess for K(r) y_guess[2, :] = 1 # Guess for h(r) # Solve the BVPs sol_first = solve_bvp(first_set_ode, first_set_bc, r_first, y_guess, max_nodes=100000, tol=1e-9, verbose=2) if not sol_first.success: raise RuntimeError("solve_bvp for the first set failed!") K_r = sol_first.y[0] h_r = sol_first.y[2] # Plot K(r) and h(r) plt.figure(figsize=(10, 5)) plt.plot(sol_first.x, K_r, label='K(r)') plt.plot(sol_first.x, h_r, label='h(r)') plt.xlabel('r') plt.ylabel('Functions') plt.title('Solutions of the First Set of Equations') plt.legend() plt.grid(True) plt.show() # Interpolate the solutions to be used in the second set K_interp = interp1d(sol_first.x, K_r, bounds_error=False, fill_value=(K_r[0], K_r[-1])) h_interp = interp1d(sol_first.x, h_r, bounds_error=False, fill_value=(h_r[0], h_r[-1])) # ----------------------------------------------------------------------- # Part 2: Set up and solve the eigenvalue problem (second set) # ----------------------------------------------------------------------- r_second = np.linspace(1e-3, 10, 500) n = len(r_second) dr = r_second[1] - r_second[0] matrix = np.zeros((4 * n, 4 * n), dtype=complex) for i in range(n): r = r_second[i] K_val = K_interp(r) h_val = h_interp(r) # Row 0 (corresponds to GL[i]) if i > 0: matrix[i, n + i - 1] += 1 / dr matrix[i, n + i] += -1 / dr + (K_val - 1) / r matrix[i, 3 * n + i] += fq * v / np.sqrt(2) * h_val # Row 1 (corresponds to FL[i]) if i > 0: matrix[n + i, i - 1] += 1 / dr matrix[n + i, i] += -1 / dr + (1 - K_val) / r matrix[n + i, 2 * n + i] += -fq * v / np.sqrt(2) * h_val # Row 2 (corresponds to GR[i]) matrix[2 * n + i, n + i] += -fq * v / np.sqrt(2) * h_val if i > 0: matrix[2 * n + i, 3 * n + i - 1] += 1 / dr matrix[2 * n + i, 3 * n + i] += -1 / dr - 2 / r # Row 3 (corresponds to FR[i]) matrix[3 * n + i, i] += fq * v / np.sqrt(2) * h_val if i < n - 1: matrix[3 * n + i, 2 * n + i + 1] += 1 / dr matrix[3 * n + i, 2 * n + i] += -1 / dr # Boundary Condition: FL(0) = FR(0) => FL[0] - FR[0] = 0 bc_row = np.zeros(4 * n) bc_row[n] = 1 bc_row[3 * n] = -1 matrix[0, :] = bc_row # Solve the eigenvalue problem eigenvalues, eigenvectors = eig(matrix) # Filter for real and positive eigenvalues tolerance = 1e-9 # Tolerance for imaginary part real_eigenvalues_mask = np.abs(eigenvalues.imag) < tolerance real_eigenvalues = eigenvalues[real_eigenvalues_mask].real positive_real_eigenvalues_mask = real_eigenvalues > 0 positive_real_eigenvalues = real_eigenvalues[positive_real_eigenvalues_mask] # Get the indices of the positive real eigenvalues in the original eigenvalues array positive_real_eigenvalue_indices = np.where((np.abs(eigenvalues.imag) < tolerance) & (eigenvalues.real > 0))[0] positive_real_eigenvectors = eigenvectors[:, positive_real_eigenvalue_indices] # ----------------------------------------------------------------------- # Plotting the energy levels # ----------------------------------------------------------------------- plt.figure(figsize=(8, 6)) if len(positive_real_eigenvalues) > 0: for eigenvalue in positive_real_eigenvalues: plt.axhline(y=eigenvalue, color='r') plt.yticks(positive_real_eigenvalues) plt.ylabel('Energy Levels (Epsilon)') plt.title('Real and Positive Energy Levels') plt.grid(axis='y') else: plt.title('No Real and Positive Energy Levels Found') plt.show() # Print the real and positive eigenvalues print("Real and positive eigenvalues (epsilon):") if len(positive_real_eigenvalues) > 0: for i, eigenvalue in enumerate(positive_real_eigenvalues): print(f"Eigenvalue {i+1}: {eigenvalue}") for i in range(1, positive_real_eigenvectors.shape[1], 5)[:5]: eigenfunction = positive_real_eigenvectors[:, i] GL = eigenfunction[0:n] FL = eigenfunction[n:2*n] GR = eigenfunction[2*n:3*n] FR = eigenfunction[3*n:4*n] plt.figure(figsize=(10, 6)) plt.plot(r_second, GL.real, label='GL') plt.plot(r_second, FL.real, label='FL') plt.plot(r_second, GR.real, label='GR') plt.plot(r_second, FR.real, label='FR') plt.xlabel('r') plt.ylabel('Eigenfunction components') plt.title(f'Eigenfunction for positive real eigenvalue {positive_real_eigenvalues[i]:.4f}') plt.legend() plt.grid(True) plt.show() else: print("No real and positive eigenvalues found.") 

My primary goal here is to verify the correctness of my method and its implementation in code.

Answer 1

My primary goal here is to verify the correctness of my method and its implementation in code

You've done some things well - you have seemingly descriptive variable names and comments.

You need to subdivide your code into functions and make their inputs and outputs well-defined (I show a little of this; they need to be further subdivided); and you need to write tests. Only you are likely to know what tests are meaningful in context here, but one way or the other, unit tests are very important.

For long integer literals like 100000, add digit group separators like 100_000.

Rather than

K = y[0] dK_dr = y[1] h = y[2] dh_dr = y[3] 

use a tuple unpack.

Add PEP484 type hints for more self-descriptive code.

Where possible (especially solve_bvp), pass arguments by keyword.

Assuming that you have enough patience, remove all of your intermediate plt.show and only call it once at the end. Among other reasons, this allows you to browse through the figures and do analysis with them all available rather than in chunks.

When the elements are already arrays, prefer np.stack(()) rather than np.array().

A light refactor looks like:

import numpy as np from scipy.integrate import solve_bvp from scipy.interpolate import interp1d from scipy.linalg import eig import matplotlib.pyplot as plt # Define the parameters (these will be set by hand) G = 1. V = 1. LAMBDA_VAL = 1. FQ = 1. def first_set_ode(r: np.ndarray, y: np.ndarray) -> np.ndarray: """Defines the system of ODEs for the first set.""" K, dK_dr, h, dh_dr = y ddK_ddr = K/r**2 * (K**2 - 1) - 4*G**2 * V**2 * h**2 * K ddh_ddr = 2/r**2 * K**2 * h - 2/r*dh_dr - 4*V**2 * LAMBDA_VAL * (1 - h**2)*h return np.stack((dK_dr, ddK_ddr, dh_dr, ddh_ddr)) def first_set_bc(ya: np.ndarray, yb: np.ndarray) -> np.ndarray: """Defines the boundary conditions for the first set.""" # h(0)=0, K(0)=1, h(inf)=1, K(inf)=0 return np.array((ya[2], ya[0] - 1, yb[2] - 1, yb[0])) def do_part_1() -> tuple[ interp1d, # K interpolator interp1d, # h interpolator ]: """Part 1: Solve the first set of equations for K(r) and h(r)""" # Define the spatial grid for the first set (adjust as needed) r_first = np.linspace(start=1e-3, stop=10, num=100) # Start slightly above 0 to avoid division by zero # Initial guess for the solution (can be improved) y_guess = np.zeros((4, r_first.size)) y_guess[0, :] = 1 # Guess for K(r) y_guess[2, :] = 1 # Guess for h(r) # Solve the BVPs sol_first = solve_bvp( fun=first_set_ode, bc=first_set_bc, x=r_first, y=y_guess, max_nodes=100_000, tol=1e-9, verbose=2, ) if not sol_first.success: raise RuntimeError("solve_bvp for the first set failed!") K_r = sol_first.y[0] h_r = sol_first.y[2] # Plot K(r) and h(r) plt.figure(figsize=(10, 5)) plt.plot(sol_first.x, K_r, label='K(r)') plt.plot(sol_first.x, h_r, label='h(r)') plt.xlabel('r') plt.ylabel('Functions') plt.title('Solutions of the First Set of Equations') plt.legend() plt.grid(True) # Interpolate the solutions to be used in the second set K_interp = interp1d(sol_first.x, K_r, bounds_error=False, fill_value=(K_r[0], K_r[-1])) h_interp = interp1d(sol_first.x, h_r, bounds_error=False, fill_value=(h_r[0], h_r[-1])) return K_interp, h_interp def do_part_2(K_interp: interp1d, h_interp: interp1d) -> tuple[ np.ndarray, # r_second (what does this mean?): 500 int, # length of... something? np.ndarray, # positive real eigenvalues: 2000, np.ndarray, # positive real eigenvectors: 2000,2000 ]: """Part 2: Set up and solve the eigenvalue problem (second set)""" r_second = np.linspace(start=1e-3, stop=10, num=500) n = len(r_second) dr = r_second[1] - r_second[0] matrix = np.zeros((4*n, 4*n), dtype=complex) for i in range(n): r = r_second[i] K_val = K_interp(r) h_val = h_interp(r) # Row 0 (corresponds to GL[i]) if i > 0: matrix[i, n + i - 1] += 1/dr matrix[i, n + i] += -1/dr + (K_val - 1)/r matrix[i, 3*n + i] += FQ * V/np.sqrt(2) * h_val # Row 1 (corresponds to FL[i]) if i > 0: matrix[n + i, i - 1] += 1/dr matrix[n + i, i] += -1/dr + (1 - K_val)/r matrix[n + i, 2*n + i] += -FQ * V/np.sqrt(2) * h_val # Row 2 (corresponds to GR[i]) matrix[2*n + i, n + i] += -FQ * V/np.sqrt(2) * h_val if i > 0: matrix[2*n + i, 3*n + i - 1] += 1/dr matrix[2*n + i, 3*n + i] += -1/dr - 2/r # Row 3 (corresponds to FR[i]) matrix[3*n + i, i] += FQ*V/np.sqrt(2) * h_val if i < n - 1: matrix[3*n + i, 2*n + i + 1] += 1/dr matrix[3*n + i, 2*n + i] += -1/dr # Boundary Condition: FL(0) = FR(0) => FL[0] - FR[0] = 0 bc_row = np.zeros(4*n) bc_row[n] = 1 bc_row[3*n] = -1 matrix[0, :] = bc_row # Solve the eigenvalue problem eigenvalues, eigenvectors = eig(matrix) # Filter for real and positive eigenvalues tolerance = 1e-9 # Tolerance for imaginary part real_eigenvalues_mask = np.abs(eigenvalues.imag) < tolerance real_eigenvalues = eigenvalues[real_eigenvalues_mask].real positive_real_eigenvalues_mask = real_eigenvalues > 0 positive_real_eigenvalues = real_eigenvalues[positive_real_eigenvalues_mask] # Get the indices of the positive real eigenvalues in the original eigenvalues array positive_real_eigenvalue_indices = np.where((np.abs(eigenvalues.imag) < tolerance) & (eigenvalues.real > 0))[0] positive_real_eigenvectors = eigenvectors[:, positive_real_eigenvalue_indices] return r_second, n, positive_real_eigenvalues, positive_real_eigenvectors def plot_energy_levels(positive_real_eigenvalues: np.ndarray) -> None: plt.figure(figsize=(8, 6)) if len(positive_real_eigenvalues) > 0: for eigenvalue in positive_real_eigenvalues: plt.axhline(y=eigenvalue, color='r') plt.yticks(positive_real_eigenvalues) plt.ylabel('Energy Levels (Epsilon)') plt.title('Real and Positive Energy Levels') plt.grid(axis='y') else: plt.title('No Real and Positive Energy Levels Found') def print_eigenvalues( n: int, r_second: np.ndarray, positive_real_eigenvalues: np.ndarray, positive_real_eigenvectors: np.ndarray, ) -> None: # Print the real and positive eigenvalues print("Real and positive eigenvalues (epsilon):") if len(positive_real_eigenvalues) > 0: for i, eigenvalue in enumerate(positive_real_eigenvalues): print(f"Eigenvalue {i+1}: {eigenvalue}") for i in range(1, positive_real_eigenvectors.shape[1], 5)[:5]: eigenfunction = positive_real_eigenvectors[:, i] GL = eigenfunction[0:n] FL = eigenfunction[n:2*n] GR = eigenfunction[2*n:3*n] FR = eigenfunction[3*n:4*n] plt.figure(figsize=(10, 6)) plt.plot(r_second, GL.real, label='GL') plt.plot(r_second, FL.real, label='FL') plt.plot(r_second, GR.real, label='GR') plt.plot(r_second, FR.real, label='FR') plt.xlabel('r') plt.ylabel('Eigenfunction components') plt.title(f'Eigenfunction for positive real eigenvalue {positive_real_eigenvalues[i]:.4f}') plt.legend() plt.grid(True) else: print("No real and positive eigenvalues found.") def main() -> None: K_interp, h_interp = do_part_1() r_second, n, eigenvalues, eigenvectors = do_part_2(K_interp=K_interp, h_interp=h_interp) plot_energy_levels(positive_real_eigenvalues=eigenvalues) print_eigenvalues( n=n, r_second=r_second, positive_real_eigenvalues=eigenvalues, positive_real_eigenvectors=eigenvectors, ) plt.show() if __name__ == '__main__': main() 

 Iteration Max residual Max BC residual Total nodes Nodes added 1 1.51e+00 5.77e-21 100 198 2 1.63e+00 9.29e-28 298 594 3 1.90e+00 1.48e-26 892 1244 4 3.47e+00 1.61e-27 2136 2162 5 7.82e-01 2.16e-29 4298 1855 6 3.58e-02 2.01e-34 6153 157 7 1.66e-03 1.34e-29 6310 284 8 7.02e-05 7.24e-32 6594 506 9 2.75e-06 1.16e-33 7100 758 10 1.04e-07 2.56e-34 7858 866 11 1.25e-08 1.42e-32 8724 615 12 1.57e-09 1.46e-38 9339 112 13 1.00e-09 4.97e-36 9451 0 Solved in 13 iterations, number of nodes 9451. Maximum relative residual: 1.00e-09 Maximum boundary residual: 4.97e-36 Real and positive eigenvalues (epsilon): Eigenvalue 1: 324.0067978329092 Eigenvalue 2: 114.28467066216616 Eigenvalue 3: 105.63759576709633 Eigenvalue 4: 102.994278603273 Eigenvalue 5: 101.82411667273881 Eigenvalue 6: 101.20103102776568 Eigenvalue 7: 100.8293002060852 Eigenvalue 8: 100.58960223537291 Eigenvalue 9: 100.42596948952942 Eigenvalue 10: 100.30925585444952 Eigenvalue 11: 100.2229687638843 Eigenvalue 12: 100.15738379152181 Eigenvalue 13: 100.10635487840631 Eigenvalue 14: 100.0659082890682 Eigenvalue 15: 100.03325019001414 Eigenvalue 16: 100.00654242984497 Eigenvalue 17: 99.98439899171278 Eigenvalue 18: 99.96585374241535 Eigenvalue 19: 99.79487812253377 Eigenvalue 20: 99.95014487967713 Eigenvalue 21: 99.81632759694592 Eigenvalue 22: 99.83643273381301 Eigenvalue 23: 99.93673639256637 Eigenvalue 24: 99.92507442910461 Eigenvalue 25: 99.87238293936727 Eigenvalue 26: 99.85514822925501 Eigenvalue 27: 99.88803669466391 Eigenvalue 28: 99.91404999967207 Eigenvalue 29: 99.90193882153434 Eigenvalue 30: 99.772142779728 Eigenvalue 31: 99.57902311831322 Eigenvalue 32: 99.61013016689884 Eigenvalue 33: 99.74814016895175 Eigenvalue 34: 99.72291797384827 Eigenvalue 35: 99.64006189931119 Eigenvalue 36: 99.66886400545937 Eigenvalue 37: 99.69647647079147 Eigenvalue 38: 99.54682891016543 Eigenvalue 39: 99.5134554405218 Eigenvalue 40: 99.47904973824294 Eigenvalue 41: 99.44343119333162 Eigenvalue 42: 99.3690021611815 Eigenvalue 43: 99.2901945252259 Eigenvalue 44: 99.3303603535861 Eigenvalue 45: 99.406869037291 Eigenvalue 46: 99.24961157971717 Eigenvalue 47: 99.2069880106213 Eigenvalue 48: 99.16476394081408 Eigenvalue 49: 99.11925265846924 Eigenvalue 50: 99.07611248020261 Eigenvalue 51: 99.02655396430973 Eigenvalue 52: 98.98437958335408 Eigenvalue 53: 98.92742220358343 Eigenvalue 54: 98.89156257426785 Eigenvalue 55: 8.658778262657613