Last active
May 16, 2024 00:47
-
-
Save marl0ny/b0689700be41f97b9c74df512e72d44e 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
import numpy as np | |
from scipy.fft import dst, idst | |
from scipy.integrate import trapezoid | |
import matplotlib.pyplot as plt | |
import matplotlib.animation as animation | |
# Constants | |
N: int = 512 | |
X = np.arange(0, N, dtype=np.float64) | |
DX: float = X[1] - X[0] | |
DT: float = 0.5 # time step | |
# Energies of the stationary states for the infinite square well. | |
# Note that the discrete sin transformation used for getting the wave function | |
# in terms of the stationary states of the infinite square well does not | |
# include the two boundary points where the wave function vanishes: | |
# hence the extra 2*DX term in the denominator. | |
# Formula for infinite square well energies from | |
# https://en.wikipedia.org/wiki/Particle_in_a_box | |
E = np.pi**2*np.arange(1, N+1)**2/(2.0*(float(N)+2.0*DX)**2) | |
def time_evolve_infinite_square_well(psi: np.ndarray, t: float) -> np.ndarray: | |
""" | |
Get the wave function in terms of infinite square well eigenfunctions by | |
doing a discrete sin transformation on it, then time evolve it using the | |
infinite square well energies. Inverse discrete sin transform to get | |
it back in its position representation. | |
This will eventually be used in the split operator method; | |
see eq 2.22 and 2.23 of [1]. | |
[1] Antoine X., Bao W., Besse C. | |
Computational methods for the dynamics of the | |
nonlinear Schrodinger/Gross-Pitaevskii equations. | |
https://arxiv.org/abs/1305.1093 | |
""" | |
# Use the definition of the type I DST to get the wave function in terms | |
# of the stationary states of the infinite square well. See | |
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.fft.dst.html | |
return idst(np.exp(-1.0j*E*t)*dst(psi, type=1), type=1) | |
def phi(x: np.ndarray, t: float) -> np.ndarray: | |
""" | |
The time varying potential phi(x, t). | |
""" | |
# This function currently implements a harmonic oscillator | |
# whose minimum oscillates in time. Note that since the DST is used | |
# in the integration method, the wave function must always | |
# vanish at the boundaries, no matter what potential is used. | |
# Formula for the harmonic oscillator Hamiltonian from | |
# https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator | |
kt: float = 0.00001*t | |
# Corresponds to omega = sqrt(2)/N | |
return ((x/float(N)) - (0.5 + 0.05*np.sin(kt + 150.0*kt**2)))**2 | |
def split_step(psi: np.ndarray, # wave function | |
phi: 'function', # potential | |
t: float, dt: float) -> np.ndarray: | |
""" | |
Advance the wave function in time using the Split step method. | |
For a time dependent potential this follows II.1 of [1], as well | |
as eq. 2.22 - 2.23 of [2]. | |
[1] Bauke H., Keitel C. Accelerating the Fourier split operator | |
method via graphics processing units. | |
https://arxiv.org/abs/1012.3911 | |
[2] Antoine X., Bao W., Besse C. | |
Computational methods for the dynamics of the | |
nonlinear Schrodinger/Gross-Pitaevskii equations. | |
https://arxiv.org/abs/1305.1093 | |
[3] Schloss J. The Split Operator Method - Arcane Algorithm Archive. | |
https://www.algorithm-archive.org/contents/split-operator_method/ | |
split-operator_method.html | |
""" | |
# Numerically integrate the time dependent potential from t to t + dt. | |
int_phi_dt = trapezoid([phi(X, t + i*dt/3.0) | |
for i in range(4)], dx=dt/3.0, axis=0) | |
ux = np.exp(-0.5j*int_phi_dt) | |
return ux*time_evolve_infinite_square_well(ux*psi, dt) | |
# Currently initialized as the ground state of the harmonic oscillator. | |
# Formula for the harmonic oscillator eigenstates from | |
# https://en.wikipedia.org/wiki/Quantum_harmonic_oscillator | |
psi = np.exp(-0.5*float(N)*np.sqrt(2.0)*(X/float(N) - 0.5)**2)*(0.0 + 1j) | |
# Now do some fancy plots and animations ##################################### | |
fig = plt.figure(figsize=(5, 7) | |
# figsize=(6, 8) | |
) | |
ax, ax2 = fig.subplots(2, 1, height_ratios=[1, 0.2]) | |
ax.set_title('Position space') | |
ax2.set_title('Momentum space') | |
VIEW_SCALE = 4.0 | |
line1, = ax.plot(X, VIEW_SCALE*phi(X, 0.0), | |
label='V(x, t)', color='gray', linewidth=2.0) | |
line2, = ax.plot(X, np.real(psi), label=r'$Re(\psi(x))$') | |
c_re = line2.get_color() | |
line3, = ax.plot(X, np.imag(psi), label=r'$Im(\psi(x))$') | |
c_im = line3.get_color() | |
line4, = ax.plot(X, np.abs(psi), label=r'$|\psi(x)|$', color='black') | |
freq = np.fft.fftshift(np.fft.fftfreq(N)) | |
psi_p = np.fft.fftshift(np.fft.fft(psi)) | |
line5, = ax2.plot(freq, np.real(psi_p), color=c_re) | |
line6, = ax2.plot(freq, np.imag(psi_p), color=c_im) | |
line7, = ax2.plot(freq, np.abs(psi_p), color='gray') | |
ax.set_xlim(X[0], X[-1]) | |
ax.set_ylim(-1.2, 1.2) | |
ax.set_ylabel(r'$\psi(x)$') | |
ax.legend(loc='lower left') | |
ax2.set_facecolor('black') | |
ax2.set_ylabel(r'$\psi(p)$') | |
ax2.set_xlim(np.amin(freq)/4.0, np.amax(freq)/4.0) | |
ax2.set_xlabel('Frequency') | |
ax2.grid(linestyle='--', linewidth=0.5) | |
ax.set_xlabel(r'Position $x$ (a.u.)') | |
fig.tight_layout() | |
animation_data = {'psi': psi, | |
'lines': [line1, line2, line3, | |
line4, line5, line6, line7], | |
't': 0.0} | |
def animation_func(*_): | |
lines = animation_data['lines'] | |
for _ in range(100): | |
t = animation_data['t'] | |
psi = animation_data['psi'] | |
animation_data['psi'] = split_step(psi, phi, t, DT) | |
animation_data['t'] += DT | |
lines[0].set_ydata(VIEW_SCALE*phi(X, t)) | |
lines[1].set_ydata(np.real(psi)) | |
lines[2].set_ydata(np.imag(psi)) | |
lines[3].set_ydata(np.abs(psi)) | |
psi_p = np.fft.fftshift(np.fft.fft(psi)) | |
lines[4].set_ydata(np.real(psi_p)) | |
lines[5].set_ydata(np.imag(psi_p)) | |
lines[6].set_ydata(np.abs(psi_p)) | |
return lines | |
a = animation.FuncAnimation(fig, animation_func, | |
# frames=1800, | |
blit=True, interval=1000//60) | |
# a.save('animation2.gif', writer='ffmpeg', fps=60, bitrate=1800) | |
plt.show() | |
plt.close() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment