Quantum Physics Visualization With Python | by Syed Sadat Nazrul | Towards Data Science

On this blog, I have decided to review some college level quantum chemistry for deriving electron orbitals. The additional fun part is that, we are going to visualize wave functions and electron probabilities.. using Python!
In 1926, Erwin Schrodinger advanced the famous wave equation that relates the energy of a system to its wave properties. Because its application to the hydrogen atom is rather complicated, we shall first use wave equation to solve the particle-in-a-box. The Schrodinger Wave equation expressing in 1D is
Now, to simplify our equation, we assume a Particle In A Box.
The particle-in-a-box problem does not correspond to any real chemical system. Its usefulness in our context is that it illustrates several quantum mechanical features. The potential energy at the barrier is set to infinity (i.e. the particle cannot escape) and the potential energy inside the barrier is set to 0. Under these conditions, classical mechanics predicts that the particle has an equal probability of being in any part of the box and the kinetic energy of the particle is allowed to have any value. Taking this assumption into consideration, we get different equations for the particle’s energy at the barrier and inside the box.
At the barrier, V is infinite and the hence, the particle does not exist:
Inside the box, V is zero and hence the wave can have any finite value:
Inside the box, we can rearrange the equation as follows:
As we can see above, the wave function would be such that if differentiated twice, should give the same function multiplied by E. The sine function possesses this behavior.
Now, we need to evaluate the values for the constants, α and A. For α, we use the wave equations at the barriers, where the wave functions equal 0.
Now plugging in the value for α:
We can determine the value of A by requiring the wave function to be normalized. This is because, the particle must exist somewhere in the box. Hence, the sum of the probability of finding the particle in the box is 1:
Plugging in the values, the final wave and energy equations are:
Visualizing the Energy and wave functions using Python:
import matplotlib.pyplot as pltimport numpy as np#Constantsh = 6.626e-34m = 9.11e-31#Values for L and xx_list = np.linspace(0,1,100)L = 1def psi(n,L,x):    return np.sqrt(2/L)*np.sin(n*np.pi*x/L)def psi_2(n,L,x):    return np.square(psi(n,L,x))plt.figure(figsize=(15,10))plt.suptitle("Wave Functions", fontsize=18)for n in range(1,4):    #Empty lists for energy and psi wave    psi_2_list = []    psi_list = []    for x in x_list:        psi_2_list.append(psi_2(n,L,x))        psi_list.append(psi(n,L,x))    plt.subplot(3,2,2*n-1)    plt.plot(x_list, psi_list)    plt.xlabel("L", fontsize=13)    plt.ylabel("Ψ", fontsize=13)    plt.xticks(np.arange(0, 1, step=0.5))    plt.title("n="+str(n), fontsize=16)    plt.grid()    plt.subplot(3,2,2*n)    plt.plot(x_list, psi_2_list)    plt.xlabel("L", fontsize=13)    plt.ylabel("Ψ*Ψ", fontsize=13)    plt.xticks(np.arange(0, 1, step=0.5))    plt.title("n="+str(n), fontsize=16)    plt.grid()plt.tight_layout(rect=[0, 0.03, 1, 0.95])
Notice how there are regions where both Ψ and Ψ* Ψ are zero on the same regions. This is known as the node. Energy levels of orbitals are not continuous. They exist at discreet levels, as represented by the location of the nodes. Also, as the n value increases, the density of the wave inside the box also increases.
Now to get the wave equation with respect to quantum numbers, it needs to be in the following 3D format:
Now, separating of variables depends on the type of atom and is too complex to cover in this blog post. Instead, we will just write the solution directly for plotting. For the following, we will be using the functions of R and Y for Hydrogen atom without deriving them.
First, let’s look at the 1s orbital:
The 1s wave function reveals that the probability of an electron appearing decreases exponentially as we move away from the nucleus. It also reveals a spherical shape.
import matplotlib.pyplot as pltimport numpy as np#Probability of 1sdef prob_1s(x,y,z):    r=np.sqrt(np.square(x)+np.square(y)+np.square(z))    #Remember.. probability is psi squared!    return np.square(np.exp(-r)/np.sqrt(np.pi))#Random coordinatesx=np.linspace(0,1,30)y=np.linspace(0,1,30)z=np.linspace(0,1,30)elements = []probability = []for ix in x:    for iy in y:        for iz in z:            #Serialize into 1D object            elements.append(str((ix,iy,iz)))            probability.append(prob_1s(ix,iy,iz))            #Ensure sum of probability is 1probability = probability/sum(probability)#Getting electron coordinates based on probabiliycoord = np.random.choice(elements, size=100000, replace=True, p=probability)elem_mat = [i.split(',') for i in coord]elem_mat = np.matrix(elem_mat)x_coords = [float(i.item()[1:]) for i in elem_mat[:,0]] y_coords = [float(i.item()) for i in elem_mat[:,1]] z_coords = [float(i.item()[0:-1]) for i in elem_mat[:,2]]#Plottingfig = plt.figure(figsize=(10,10))ax = fig.add_subplot(111, projection='3d')ax.scatter(x_coords, y_coords, z_coords, alpha=0.05, s=2)ax.set_title("Hydrogen 1s density")plt.show()
A little hard to see from the electron density plot above. However, you can see that it has a spherical shape. As we move further away from the center, the density decreases. Generally, the cut off point is when the probability of the electron appearing is at 99%.
The same density plots can also be derived for the other spdf orbitals. Hopefully, this blog has motivated you to have fun with Quantum Physics and Python programming! For more on python based visualizations, check out my blog: Python based Plotting with Matplotlib