!sudo apt update
!sudo apt install ffmpeg -y
# Import the required libraries
import scipy
from scipy.fftpack import fft, ifft, fftshift, ifftshift
import scipy.special as scsp
import numpy as np
import matplotlib.pyplot as plt
from functools import partial
from IPython.display import Video
num_points = 10001
left, right = -200, 200
xspace = np.linspace(left, right, num_points)
delta_x = (right - left) / (num_points)
delta_k = 2 * np.pi / (right - left)
k0 = -0.5 * num_points * delta_k
k1 = -k0
kspace = np.linspace(k0, k1, num_points)
omega = 0.5
x0 = 5
num_tsteps = 1000
# Set up the potentials and the eigenstat
def V_i(omega: float, x: np.ndarray) -> np.ndarray:
return 0.5 * omega ** 2 * x ** 2
def V_T(omega: float, x: np.ndarray, x0: float) -> np.ndarray:
return 0.5 * omega ** 2 * (x - x0) ** 2
def V(t: float, t_total: float) -> np.ndarray:
if t < t_total:
return (1 - t / t_total) * V_i(omega, xspace) + t / t_total * V_T(omega, xspace, x0)
else:
return V_T(omega, xspace, x0)
def psi_n(n: int, omega: float, x0: float = 0) -> np.ndarray:
"""Generates the nth eigenstate of the quantum potential, uses the known solution"""
hermiten = scsp.hermite(n)
array = (
1 / (np.sqrt(2 ** n * scsp.factorial(n)))
* (omega / np.pi) ** (1 / 4)
* np.exp(-omega * (xspace - x0) ** 2 / 2)
* hermiten(np.sqrt(omega) * (xspace - x0))
)
array = array.astype(np.complex128)
return array
psi_x_0 = psi_n(0, 1)
xlabels = [
"Pr(Ground)",
"Pr(Excited_1)",
"Pr(Excited_2)",
"Pr(Excited_3)",
"Pr(Excited_4)",
"Pr(Excited_5)",
"Pr(Excited_6)",
"Pr(Rest)",
]
def calculate_inner(psi_x, n, t, t_total, omega):
"""Check how much probability is in the nth eigenstate"""
if t < t_total:
return np.abs(delta_x * np.inner(psi_n(n, omega, t/t_total * x0), psi_x)) ** 2
else:
return np.abs(delta_x * np.inner(psi_n(n, omega, x0), psi_x)) ** 2
def progress_x(psi_x, dt, t, t_total):
"""Progress half a step in x"""
return psi_x * np.exp(-0.5 * 1j * dt * V(t, t_total))
def progress_k(psi_k, dt):
"""Progress a full step in k"""
return psi_k * np.exp(-0.5 * 1j * dt * kspace * kspace)
def progress_step(psi_x, dt, t, t_total):
"""Progress a full step in position and momentum"""
psi_x = progress_x(psi_x, dt, t, t_total)
psi_k = fftshift(fft(psi_x))
psi_k = progress_k(psi_k, dt)
psi_x = ifft(ifftshift(psi_k))
psi_x = progress_x(psi_x, dt, t + dt, t_total)
return psi_x
# Run the slow evolution simulation
max_t = 70
t_total = 60
# SIMULATE
ts = np.linspace(0, max_t, num_tsteps)
dt = t_total / num_tsteps
# Generates the inner product of the nth eigenstate with the wavefunction
inner_calcs = [partial(calculate_inner, n=n, t_total=t_total, omega=omega) for n in range(7)]
psi_x_0 = psi_n(0, omega)
psi_xs = []
psi_inners = np.zeros((num_tsteps, len(xlabels)))
for i_idx, t in enumerate(ts):
psi_xs.append(psi_x_0)
psi_x_0 = progress_step(psi_x_0, dt, t, t_total)
for j_idx, calc in enumerate(inner_calcs):
psi_inners[i_idx, j_idx] = calc(psi_x_0, t=t)
psi_inners[i_idx, -1] = 1 - np.sum(psi_inners[i_idx, :-1])
psi_xs = np.array(psi_xs)
plot_anim(psi_xs, psi_inners, "slow_evolution")
Video('./slow_evolution.mp4', embed=True)
# Run the fast evolution simulation
max_t = 70
t_total = 20
ts = np.linspace(0, max_t, num_tsteps)
dt = t_total / num_tsteps
inner_calcs = [partial(calculate_inner, n=n, t_total=t_total, omega=omega) for n in range(7)]
psi_x_0 = psi_n(0, omega)
psi_xs = []
psi_inners = np.zeros((num_tsteps, len(xlabels)))
for i_idx, t in enumerate(ts):
psi_xs.append(psi_x_0)
psi_x_0 = progress_step(psi_x_0, dt, t, t_total)
for j_idx, calc in enumerate(inner_calcs):
psi_inners[i_idx, j_idx] = calc(psi_x_0, t=t)
psi_inners[i_idx, -1] = 1 - np.sum(psi_inners[i_idx, :-1])
psi_xs = np.array(psi_xs)
plot_anim(psi_xs, psi_inners, "fast_evolution")
Video('./fast_evolution.mp4', embed=True)