# ————————————————————————————————————————————————————————————————————————————————————————————————
# Import Libraries
# ————————————————————————————————————————————————————————————————————————————————————————————————
import cmocean as cmo
import numpy as np
from scipy.optimize import curve_fit
from scipy.ndimage import convolve, generate_binary_structure
from scipy import sparse
from scipy.sparse.linalg import eigsh
from scipy.sparse.linalg import eigs
import scipy.constants as spc
import matplotlib as mpl
import matplotlib.pyplot as plt
from ipywidgets import interactive, fixed
import matplotlib.animation
from matplotlib.animation import FuncAnimation
from IPython.core.display import HTML
from IPython.display import display, Math

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
plt.rcParams.update({
    'mathtext.fontset':'cm', 
    'font.family':'serif', 
    'font.sans-serif':'Times New Roman', 
    'figure.figsize':[6,6], 
    'figure.titlesize':22,
    "figure.dpi":90, 
    'savefig.dpi':90,
    'axes.titlesize':20, 
    'axes.labelsize':18, 
    'axes.titley': 1.0, 
    'axes.titlepad': 5.0,
    'axes.edgecolor':'black', 
    'axes.grid': True, 
    'grid.alpha': .5,
    'xtick.labelsize':14,'ytick.labelsize':14,
    'xtick.major.size':6,'ytick.major.size':6,
    'xtick.major.width':1.25, 'ytick.major.width':1.25, 
    'xtick.direction':'inout','ytick.direction':'inout',
    'xtick.top':False, 'ytick.right':False,
    'legend.title_fontsize':14, 'legend.fontsize':14,
    'legend.borderaxespad': 1, 'legend.borderpad': 0.5,
    'legend.framealpha': 1,
    'legend.handleheight': 0.5, 'legend.handlelength': 2.0, 'legend.handletextpad':0.5,
    'legend.labelspacing': 0.25,
    'legend.fancybox':False,
    'legend.edgecolor': '0',
    'legend.frameon': True,
    'legend.markerscale': 1.25,
    'animation.embed_limit':2**128,
    'animation.html': 'jshtml'})

Making the transformations: \begin{align} x \to \ell x \quad &\mathrm{with}\quad \ell = 1\\ t \to \tau t \quad &\mathrm{with}\quad \tau = \frac{m}{\hbar}\\ \mathcal{V} \to V_0 \mathcal{V} \quad &\mathrm{with}\quad V_0 = \frac{\hbar^2}{m}\\ \end{align}

the Schrodinger equation can be expressed as:

\begin{align} \Psi_{t} = \frac{i}{2}\left(\Psi_{xx} + \Psi_{yy}\right) - i \mathcal{V} \,\Psi \end{align}

In finite differences: \begin{align} \frac{\partial\Psi_{i,j}}{\partial t} = \frac{i}{2(\Delta x)^2}\left({\Psi_{i+1,j} + \Psi_{i-1,j} } + {\Psi_{i,j+1} + \Psi_{i,j-1} - 4 \Psi_{i,j}} \right) - i \mathcal{V}_{i,j} \,\Psi_{i,j} \end{align}

def normalize(f, dx, dy):
    return f/(np.sum(f**2 * dx * dy))

def zero_at_walls(f):
    f[0,:]=0
    f[:,0]=0
    f[-1,:]=0
    f[:,-1]=0
    return f

Nx, Ny, Nt = 64, 64, 1000

dx = 0.1
dy = dx
dt = dx**2/4

x    = np.arange(Nx+1)*dx
y    = np.arange(Ny+1)*dy
t    = np.arange(Nt+1)*dt
psi  = np.zeros((Nt+1, Nx+1, Ny+1)).astype(complex)
X, Y = np.meshgrid(x, y)
psi0 = np.zeros((Nx+1,Ny+1))

psi0[(X-np.median(X))**2 + (Y-np.median(Y))**2 < 1] = 1
psi[0]       = psi0.astype(complex)
psi[0][0,:]  = 0
psi[0][:,0]  = 0
psi[0][-1,:] = 0
psi[0][:,-1] = 0
psi[0]       = psi[0]/(np.sum(np.abs(psi[0])**2 * dx * dy))

def dpsi_dt(Psi):
    ker = generate_binary_structure(2,1).astype(int)
    ker[1,1] = -4
    dpsi_dt = 1j/(2*dx**2)*convolve(Psi, ker, mode = 'wrap')
    return dpsi_dt

for i in range(Nt):
    k1 = dpsi_dt(psi[i])
    k2 = dpsi_dt(psi[i] + k1 * dt/2)
    k3 = dpsi_dt(psi[i] + k2 * dt/2)
    k4 = dpsi_dt(psi[i] + k3 * dt)
    psi[i+1][1:-1,1:-1] = psi[i][1:-1,1:-1] + (dt/6) * (k1 + 2*k2 + 2*k3 + k4)[1:-1,1:-1] 
    psi[i+1] = psi[i+1]/(np.sum(np.abs(psi[i+1])**2 * dx * dy))

%config InlineBackend.figure_format = 'retina'
fig, ax = plt.subplots(figsize = (6,6),facecolor = 'white', constrained_layout = True)
wf = ax.imshow(np.abs(psi[0])**2, cmap = 'seismic', interpolation='gaussian')
ax.axis(False)
def animate(i):
    wf.set_array(np.abs(psi[i])**2)
plt.close()
FPS = 24
N_FRAMES = 144
anim = FuncAnimation(fig, animate, interval = int(1e3/FPS), frames = np.arange(0, Nt, int(Nt/N_FRAMES)), repeat = False)
anim

Once Loop Reflect