quantum-2d-square-potential.py

quantum-2d-square-potential.py

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

if __name__ == '__main__':
    # Constants and parameters
    hbar = 1.0  # Reduced Planck's constant
    m = 1.0     # Particle mass
    V0 = 1.0    # Barrier height

    # Spatial domain
    xmin, xmax = -50*3, 50*3
    ymin, ymax = -20*3, 20*3
    Nx, Ny = 256, 128  # Grid points
    x = np.linspace(xmin, xmax, Nx)
    y = np.linspace(ymin, ymax, Ny)
    dx = x[1] - x[0]
    dy = y[1] - y[0]

    X, Y = np.meshgrid(x, y, indexing='ij')

    # Potential barrier (centered rectangle)
    x1, x2 = 0, 5   # Barrier x-range
    y1, y2 = -5, 5  # Barrier y-range
    V = np.zeros_like(X)
    V[(X >= x1) & (X <= x2) & (Y >= y1) & (Y <= y2)] = V0

    # Initial Gaussian wave packet
    sigma_x, sigma_y = 2.0, 2.0
    x0, y0 = -20, 0  # Initial position
    kx = 1.0          # Initial momentum (ħk)
    ky = 0.0

    psi = np.exp(-((X - x0)**2/(2*sigma_x**2) + (Y - y0)**2/(2*sigma_y**2))) * np.exp(1j*(kx*X + ky*Y))
    psi /= np.sqrt(np.sum(np.abs(psi)**2) * dx * dy)  # Normalize

    # Precompute evolution operators
    dt = 0.1
    # Potential operator (half steps)
    potential_half = np.exp(-0.5j * V * dt / hbar)
    # Kinetic operator in Fourier space
    kx = 2 * np.pi * np.fft.fftfreq(Nx, dx)
    ky = 2 * np.pi * np.fft.fftfreq(Ny, dy)
    KX, KY = np.meshgrid(kx, ky, indexing='ij')
    K_squared = KX**2 + KY**2
    kinetic = np.exp(-1j * hbar * K_squared * dt / (2 * m))

    # Set up the figure
    fig, ax = plt.subplots()
    prob = np.abs(psi)**2
    im = ax.imshow(prob.T, extent=[xmin, xmax, ymin, ymax], origin='lower', aspect='auto', cmap='hot')
    plt.colorbar(im, label='Probability Density')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_title('Quantum Tunneling Through a Square Barrier')

    # Add barrier rectangle to the plot
    barrier_rect = plt.Rectangle((x1, y1), x2-x1, y2-y1, color='white', alpha=0.3)
    ax.add_patch(barrier_rect)

    def animate(frame):
        global psi
        # Advance 5 time steps per frame
        for _ in range(5):
            psi *= potential_half
            psi = np.fft.ifft2(np.fft.fft2(psi) * kinetic)
            psi *= potential_half
        prob = np.abs(psi)**2
        im.set_data(prob.T)
        im.set_clim(vmax=prob.max())  # Adjust color scale dynamically
        return im,

    # Create animation
    ani = FuncAnimation(fig, animate, frames=100*3, interval=50*3, blit=True)
    plt.show()