import numpy as np
import matplotlib.pyplot as plt
from scipy.special import hermite
from scipy.special import factorial
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
# --- Set up plot styling ---
plt.rcParams['text.usetex'] = True
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.size'] = 12
# Parameters
m = 1.0 # mass
omega = 1.0 # angular frequency
hbar = 1.0 # reduced Planck constant
x = np.linspace(-5, 5, 1000)
# Harmonic oscillator potential
V = 0.5 * m * omega**2 * x**2
# Wavefunction for the n-th energy level
def psi_n(n, x, m=1.0, omega=1.0, hbar=1.0):
"""Calculate the n-th energy eigenstate wavefunction"""
alpha = np.sqrt(m * omega / hbar)
normalization = (alpha / np.pi)**0.25 / np.sqrt(2**n * factorial(n))
H_n = hermite(n)
return normalization * np.exp(-alpha**2 * x**2 / 2) * H_n(alpha * x)
# Energy levels
def E_n(n, omega=1.0, hbar=1.0):
"""Energy of the n-th level"""
return hbar * omega * (n + 0.5)
# Coherent state coefficients
def coherent_state_coeff(n, alpha):
"""Coefficient c_n for coherent state |alpha>"""
return np.exp(-abs(alpha)**2 / 2) * alpha**n / np.sqrt(factorial(n))
# Coherent state wavefunction
def coherent_state_psi(x, alpha, t, n_max=30, m=1.0, omega=1.0, hbar=1.0):
"""Time-dependent coherent state wavefunction"""
psi = np.zeros_like(x, dtype=complex)
for n in range(n_max):
c_n = coherent_state_coeff(n, alpha)
E = E_n(n, omega, hbar)
phase = np.exp(-1j * E * t / hbar)
psi += c_n * psi_n(n, x, m, omega, hbar) * phase
return psi
# Find alpha for desired coherent state energy
def find_alpha_for_energy(E_coh, omega=1.0, hbar=1.0):
"""Find alpha such that <E> = E_coh"""
# For coherent state: <E> = hbar*omega*(|alpha|^2 + 1/2)
# E_coh = hbar*omega*(|alpha|^2 + 1/2)
# |alpha|^2 = E_coh/(hbar*omega) - 1/2
alpha_squared = E_coh / (hbar * omega) - 0.5
return np.sqrt(max(0, alpha_squared))
# Setup
n_levels = 8
E_coh = (E_n(0) + E_n(3)) / 2 # Energy between E_0 and E_7
alpha = find_alpha_for_energy(E_coh, omega, hbar)
#print(f"Coherent state energy: E_coh = {E_coh:.2f}")
#print(f"Alpha parameter: α = {alpha:.2f}")
#print(f"Average quantum number: <n> = |α|² = {alpha**2:.2f}")
# Create figure
fig, ax = plt.subplots(figsize=(7.5, 6))
# Plot potential
ax.plot(x, V, 'k-', linewidth=1.5, alpha=0.3)
# Colors for eigenstates
colors = plt.cm.Set2(np.linspace(0, 1, n_levels))
# Plot energy levels and scaled wavefunctions
scale = 0.3
eigenstate_plots = []
for n in range(n_levels):
E = E_n(n, omega, hbar)
# Plot energy level line
ax.axhline(y=E, color='gray', linestyle='-', linewidth=0.8, alpha=0.5)
# Calculate wavefunction and coherent state coefficient
psi = psi_n(n, x, m, omega, hbar)
c_n = coherent_state_coeff(n, alpha)
# Scale by coherent state coefficient (use absolute value for scaling)
scaling_factor = abs(c_n) * 5 # Extra factor for visibility
psi_scaled = scale * scaling_factor * psi + E
# Fill and plot
fill = ax.fill_between(x, E, psi_scaled, alpha=0.6, color=colors[n])
line, = ax.plot(x, psi_scaled, color=colors[n], linewidth=1.5, alpha=0.8)
eigenstate_plots.append((fill, line))
# Add energy level labels
ax.text(-5.2, E+0.07, f'$E_{n}$', fontsize=12, va='bottom', ha='right')
# Add wavefunction labels
ax.text(5.2, E+0.07, f'$\\psi_{n}(x)$', fontsize=11, va='bottom', ha='left')
# Add energy spacing annotations
E0 = E_n(0, omega, hbar)
E1 = E_n(1, omega, hbar)
ax.annotate('', xy=(-4.5, E1), xytext=(-4.5, E0),
arrowprops=dict(arrowstyle='<->', color='black', lw=1.5))
ax.text(-4.8, (E0 + E1)/2, r'$\hbar\omega$', fontsize=13, va='center', ha='right')
ax.annotate('', xy=(-4.5, E0), xytext=(-4.5, 0),
arrowprops=dict(arrowstyle='<->', color='black', lw=1.5))
ax.text(-4.8, E0/2, r'$\hbar\omega/2$', fontsize=13, va='center', ha='right')
# Mark coherent state energy level
ax.axhline(y=E_coh, color='red', linestyle='--', linewidth=0.5, alpha=0.7)
ax.text(-5.2, E_coh+0.07, r'$E_{{coh}}=\langle \hat H \rangle$', fontsize=12, va='bottom', ha='right', color='red')
# Initialize coherent state plot
coh_scale = 3 # Scale for coherent state probability density
psi_coh_t0 = coherent_state_psi(x, alpha, 0, n_max=30, m=m, omega=omega, hbar=hbar)
prob_density = np.abs(psi_coh_t0)**2
prob_scaled = coh_scale * prob_density + E_coh
coherent_fill = ax.fill_between(x, E_coh, prob_scaled, alpha=0.7, color='red', label='Coherent state')
coherent_line, = ax.plot(x, prob_scaled, 'r-', linewidth=2, alpha=0.9)
# Axes formatting
ax.plot([0,0], [0,E_n(n_levels-1) + 0.7], color='black', linewidth=1, alpha=0.7)
ax.axhline(y=0, color='black', linewidth=1.0, alpha=0.5)
ax.set_xlim(-5.5, 5.5)
ax.set_ylim(-0.5, E_n(n_levels-1) + 0.7)
ax.set_xlabel('x', fontsize=14, loc='right')
#ax.set_ylabel('$\\psi(x)$', fontsize=14, loc='top', rotation=0)
ax.yaxis.set_label_coords(0.02, 0.98)
ax.set_xticks([0])
ax.set_yticks([])
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
ax.spines['left'].set_visible(False)
ax.spines['bottom'].set_position('zero')
ax.plot(1, 0, ">k", transform=ax.get_yaxis_transform(), clip_on=False, markersize=8)
time_text = ax.text(0.02, 1, '', transform=ax.transAxes, fontsize=12,
verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
#ax.set_title('Quantum Harmonic Oscillator: Coherent State Evolution',
# fontsize=14, pad=20, fontweight='bold')
# Animation
def animate(frame):
t = frame # Time step
# Update coherent state
psi_coh = coherent_state_psi(x, alpha, t, n_max=30, m=m, omega=omega, hbar=hbar)
prob_density = np.abs(psi_coh)**2
prob_scaled = coh_scale * prob_density + E_coh
# Update coherent state plot
coherent_line.set_ydata(prob_scaled)
# Update fill
global coherent_fill
coherent_fill.remove()
coherent_fill = ax.fill_between(x, E_coh, prob_scaled, alpha=0.7, color='red')
# Update time text
time_text.set_text(f't = {t:.1f}')
return coherent_line, time_text
t_max = 37.8
frames = [i * 0.2 for i in range(int(t_max/0.2)+1)]
anim = FuncAnimation(fig, animate, frames=frames, interval=50, blit=False)
plt.tight_layout()
plt.close()
# Display as HTML5 video in Jupyter
HTML(anim.to_html5_video())