# https://towardsdatascience.com/quantum-physics-visualization-with-python-35df8b365ff

import matplotlib.pyplot as plt
import numpy as np#Constants

h = 6.626e-34
m = 9.11e-31#Values for L and x

x_list = np.linspace(0,1,100)
L = 1

def 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])

plt.show()