Computational Demo: Harmonic oscillator

import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import interact, Dropdown, FloatSlider

# Physical constants
h = 6.62607015e-34        # Planck constant (J s)
hbar = h / (2 * np.pi)
c = 2.99792458e10         # speed of light (cm / s)

# Atomic masses in kg
atomic_masses = {
    "H": 1.6735575e-27,
    "C": 1.9944235e-26,
    "N": 2.3258671e-26,
    "O": 2.6566962e-26,
    "F": 3.1547626e-26
}

def reduced_mass(m1, m2):
    return (m1 * m2) / (m1 + m2)

def frequency(k, mu):
    return (1 / (2 * np.pi)) * np.sqrt(k / mu)

def energy_levels(n, nu):
    return h * nu * (n + 0.5)

def turning_point(E, k):
    return np.sqrt(2 * E / k)

def harmonic_oscillator_plot(atom1, atom2, k):
    # Masses and reduced mass
    m1 = atomic_masses[atom1]
    m2 = atomic_masses[atom2]
    mu = reduced_mass(m1, m2)
    
    # Frequency
    nu = frequency(k, mu)
    nu_cm = nu / c
    
    # Energy levels
    v_levels = np.arange(4)
    energies = energy_levels(v_levels, nu)
    
    # Position grid (set by highest turning point)
    x_max = turning_point(energies[-1], k) * 1.3
    x = np.linspace(-x_max, x_max, 1000)
    
    # Harmonic potential
    V = 0.5 * k * x**2
    
    # Plot
    plt.figure(figsize=(8,6))
    plt.plot(x * 1e10, V, label="Harmonic Potential", color="black")
    
    # Plot energy levels and turning points
    for v, E in zip(v_levels, energies):
        xtp = turning_point(E, k)
        plt.hlines(E, -xtp*1e10, xtp*1e10, linestyles="dashed")
        plt.text(xtp*1e10*1.05, E, f"$v={v}$", verticalalignment="center")
    
    # Fundamental transition arrow
    plt.annotate(
        "",
        xy=(0, energies[1]),
        xytext=(0, energies[0]),
        arrowprops=dict(arrowstyle="<->", linewidth=2)
    )
    plt.text(0.1, 0.5*(energies[0]+energies[1]), "Fundamental\nTransition")
    
    # Labels
    plt.xlabel("Displacement $x$ (Å)")
    plt.ylabel("Energy (J)")
    plt.title(f"Harmonic Oscillator: {atom1}-{atom2}")
    plt.legend()
    plt.grid(True)
    plt.show()
    
    # Print frequencies
    print(f"Reduced mass μ = {mu:.3e} kg")
    print(f"Fundamental frequency ν = {nu:.3e} Hz")
    print(f"Fundamental frequency = {nu_cm:.1f} cm⁻¹")

interact(
    harmonic_oscillator_plot,
    atom1=Dropdown(options=atomic_masses.keys(), value="H", description="Atom 1"),
    atom2=Dropdown(options=atomic_masses.keys(), value="F", description="Atom 2"),
    k=FloatSlider(
        value=1000,
        min=500,
        max=1500,
        step=50,
        description="Force constant (N/m)",
        continuous_update=True
    )
);