One of the most commons math problems that I stumbled across in grad school were Ordinary Differential Equations, otherwise known as ODEs and one of the challenges that had me stumped for a while was, how do I solve ODEs in Python? Thankfully, I was able to stumble across two methods, the Runge-Kutta method and SciPy’s built-in function.
Runge-Kutta Method The Runge-Kutta method was a numerical approximation for ODE’s, developed by Carl Runge and Wilhelm Kutta. By using four slope values within an interval, that do not necessarily fall on the actual solution, and averaging out the slopes, one can get a pretty nice approximation of the solution. For a more in detail explanation of the Runge-Kutta method and its variations, I highly suggest researching the history, derivation and applications using your favorite textbook/website. Now for this example, we will be focusing on the Fourth Order Runge-Kutta Method to help us solve the 1D scattering problem.
Coding To start off our code, we are going to import some packages that will help us with the math and the visualization.
import cmath #To help us out with the complex square root import numpy as np #For the arrays import matplotlib.pyplot as plt #Visualization From here, we then start defining our initial parameters for the equations.
mass = 1.0 #Mass, one for simplicity hbar = 1.0 #HBar, one for simplicity v0 = 2.0 #Initial potential value alpha = 0.5 #Value for our potential equation E = 3.0 #Energy i = 1.0j #Defining imaginary number x = 10.0 #Initial x-value xf = -10.0 #Final x-value h = -.001 #Step value xaxis = np.array([], float) #Empty array to fill with out x-values psi = np.array([], complex) #Empty array to fill with the values for the initial equation we are trying to solve, defined array as complex to fill with complex numbers psiprime = np.array([], complex) #Empty array to fill with the values for the first derivative equation we are trying to solve, defined array as complex to fill with complex numbers Once we have our initial values, we then start to work on our functions that define the equations we are about to use. The main equation we have is k(x), which is reworked version of Schrödinger's equation to solve for the variable k, as well as our Ψ equations, which will be defined by psione(x) and psitwo(x). To further explore the equation, MIT OCW has free lectures that anyone can access for free.
def v(x): #Potential equation we will be using for this example return v0/2.0 * (1.0 + np.tanh(x/alpha)) def k(x): #Reworked Schrödinger's equation to solve for k return cmath.sqrt((2_mass/(hbar**2))_(E - v(x))) def psione(x): #PSI, wavefunction equation return np.exp(i*k(x)_x) def psitwo(x): #Derivative of the psione equation return i_k(x)_np.exp(i_k(x)*x) Now this is where we get to the good part. First, we need to define an array that contains our initial condition wavefunctions.
r = np.array([psione(x), psitwo(x)]) #Array with wavefunctions, usually this is where our initial condition equations go. With these equations set in our array, we can iterate both of these equations through the Runge-Kutta method, which will be defined below, and have them give us the approximate solutions for the equation we defined for psione(x) and psitwo(x). But before we reach the main part of the equation, we need to define one more important function.
def deriv(r,x): return np.array([r[1],-(2.0*mass/(hbar**2) * (E - v(x))*r[0])], complex) #The double star, **, is for exponents The deriv function is where the outputs from the Runge-Kutta get passed through, this function grabs our values from the array r and then pushes it through these conditions. For the first value that gets returned, it’s quite simple, our x value will just be inputted into the second equation of the array. The second value, however, will be going through a different treatment. This time, the x value will be going through another iteration of Schrodinger’s equation, one that takes into consideration the wave function psione(x).
#While loop to iterate through the Runge-Kutta. This particular version, the Fourth Order, will have four slope values that help approximate then next slope value, from k1 to k2, k2 to k3, and k3 to k4. #This loop also appends that values, starting with the initial values, to the empty arrays that we've initialized earlier. while (x >= xf ): xaxis = np.append(xaxis, x) psi = np.append(psi, r[0]) psiprime = np.append(psiprime, r[1]) k1 = h_deriv(r,x) k2 = h_deriv(r+k1/2,x+h/2) k3 = h_deriv(r+k2/2,x+h/2) k4 = h_deriv(r+k3,x+h) r += (k1+2_k2+2_k3+k4)/6 x += h #The += in this line, and the line above, is the same thing as telling the code to x = x + h, which updates x, using the previous x with the addition of the step value. Here, the loop pretty much goes over the whole process of the Runge-Kutta. By using approximations of the slopes, as defined by the k values, each k value helps approximate the next slope, bringing us one step closer to solving for f(x). Furthermore, after getting each of our slopes, we then obtain a weighted average and update our array with these new values to get the ready for the next iteration. This process will continue for our defined range in the x-axis, which will end up giving us the necessary values for plotting our soon to be solved ODE.
Overview The Runge-Kutta Method can be easily adapted to plenty of other equations, most of the time we just have to adjust the deriv function, and our initial condition equations. Other examples include the pendulum ODEs and planetary motion ODEs. Down below, one can now find the full code, along with extra steps, such as the functions to solve for the reflection and transmission values, as well as how to plot our values.
import cmath #To help us out with the complex square root import numpy as np #For the arrays import matplotlib.pyplot as plt #Visualization mass = 1.0 #Mass, one for simplicity hbar = 1.0 #HBar, one for simplicity v0 = 2.0 #Initial potential value alpha = 0.5 #Value for our potential equation E = 3.0 #Energy i = 1.0j #Defining imaginary number x = 10.0 #Initial x-value xf = -10.0 #Final x-value h = -.001 #Step value xaxis = np.array([], float) #Empty array to fill with out x-values psi = np.array([], complex) #Empty array to fill with the values for the initial equation we are trying to solve, defined array as complex to fill with complex numbers psiprime = np.array([], complex) #Empty array to fill with the values for the first derivative equation we are trying to solve, defined array as complex to fill with complex numbers
def v(x): #Potential equation we will be using for this example return v0/2.0 * (1.0 + np.tanh(x/alpha)) def k(x): #Reworked Schrödinger's equation to solve for k return cmath.sqrt((2_mass/(hbar**2))_(E - v(x))) def psione(x): #PSI, wavefunction equation return np.exp(i_k(x)x) def psitwo(x): #Derivative of the psione equation return ik(x)np.exp(ik(x)x) r = np.array([psione(x), psitwo(x)]) #Array with wavefunctions, usually this is where our initial condition equations go. def deriv(r,x): return np.array([r[1],-(2.0mass/(hbar**2) * (E - v(x))r[0])], complex) #The double star, **, is for exponents #While loop to iterate through the Runge-Kutta. This particular version, the Fourth Order, will have four slope values that help approximate then next slope value, from k1 to k2, k2 to k3, and k3 to k4. #This loop also appends that values, starting with the initial values, to the empty arrays that we've initialized earlier. while (x >= xf ): xaxis = np.append(xaxis, x) psi = np.append(psi, r[0]) psiprime = np.append(psiprime, r[1]) k1 = hderiv(r,x) k2 = hderiv(r+k1/2,x+h/2) k3 = hderiv(r+k2/2,x+h/2) k4 = h_deriv(r+k3,x+h) r += (k1+2_k2+2_k3+k4)/6 x += h #The += in this line, and the line above, is the same thing as telling the code to x = x + h, which updates x, using the previous x with the addition of the step value. #Grabbing the last values of the arrays and redefining our x-axis psi1 = psi[20000]; psi2 = psiprime[20000]; x = 10; xf = -10 def reflection(x, y): aa = (psi1 + psi2/(i_k(y)))/(2_np.exp(i_k(y)y)) bb = (psi1 - psi2/(ik(y)))/(2_np.exp(-i*k(y)_y)) return (np.abs(bb)/np.abs(aa))*2 def transmission(x,y): aa = (psi1 + psi2/(ik(y)))/(2.0*np.exp(i_k(y)*y)) return k(x)/k(y) * 1.0/(np.abs(aa))**2 print('reflection = ',reflection(x,xf)) print('transmission = ', transmission(x,xf)) print('r + t = ', reflection(x,xf) + transmission(x,xf)) #Outputs for the print command #reflection = 0.007625630800891285 #transmission = (0.9923743691991354+0j) #r + t = (1.0000000000000266+0j) #Ideally, r + t should give us one, a bit stumped if the precision that's present in Python can lead to the small discrepancy, without considering formatting the answer to a set amount of decimal values. #Plotting the graphs side by side, including the imaginary values. fig, ax = plt.subplots(1,2, figsize = (15,5)) ax[0].plot(xaxis, psi.real, xaxis, psi.imag, xaxis, v(xaxis)) ax[1].plot(xaxis, psiprime.real, xaxis, psiprime.imag, xaxis, v(xaxis)) plt.show()
For those interested in the built-in SciPy version of this code, here you go.
import cmath import numpy as np import matplotlib.pyplot as plt from scipy.integrate import odeint, solve_ivp E = 3; m = 1; h = 1; alpha = .5; v0=2; i = 1.0j; xi = 10; xf = -10 def v(x): return v0/2.0 * (1.0 + np.tanh(x/alpha)) def k(x): return cmath.sqrt((2_m/(h**2))_(E - v(x))) def psione(x): return np.exp(i*k(x)_x) def psitwo(x): return i_k(x)_np.exp(i_k(x)_x) def deriv(x, y): return [y[1], -(2.0_m/(h**2.0) * (E - v(x))*y[0])]
solve_ivp is a built in rk45step solver
values = solve_ivp(deriv, [10, -10], [psione(xi), psitwo(xi)], first_step = .001, max_step = .001) psi1 = values.y[0,20000]; psi2 = values.y[1,20000]; x = 10; xf = -10 def reflection(x, y): aa = (psi1 + psi2/(i_k(y)))/(2_np.exp(i_k(y)y)) bb = (psi1 - psi2/(ik(y)))/(2_np.exp(-i*k(y)_y)) return (np.abs(bb)/np.abs(aa))*2 def transmission(x,y): aa = (psi1 + psi2/(ik(y)))/(2.0*np.exp(i_k(y)*y)) return k(x)/k(y) * 1.0/(np.abs(aa))**2 print('reflection = ',reflection(x,xf)) print('transmission = ', transmission(x,xf)) print('r + t = ', reflection(x,xf) + transmission(x,xf)) fig, ax = plt.subplots(1,2, figsize = (15,5)) ax[0].plot(values.t, values.y[0].real, values.t, values.y[0].imag, values.t, v(values.t)) ax[1].plot(values.t, values.y[1].real, values.t, values.y[1].imag, values.t, v(values.t)) plt.show() Notes A bit of a disclaimer, most of this code was adapted from the coursework from my computational class, which focused on using FORTRAN90 instead of Python, so code may not be efficient since I started teaching myself Python by transferring code from FORTRAN. Plus, a nice shoutout to MIT OCW for giving me a small refresher on the equations and methods used in this code.