Python and Physics: Runge-Kutta Method

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 i*k(x)*np.exp(i*k(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.0*mass/(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 = 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.
#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/(i*k(y)))/(2*np.exp(-i*k(y)*y))
    return (np.abs(bb)/np.abs(aa))**2
def transmission(x,y):
    aa = (psi1 + psi2/(i*k(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()

Visual Output for this code. Note how the wavefunction changes as it enters the potential barrier.

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/(i*k(y)))/(2*np.exp(-i*k(y)*y))
    return (np.abs(bb)/np.abs(aa))**2
def transmission(x,y):
    aa = (psi1 + psi2/(i*k(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.