qm1d_chebyshev.py

"""
Using the Chebyshev method purely with finite differences
to numerically integrate the Schrodinger equation in 1D.
This is primarily based on the following article:

Tal-Ezer H., Kosloff R. (1984). 
An accurate and efficient scheme for propagating 
the time dependent Schrodinger equation.
J. Chem. Phys. 81(9), 3967-3971.

The article is referenced in these notes:
https://www.quantphys.com/2021/03/chebyshev-polynomial-expansion-of-time.html

Reference for discrete Laplacian stencils:
Fornberg, B. (1988).
Generation of Finite Difference Formulas on Arbitrarily Spaced Grids.
Mathematics of Computation, 51(184), 699-706. 
https://doi.org/10.1090/S0025-5718-1988-0935077-0

"""
from time import perf_counter
import numpy as np
from scipy import sparse
from scipy import special
from scipy.sparse import linalg
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# Constants
N = 256 # Number of points to use
L = 45*(N/256) # Extent of simulation
X = L*np.linspace(-0.5, 0.5, N)
DX = X[1] - X[0] # Spatial step size
DT = 0.1
HBAR = 1.0
M_E = 1.0


sigma = 0.05
k = 4.0
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.sum(psi*np.conj(psi)))
psi = normalize(psi)


V = 20.0*(X/L)**2
R = DT*(np.amax(V) - np.amin(V)
        + np.pi**2/(2.0*DX**2*M_E)
        )/2.0
print(R)
G = np.amin(V)*DT


diag = ((-HBAR**2/(2.0*M_E))/DX**2)*np.ones([N])
# diags = np.array([diag, -2.0*diag, diag])
# T = sparse.spdiags(diags,
#                    np.array([1.0, 0.0, -1.0]),
#                    N, N)
# 4th order discrete Laplacian, psi = 0 at boundaries
diags = np.array([-diag/12.0, 4.0*diag/3.0, -5.0/2.0*diag, 
                  4.0*diag/3.0, -diag/12.0])
T = sparse.spdiags(diags,
                   np.array([2.0, 1.0, 0.0, -1.0, -2.0]),
                   N, N)
H = T + sparse.diags(V, (0))


phi_prev = [sparse.identity(N), (-1.0j*DT/(R*HBAR))*H]
# a = [1.0 if k < 2 else ]
# np.exp(1.0j*(R + G))*
# special.jn(0, R)*
U = special.jn(0, R)*phi_prev[0] + 2.0*special.jn(1, R)*phi_prev[1]
for k in range(2, 200):
    ak = 2.0*special.jn(k, R)
    phi = 2.0*(-1.0j*DT/(R*HBAR))*H @ phi_prev[-1] + phi_prev[-2]
    phi_prev[0], phi_prev[1] = phi_prev[1], phi
    U += ak*phi


fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
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', label=r'|$\psi(x, t)$|')
pot_plot, = ax.plot(X, V*np.amax(np.abs(psi))/np.amax(np.abs(V)), 
                    color='grey', linestyle='--', label='Potential V(x)')
ax.set_xlim(X[0], X[-1])
rnge = np.amax(np.real(psi)) - np.amin(np.real(psi))
ax.set_ylim(-rnge, rnge)
ax.set_xlabel('x (a.u.)')
ax.set_title('Wavefunction')

data = {'psi': psi, 't': 0.0}

def animation_func(*arg):
    for _ in range(3):
        data['psi'] = U @ data['psi']
    real_plot.set_ydata(np.real(data['psi']))
    imag_plot.set_ydata(np.imag(data['psi']))
    abs_plot.set_ydata(np.abs(data['psi']))
    data['t'] += DT
    return real_plot, imag_plot, abs_plot, pot_plot


ani = animation.FuncAnimation(fig, animation_func, blit=True, 
                              interval=1000.0/60.0)
plt.legend()
plt.show()
print(np.sum(np.abs(data['psi'])**2))

qm2d_chebyshev.py

"""
Using the Chebyshev method purely with finite differences
to numerically integrate the Schrodinger equation in 2D.
This is primarily based on the following article:

Tal-Ezer H., Kosloff R. (1984). 
An accurate and efficient scheme for propagating 
the time dependent Schrodinger equation.
J. Chem. Phys. 81(9), 3967-3971.

The article is referenced in these notes:
https://www.quantphys.com/2021/03/chebyshev-polynomial-expansion-of-time.html

Reference for discrete Laplacian stencils:
Fornberg, B. (1988).
Generation of Finite Difference Formulas on Arbitrarily Spaced Grids.
Mathematics of Computation, 51(184), 699-706. 
https://doi.org/10.1090/S0025-5718-1988-0935077-0

"""
from time import perf_counter
import numpy as np
from scipy import sparse
from scipy import special
from scipy.sparse import linalg
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# Constants
N = 256 # Number of points to use
L = 45*(N/256) # Extent of simulation
S = L*np.linspace(-0.5, 0.5, N)
X, Y = np.meshgrid(S, S)
DX = S[1] - S[0] # Spatial step size
DT = 0.05
HBAR = 1.0
M_E = 1.0

# The wavefunction
kx, ky = 0.0, 0.0
# kx, ky = 0.0, -30.0 # Sets the initial momentum of the wavefunction
sigma = 0.056568
bx, by = -0.25, 0.25
# bx, by = -0.0, 0.25
psi = np.exp(-((X/L-bx)/sigma)**2/2.0 -((Y/L-by)/sigma)**2/2.0
            )*(1.0 + 0.0j)*np.exp(2.0j*np.pi*(kx*X/L + ky*Y/L))
normalize = lambda psi: psi/np.sqrt(np.sum(psi*np.conj(psi)))
psi = normalize(psi)


# The potential

# Simple Harmonic Oscillator
V = 20.0*((X/L)**2 + (Y/L)**2)

# Zero potential everywhere (except at boundaries)
# V = np.zeros([N])

# Double Slit
# V = np.zeros([N, N])
# y0, yf = 11*N//20 - 3, 11*N//20 + 3
# V[y0: yf, :] = 40.0 # Barrier
# V[y0: yf, 52*N//128: 56*N//128] = 0.0 # Make the left slit
# V[y0: yf, 72*N//128: 76*N//128] = 0.0 # Make the right slit

R = DT*(np.amax(V) - np.amin(V)
        + np.pi**2/(DX**2*M_E)
        )/2.0
print(R)
G = np.amin(V)*DT


diag = ((-HBAR**2/(2.0*M_E))/DX**2)*np.ones([N])
# 4th order discrete Laplacian, psi = 0 at boundaries
diags = np.array([-diag/12.0, 4.0*diag/3.0, -5.0/2.0*diag, 
                  4.0*diag/3.0, -diag/12.0])
kinetic_1d = sparse.spdiags(diags, np.array([2.0, 1.0, 0.0, -1.0, -2.0]),
                            N, N)
T = sparse.kronsum(kinetic_1d, kinetic_1d)
H = T + sparse.diags(V.reshape(N*N), (0))
I = sparse.identity(N*N)
Z = (-1.0j*DT/(R*HBAR))*H

phi_prev = [I, Z]

# U = special.jn(0, R)*phi_prev[0] + 2.0*special.jn(1, R)*phi_prev[1]
# for k in range(2, 12):
#     ak = 2.0*special.jn(k, R)
#     phi = 2.0*Z @ phi_prev[-1] + phi_prev[-2]
#     phi_prev[0], phi_prev[1] = phi_prev[1], phi
#     U += ak*phi

# plt.imshow(np.abs(U[0:100, 0:100].toarray()))
# plt.show()
# plt.close()

chebychev_coeffs = np.array([2.0*special.jn(k, R) for k in range(100)])
chebychev_coeffs[0] *= 0.5

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


def animation_func(*args):
    """
    Animation function.
    """
    steps_per_frame = 1
    number_of_terms = 40
    for i in range(steps_per_frame):
        psi0, psi1 = data['psi'].copy(), Z @ data['psi']
        psi_prev = [psi0, psi1]
        data['psi'] = (chebychev_coeffs[0]*psi_prev[0]
                        + chebychev_coeffs[1]*psi_prev[1])
        for k in range(2, number_of_terms):
            ak = chebychev_coeffs[k] 
            psi = 2.0*Z @ psi_prev[-1] + psi_prev[-2]
            data['psi'] += ak*psi
            psi_prev[0], psi_prev[1] = psi_prev[1], psi
        # data['psi'] = U @ data['psi']
    # data['psi'] = normalize(data['psi'])
    psi = np.flip(data['psi'].reshape([N, N]), axis=0)
    im.set_data(np.angle(psi))
    abs_wavefunc2 = np.real(psi*np.conj(psi))
    alpha_map = 5.0*abs_wavefunc2/np.amax(abs_wavefunc2)
    im.set_alpha(np.where(alpha_map > 1.0, 1.0, alpha_map))
    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)}')
print(f'|psi(x)| = {np.sum(np.abs(data["psi"])**2)}')

qm2d_chebyshev_fft.py

"""
Using the Chebyshev method to numerically integrate the
Schrodinger equation in 2D with vector potential terms.
The momentum terms are handled with spectral methods.
This is primarily based on the following article:

Tal-Ezer H., Kosloff R. (1984). 
An accurate and efficient scheme for propagating 
the time dependent Schrodinger equation.
J. Chem. Phys. 81(9), 3967-3971.

The article is referenced in these notes:
https://www.quantphys.com/2021/03/chebyshev-polynomial-expansion-of-time.html

"""
from time import perf_counter
import numpy as np
from scipy import special
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# Constants (Hartree Atomic Units)
N = 128 # Number of points to use
L = 45*(N/256) # Extent of simulation
S = L*np.linspace(-0.5, 0.5, N)
X, Y = np.meshgrid(S, S)
DX = S[1] - S[0] # Spatial step size
f = np.fft.fftfreq(N)
P_1D = np.pi*f*N/L # Momenta in 1D
P_SHIFT = np.fft.fftshift(P_1D)
PX, PY = np.meshgrid(P_1D, P_1D) # Momenta in the x and y directions
P2 = PX**2 + PY**2
DT = 0.3
HBAR = 1.0
M_E = 1.0 # Mass of the electron
Q_E = 1.0 # Charge of the electron
C = 137.036 # Speed of light


# The wavefunction
kx, ky = 0.0, -0.0 # Sets the initial momentum of the wavefunction
sigma = 0.056568
bx, by = -0.0, 0.25 # Initial position
psi = np.exp(-((X/L-bx)/sigma)**2/2.0 -((Y/L-by)/sigma)**2/2.0
            )*(1.0 + 0.0j)*np.exp(2.0j*np.pi*(kx*X/L + ky*Y/L))
normalize = lambda psi: psi/np.sqrt(np.sum(psi*np.conj(psi)))
psi = normalize(psi)

# The potential

# Zero potential everywhere (except at boundaries)
V = np.zeros([N, N])

# The vector potential
AX, AY = C*10.0*(Y/L), -C*10.0*(X/L)
V += Q_E/(2.0*M_E*C**2)*(AX**2 + AY**2)

# Define constants that are intrinsic to the numerical method being used
R = DT*(np.amax(V) - np.amin(V)
        + np.amax(P2)*(2*M_E)
        )/2.0
CHEBYCHEV_COEFFS = np.array([2.0*special.jn(k, R) for k in range(200)])
CHEBYCHEV_COEFFS[0] *= 0.5


def hamiltonian(psi):
    psi_p = np.fft.fftn(psi)
    grad_x_psi = np.fft.ifftn(-1.0j*PX*psi_p)
    grad_y_psi = np.fft.ifftn(-1.0j*PY*psi_p)
    kinetic = (np.fft.ifftn((P2/(2.0*M_E))*psi_p)
                + 1.0j*HBAR*Q_E/(M_E*C)*(AX*grad_x_psi + AY*grad_y_psi))
    return kinetic + V*psi


def time_step(psi_i, dt, number_of_terms):
    psi0 = psi_i.copy()
    psi1 = (-1.0j*dt/(R*HBAR))*hamiltonian(psi_i)
    psi = CHEBYCHEV_COEFFS[0]*psi0 + CHEBYCHEV_COEFFS[1]*psi1
    psi_prev = [psi0, psi1]
    for k in range(2, number_of_terms):
        psi_k = (2.0*(-1.0j*dt/(R*HBAR))*hamiltonian(psi_prev[-1])
                  + psi_prev[-2])
        psi += CHEBYCHEV_COEFFS[k]*psi_k
        psi_prev[0], psi_prev[1] = psi_prev[1], psi_k
    return psi
    

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


def animation_func(*args):
    steps_per_frame = 2
    number_of_terms = 150
    for i in range(steps_per_frame):
        data['psi'] = time_step(data['psi'], DT, number_of_terms)
    psi = np.flip(data['psi'].reshape([N, N]), axis=0)
    im.set_data(np.angle(psi))
    abs_wavefunc2 = np.real(psi*np.conj(psi))
    alpha_map = 4.0*abs_wavefunc2/np.amax(abs_wavefunc2)
    im.set_alpha(np.where(alpha_map > 1.0, 1.0, alpha_map))
    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)}')
print(f'|psi(x)|^2 = {np.sum(np.abs(data["psi"])**2)}')
print(f'error: {1.0 - np.sum(np.abs(data["psi"])**2)}')