4. Quantum Harmonic Oscillator

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import hermite
from scipy.special import factorial

# Physical constants
hbar = 1.055e-34  # J·s
m = 1.67e-27      # proton mass (kg)
omega = 1e14      # angular frequency (rad/s)

# Characteristic length scale
x0 = np.sqrt(hbar / (m * omega))

# Position array (in units of x0)
x = np.linspace(-5, 5, 1000) * x0

def psi_n(n, x_vals):
    """Wave function for quantum harmonic oscillator state n"""
    # Normalization constant
    N_n = 1 / np.sqrt(2**n * factorial(n)) * (m*omega/(np.pi*hbar))**(1/4)

    # Hermite polynomial of order n
    H_n = hermite(n)

    # Dimensionless position
    xi = x_vals / x0

    # Wave function
    return N_n * H_n(xi) * np.exp(-xi**2 / 2)

# Calculate energies
def E_n(n):
    """Energy of nth level"""
    return hbar * omega * (n + 0.5)

# Display energies
print("Quantum Harmonic Oscillator")
print(f"m = {m:.2e} kg, ω = {omega:.2e} rad/s")
print(f"Characteristic length x₀ = {x0*1e12:.3f} pm")
print("=" * 60)
print("Energy Levels:")
for n in range(5):
    E = E_n(n)
    E_eV = E / 1.602e-19
    print(f"  n = {n}: E = {E:.3e} J = {E_eV:.4f} eV")
    print(f"          E_{n}/E_0 = {(n + 0.5) / 0.5:.1f}")

# Plot wave functions and probability densities
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

# Plot wave functions
colors = ['blue', 'red', 'green', 'purple']
for n in range(4):
    psi = psi_n(n, x)
    # Offset by energy level for clarity
    offset = E_n(n) / (hbar * omega)
    ax1.plot(x/x0, psi*5 + offset, label=f'n={n}', color=colors[n])
    # Draw energy level line
    ax1.axhline(y=offset, color=colors[n], linestyle='--', alpha=0.3, linewidth=1)

# Plot potential
x_pot = np.linspace(-5, 5, 100)
V = 0.5 * x_pot**2
ax1.plot(x_pot, V, 'k-', linewidth=2, label='V(x)')
ax1.set_xlabel('Position (x/x₀)')
ax1.set_ylabel('Energy (ℏω)')
ax1.set_title('Wave Functions with Energy Levels')
ax1.legend(fontsize=9)
ax1.grid(True, alpha=0.3)
ax1.set_ylim(-0.5, 5)

# Plot probability densities
for n in range(4):
    psi = psi_n(n, x)
    prob = np.abs(psi)**2
    ax2.plot(x/x0, prob*x0, label=f'n={n}', color=colors[n])

ax2.set_xlabel('Position (x/x₀)')
ax2.set_ylabel('|ψ(x)|²')
ax2.set_title('Probability Densities')
ax2.legend()
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('harmonic_oscillator.png', dpi=150, bbox_inches='tight')
print("\nPlot saved as 'harmonic_oscillator.png'")