Quantum Harmonic Oscillator

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

# Quantum grid
x = np.linspace(-6, 6, 300)
y = np.linspace(-6, 6, 300)
X, Y = np.meshgrid(x, y)

# Quantum numbers
n, m = 3, 2

# Hermite polynomials
Hn = hermite(n)
Hm = hermite(m)

# Wavefunction
psi = (1 / np.sqrt(2**n * factorial(n) * np.sqrt(np.pi)) * 
       1 / np.sqrt(2**m * factorial(m) * np.sqrt(np.pi)) *
       Hn(X) * Hm(Y) * np.exp(-(X**2 + Y**2) / 2))
prob = np.abs(psi)**2

# Quantum neon theme
plt.style.use('dark_background')
fig, ax = plt.subplots(figsize=(10, 8))
fig.patch.set_facecolor('#050510')
ax.set_facecolor('#050510')

# Custom colormap - quantum blue-cyan-white
colors = ['#050510', '#0a0a30', '#0a1a50', '#0a3a70', '#0a5a90', 
          '#1a8ab0', '#3abad0', '#6aeaff', '#aaffff', '#ffffff']
cmap = LinearSegmentedColormap.from_list('quantum', colors, N=256)

# Filled contour
cf = ax.contourf(X, Y, prob, levels=35, cmap=cmap)

# Nodal lines (magenta for contrast)
cs = ax.contour(X, Y, prob, levels=[0.00005], colors='#ff00ff', 
                linewidths=2, alpha=0.8)

# Subtle contour lines
ax.contour(X, Y, prob, levels=12, colors='#6aeaff', linewidths=0.3, alpha=0.3)

# Quantum state label
ax.text(0.03, 0.97, '|n=%d, m=%d>' % (n, m), transform=ax.transAxes, fontsize=14, 
        color='#ff00ff', fontweight='bold', va='top', fontfamily='monospace',
        bbox=dict(boxstyle='round', facecolor='#050510', edgecolor='#ff00ff', alpha=0.8))

# Styled colorbar
cbar = plt.colorbar(cf, ax=ax, pad=0.02, shrink=0.85)
cbar.ax.set_facecolor('#050510')
cbar.set_label('Probability Density', fontsize=11, color='#c0e0ff', labelpad=10)
cbar.ax.yaxis.set_tick_params(color='#c0e0ff')
cbar.outline.set_edgecolor('#2a3a5a')
plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='#c0e0ff', fontsize=9)

# Labels
ax.set_xlabel('x (units of a)', fontsize=13, color='#c0e0ff', fontweight='600', labelpad=10)
ax.set_ylabel('y (units of a)', fontsize=13, color='#c0e0ff', fontweight='600', labelpad=10)
ax.set_title('2D Quantum Harmonic Oscillator', fontsize=16, color='white', 
             fontweight='bold', pad=20)

# Style axes
ax.tick_params(colors='#80b0d0', labelsize=10, length=0)
for spine in ax.spines.values():
    spine.set_visible(False)

ax.set_aspect('equal')
plt.tight_layout()
plt.show()