cn_method2d.py

#!/usr/bin/python3
"""
Cranck-Nicolson method for the Schrödinger equation in 2D.
 
Jacobi iteration can be used to solve for the system of equations that occurs
when using the Cranck-Nicolson method. This is following an
article found here https://arxiv.org/pdf/1409.8340 by Sadovskyy et al.,
where the Cranck-Nicolson method with Jacobi iteration is 
used to solve the Ginzburg-Landau equations, which is similar to 
the Schrödinger equation but contains nonlinear terms and
couplings with the four vector potential.
"""
from time import perf_counter
import numpy as np
import scipy.constants as const
from scipy import sparse
from scipy.sparse import linalg
import matplotlib.pyplot as plt
import matplotlib.animation as animation


def jacobi(inv_diag: sparse.dia_matrix, lower_upper: sparse.dia.dia_matrix,
           b: np.ndarray, n_iter: int = 3):
    """
    Given the inverse diagonals of a matrix A and its lower and upper parts
    without its diagonal, find the solution x for the system Ax = b.

    Reference for Jacobi iteration:
    https://en.wikipedia.org/wiki/Jacobi_method
    """
    x = b.copy()
    for i in range(n_iter):
        # TODO: this can be parallelized
        x = inv_diag.dot((b.T - lower_upper.dot(x.T)).T)
    return x


# Constants
N = 256 # Number of points to use
L = 1e-8 # Extent of simulation (in meters)
S = L*np.linspace(-0.5, 0.5, N)
X, Y = np.meshgrid(S, S)
DX = S[1] - S[0] # Spatial step size
DT = 1e-17 # Time step size

# The wavefunction
k = 0.0 # Modify this value to change the initial momentum of the wavefunction
sigma = 0.056568
psi = np.exp(-((X/L+0.25)/sigma)**2/2.0
             -((Y/L-0.25)/sigma)**2/2.0)*(1.0 + 0.0j) #*np.exp(2.0j*np.pi*k*X/L)
normalize = lambda psi: psi/np.sqrt(np.sum(psi*np.conj(psi)))
psi = normalize(psi)

# The potential
V = 6*1e-18*((X/L)**2 + (Y/L)**2) # Simple Harmonic Oscillator
# V = np.zeros([N]) # Zero potential everywhere (except at boundaries)
# V = 0.9*1e-18*np.array([1.0 if i >= 0.48*N and i < 0.52*N else 0.0 
#                         for i in range(N)]) # Potential Barrier

# Construct H
ALPHA = -const.hbar**2/(2.0*const.m_e)
diag = (ALPHA/DX**2)*np.ones([N])
diags = np.array([diag, -2.0*diag, diag])
kinetic_1d = sparse.spdiags(diags, np.array([1.0, 0.0, -1.0]), N, N)
T = sparse.kronsum(kinetic_1d, kinetic_1d)
U = sparse.diags(V.reshape(N*N), (0))
H = T + U

# Construct A and B matrices,
# where A|psi(t + dt)> = B|psi(t)>
# DT = DT*(1.0 - 1.0j)/np.sqrt(2)
# DT = -1.0j*DT
BETA = 0.5j*DT/const.hbar
I = sparse.identity(N*N)
A = I + BETA*H
B = I - BETA*H
D = sparse.diags(B.diagonal(0), (0))
INV_D = sparse.diags(1.0/B.diagonal(0), (0))
L_PLUS_U = B - D

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
im = ax.imshow(np.angle(X + 1.0j*Y),
               alpha=np.abs(psi)/np.amax(np.abs(psi)),
               extent=(X[0, 1], X[0, -1], Y[0, 0], Y[-1, 0]),
               interpolation='nearest',
               cmap='hsv',
               )
im2 = ax.imshow(V,
                extent=(X[0, 1], X[0, -1], Y[0, 0], Y[-1, 0]),
                interpolation='bilinear', cmap='gray')
ax.set_xlabel('x (m)')
ax.set_ylabel('y (m)')
ax.set_title('Wavefunction')
data = {'psi': psi.reshape([N*N]), 't': 0.0, 'frames': 0}
t0 = perf_counter()

def animation_func(*args):
    """
    Animation function.
    """
    wavefunc = data['psi']
    steps_per_frame = 5
    for i in range(steps_per_frame):
        wavefunc = A.dot(wavefunc)
        # wavefunction = linalg.spsolve(B, wavefunc.T)
        # wavefunc = linalg.gcrotmk(B, wavefunc.T)[0]
        wavefunc = jacobi(INV_D, L_PLUS_U, wavefunc, n_iter=10)
    wavefunc_im = wavefunc.reshape([N, N])
    im.set_data(np.angle(wavefunc_im))
    abs_wavefunc = np.abs(wavefunc_im)
    # print(abs_wavefunc/np.amax(abs_wavefunc))
    im.set_alpha(abs_wavefunc/np.amax(abs_wavefunc))
    data['psi'] = wavefunc
    data['t'] += steps_per_frame*DT
    data['frames'] += 1
    return (im2, im)

ani = animation.FuncAnimation(fig, animation_func, blit=True, 
                              interval=1000.0/60.0)
plt.show()
fps = 1.0/((perf_counter() - t0)/(data["frames"]))
print(f'fps: {np.round(fps, 1)}')

cn_method_nonlinear.py

#!/usr/bin/python3
"""
Cranck-Nicholson method to numerically solve Schrödinger equation
with nonlinear terms in 1D for a single particle in a finite domain.
This code is an extension of exercise 9.8 in chapter 9 of 
Mark Newman's Computational Physics:
http://www-personal.umich.edu/~mejn/cp/exercises.html
"""
import numpy as np
import scipy.constants as const
from scipy import sparse
from scipy.sparse import linalg
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Constants
N = 1025 # Number of points to use
L = 1e-8 # Extent of simulation (in meters)
X = L*np.linspace(-0.5, 0.5, N)
DX = X[1] - X[0] # Spatial step size
DT = 1e-17 # Time step size

# The wavefunction
k = 0.0 # Modify this value to change the wavefunction's initial momentum
sigma = 0.056568
psi = np.exp(-((X/L+0.25)/sigma)**2/2.0)*np.exp(2.0j*np.pi*k*X/L)
# psi += np.exp(-((X/L-0.25)/sigma)**2/2.0)*np.exp(-2.0j*np.pi*k*X/L)
normalize = lambda psi: psi/np.sqrt(np.dot(psi, np.conj(psi))) 
psi = normalize(psi)

# The potential
V = 6*1e-18*(X/L)**2 # Simple Harmonic Oscillator
# V = np.zeros([N]) # Zero potential everywhere (except at boundaries)
# V = 0.9*1e-18*np.array([1.0 if i >= 0.48*N and i < 0.52*N else 0.0 
#                         for i in range(N)]) # Potential Barrier

# Construct H
ALPHA = -const.hbar**2/(2.0*const.m_e)
d0 = -2.0*ALPHA/DX**2*np.ones([N]) + V
d1 = ALPHA/DX**2*np.ones([N])
diags = np.array([d1, d0, d1])
H = sparse.spdiags(diags, [-1, 0, 1], N, N)

# Construct A and B matrices,
# where A|psi(t + dt)> = B|psi(t)>
# DT = DT*(1.0 - 1.0j)/np.sqrt(2)
# DT = -1.0j*DT
BETA = 0.5j*DT/const.hbar
I = sparse.identity(N)
A = I + BETA*H
B = I - BETA*H

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim(X[0], X[-1])
ax.set_ylim(-6.0, 6.0)
ax.set_xlabel('Position (m)')
ax.set_ylabel('Potential Energy (eV)')
data = {'psi': psi, 't': 0.0}
real_plot, = ax.plot(X, np.real(psi), label=r'Re|$\psi(x, t)$|')
imag_plot, = ax.plot(X, np.imag(psi), label=r'Im|$\psi(x, t)$|')
abs_plot, = ax.plot(X, np.abs(psi), color='black', linewidth=1.0, 
                    label=r'|$\psi(x, t)$|')
potential_plot, = ax.plot(X, V/const.electron_volt, color='gray', 
                          label='Potential V(x)', linewidth=2.0)

def animation_func(*args):
    psi = data['psi']
    nonlinear = np.zeros([N, N], np.complex128)
    nonlinear = sparse.spdiags((1.0e-17*BETA)*np.abs(psi)**2, 
                               [0], N, N)
    psi = (A + nonlinear).dot(psi)
    psi = linalg.spsolve(B - nonlinear, psi.T)
    psi_scaled = psi*40.0*(N/1025)
    real_plot.set_ydata(np.real(psi_scaled))
    imag_plot.set_ydata(np.imag(psi_scaled))
    abs_plot.set_ydata(np.abs(psi_scaled))
    data['psi'] = psi
    data['t'] += DT
    return real_plot, imag_plot, abs_plot

ani = animation.FuncAnimation(fig, animation_func, blit=True, 
                              interval=1000.0/60.0)
plt.grid(linestyle='--', linewidth=0.5)
plt.legend()
plt.show()
plt.close()

cn_method_qm.py

#!/usr/bin/python3
"""
Using the Cranck-Nicholson method to numerically solve Schrödinger equation
in 1D for a single particle in a finite domain. This code is an extension
of exercise 9.8 in chapter 9 of Mark Newman's Computational Physics:
http://www-personal.umich.edu/~mejn/cp/exercises.html
"""
import numpy as np
import scipy.constants as const
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Constants
N = 1025 # Number of points to use
L = 1e-8 # Extent of simulation (in meters)
X = L*np.linspace(-0.5, 0.5, N)
DX = X[1] - X[0] # Spatial step size
DT = 1e-17 # Time step size

# The wavefunction
k = 0.0 # Modify this value to change the wavefunction's initial momentum
sigma = 0.056568
psi = np.exp(-((X/L+0.25)/sigma)**2/2.0)*np.exp(2.0j*np.pi*k*X/L)
norm_factor = 1.0/np.sqrt(np.dot(psi, np.conj(psi)))
psi *= norm_factor

# The potential
V = 6*1e-18*(X/L)**2 # Simple Harmonic Oscillator
# V = np.zeros([N]) # Zero potential everywhere (except at boundaries)
# V = 0.9*1e-18*np.array([1.0 if i >= 0.48*N and i < 0.52*N else 0.0 
#                         for i in range(N)]) # Potential Barrier

# Construct A and B matrices,
# where A|psi(t + dt)> = B|psi(t)>
A = np.zeros([N, N], np.complex128)
B = np.zeros([N, N], np.complex128)
K = (DT*1.0j*const.hbar)/(4*const.m_e*DX**2)
J = (DT*1.0j)/(2*const.hbar)
a1 = 1 + 2*K
a2 = -K
b1 = 1 - 2*K
b2 = K
for i in range(N-1):
    A[i][i] = a1 + J*V[i]
    B[i][i] = b1 - J*V[i]
    A[i][i+1] = a2
    A[i+1][i] = a2
    B[i][i+1] = b2
    B[i+1][i] = b2
A[N-1][N-1] = a1 + J*V[N-1]
B[N-1][N-1] = b1 - J*V[N-1]

# Find U, where U(dt)|psi(t)> = |psi(t + dt)>.
U = np.dot(np.linalg.inv(A), B)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim(X[0], X[-1])
ax.set_ylim(-6.0, 6.0)
ax.set_xlabel('Position (m)')
ax.set_ylabel('Potential Energy (eV)')
data = {'psi': psi, 't': 0.0}
real_plot, = ax.plot(X, np.real(psi), label=r'Re|$\psi(x, t)$|')
imag_plot, = ax.plot(X, np.imag(psi), label=r'Im|$\psi(x, t)$|')
abs_plot, = ax.plot(X, np.abs(psi), color='black', linewidth=1.0, 
                    label=r'|$\psi(x, t)$|')
potential_plot, = ax.plot(X, V/const.electron_volt, color='gray', 
                          label='Potential V(x)', linewidth=2.0)

def animation_func(*args):
    psi = data['psi']
    psi = np.dot(U, psi)
    psi_scaled = psi*40.0
    real_plot.set_ydata(np.real(psi_scaled))
    imag_plot.set_ydata(np.imag(psi_scaled))
    abs_plot.set_ydata(np.abs(psi_scaled))
    data['psi'] = psi
    data['t'] += DT
    return real_plot, imag_plot, abs_plot

animation.FuncAnimation(fig, animation_func, blit=True, interval=1.0)
plt.grid(linestyle='--', linewidth=0.5)
plt.legend()
plt.show()
plt.close()
print(f"Total time elapsed: {data['t']}s")

print((U - U.T)/DT)
# psi = data['psi']
# prob_density = psi*psi.conjugate() 
# print(np.sum(prob_density))
# print(np.sum(prob_density[0:N//2]), np.sum(prob_density[N//2:N]))