A Python Code for Animating Free Particle Wave Packet Evolution

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Constants
hbar = 1.0545718e-34  # Reduced Planck constant (Joule second)
m_e = 9.10938356e-31  # Mass of electron (kg)
wavelength = 20e-9  # de Broglie wavelength (meters)
k0 = 2 * np.pi / wavelength  # Wave number

# Grid parameters
Lx, Ly = 1e-6, 1e-6  # Simulation domain size (meters)
Nx, Ny = 512, 512  # Number of grid points
dx, dy = Lx / Nx, Ly / Ny  # Grid spacing
x = np.linspace(-Lx / 2, Lx / 2, Nx)
y = np.linspace(-Ly / 2, Ly / 2, Ny)
X, Y = np.meshgrid(x, y)

# Time parameters
dt = 3e-14  # Time step (seconds)
Nt = 451  # Number of time steps

# Wave packet parameters
sigma_x = sigma_y = 15e-9  # Width of the Gaussian packet (meters)
x0, y0 = -2e-7, 0  # Initial position of the packet (meters)
kx0, ky0 = k0, 0  # Initial wave vector

# Initial Gaussian wave packet
psi0 = np.exp(-(X - x0) ** 2 / (2 * sigma_x ** 2)) * np.exp(-(Y - y0) ** 2 / (2 * sigma_y ** 2))
psi0 = psi0.astype(np.complex128)
psi0 *= np.exp(1j * (kx0 * X + ky0 * Y))

# Normalize the initial wave packet
psi0 /= np.sqrt(np.sum(np.abs(psi0) ** 2))

# Fourier space components
kx = np.fft.fftfreq(Nx, dx) * 2 * np.pi
ky = np.fft.fftfreq(Ny, dy) * 2 * np.pi
KX, KY = np.meshgrid(kx, ky)
K2 = KX ** 2 + KY ** 2

# Split-step Fourier method
psi = psi0.copy()

# Prepare for animation
fig = plt.figure(figsize=(12, 9))

# Main plot (takes more space)
ax_main = fig.add_axes([0.1, 0.1, 0.65, 0.65])

# X-projection plot (at the top)
ax_xproj = fig.add_axes([0.18, 0.8, 0.49, 0.15])

# Y-projection plot (on the right)
ax_yproj = fig.add_axes([0.8, 0.1, 0.15, 0.65])

def update(t):
    global psi
    # Full step in momentum space with kinetic energy operator
    psi_k = np.fft.fft2(psi)
    psi_k *= np.exp(-1j * hbar * K2 * dt / (2 * m_e))
    psi = np.fft.ifft2(psi_k)

    # Visualization of wave packet evolution
    ax_main.clear()
    ax_main.set_xlabel(r'$x$ (m)', fontsize=14)
    ax_main.set_ylabel(r'$y$ (m)', fontsize=14)
    ax_main.tick_params(which="major", axis="both", direction="in", top=True, right=True, length=5, width=1, labelsize=12)
    Rho = np.abs(psi) ** 2 
    im = ax_main.imshow(Rho, extent=(-Lx / 2, Lx / 2, -Ly / 2, Ly / 2), cmap='hot', animated=True)
    ax_main.set_title(f'Time step: {t}', fontsize=14)

    # Projection along x-axis
    ax_xproj.clear()
    ax_xproj.plot(x, np.sum(np.abs(psi) ** 2, axis=0))
    ax_xproj.set_xlim(-Lx / 2, Lx / 2)
    ax_xproj.set_ylabel(r'$|\psi|^2$', fontsize=14)
    ax_xproj.set_title('Projection along X-axis', fontsize=14)
    ax_xproj.tick_params(which="major", axis="both", direction="in", top=True, right=True, length=5, width=1, labelsize=12)

    # Projection along y-axis
    ax_yproj.clear()
    ax_yproj.plot(np.sum(np.abs(psi) ** 2, axis=1), y)
    ax_yproj.set_ylim(-Ly / 2, Ly / 2)
    ax_yproj.set_xlabel(r'$|\psi|^2$', fontsize=14)
    ax_yproj.set_title('Projection along Y-axis', fontsize=14, loc='left', pad=20)
    ax_yproj.yaxis.tick_right()
    ax_yproj.tick_params(which="major", axis="both", direction="in", top=True, right=True, length=5, width=1, labelsize=12)

    return [im]

ani = animation.FuncAnimation(fig, update, frames=Nt, repeat=False)

# Save the animation
ani.save('free_particle_evolution_with_projections.mp4', writer='ffmpeg', fps=20)

plt.show()