Created
July 25, 2022 01:10
-
-
Save marl0ny/7793dd975eb7d03594ed3a31b84e38d6 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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)}') |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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)}') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment