quantum_mechanics.py

#!/usr/bin/env python3
"""
Quantum Mechanics Simulations for Computational Physics
========================================================

This module contains implementations of fundamental quantum mechanics simulations:
- Time-dependent Schrödinger equation solver
- Quantum harmonic oscillator eigenstates
- Particle in a box solutions
- Quantum tunneling demonstrations

Keywords: quantum mechanics python, schrodinger equation solver, quantum harmonic oscillator,
particle in a box simulation, quantum tunneling python code

Author: CoCalc Scientific Templates
License: MIT
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse import diags
from scipy.linalg import expm
from scipy.special import hermite, factorial

# Set reproducible random seed
np.random.seed(42)

class SchrodingerSolver:
    """
    Time-dependent Schrödinger equation solver using finite difference method.

    Solves: iℏ ∂ψ/∂t = Ĥ ψ where Ĥ = T̂ + V̂
    """

    def __init__(self, L=10.0, N=1000, hbar=1.0, m=1.0):
        """
        Initialize the solver.

        Parameters:
        -----------
        L : float
            Length of simulation box
        N : int
            Number of grid points
        hbar : float
            Reduced Planck constant (natural units)
        m : float
            Particle mass (natural units)
        """
        self.L = L
        self.N = N
        self.hbar = hbar
        self.m = m
        self.dx = L / N
        self.x = np.linspace(0, L, N)

        # Kinetic energy operator (finite difference)
        kinetic_diag = -2 * np.ones(N)
        kinetic_off = np.ones(N-1)
        self.T_matrix = -(hbar**2 / (2*m*self.dx**2)) * (
            diags([kinetic_off, kinetic_diag, kinetic_off], [-1, 0, 1], shape=(N, N))
        )

    def set_potential(self, V):
        """Set the potential energy function V(x)."""
        self.V = V
        self.H = self.T_matrix + diags(V, 0)

    def gaussian_wavepacket(self, x0, sigma, k0):
        """
        Create a Gaussian wave packet initial condition.

        Parameters:
        -----------
        x0 : float
            Initial position
        sigma : float
            Wave packet width
        k0 : float
            Initial momentum
        """
        psi = np.exp(-(self.x - x0)**2 / (2*sigma**2)) * np.exp(1j * k0 * self.x)
        return psi / np.sqrt(np.trapezoid(np.abs(psi)**2, self.x))

    def evolve(self, psi_initial, dt, t_final):
        """
        Time evolution using matrix exponentiation.

        Returns:
        --------
        times : array
            Time points
        psi_evolution : list
            Wavefunction at each time step
        """
        times = np.arange(0, t_final, dt)
        U = expm(-1j * self.H.toarray() * dt / self.hbar)

        psi = psi_initial.copy()
        psi_evolution = [psi.copy()]

        for t in times[1:]:
            psi = U @ psi
            psi_evolution.append(psi.copy())

        return times, psi_evolution

class HarmonicOscillator:
    """
    Quantum harmonic oscillator analytical solutions.

    Energy eigenvalues: E_n = ℏω(n + 1/2)
    Wavefunctions: ψ_n(x) = N_n * exp(-ξ²/2) * H_n(ξ)
    where ξ = x/x₀ and x₀ = √(ℏ/mω)
    """

    def __init__(self, omega=1.0, m=1.0, hbar=1.0):
        """
        Parameters:
        -----------
        omega : float
            Angular frequency
        m : float
            Mass
        hbar : float
            Reduced Planck constant
        """
        self.omega = omega
        self.m = m
        self.hbar = hbar
        self.x0 = np.sqrt(hbar / (m * omega))  # Length scale

    def energy_level(self, n):
        """Energy eigenvalue for quantum number n."""
        return self.hbar * self.omega * (n + 0.5)

    def wavefunction(self, x, n):
        """
        Harmonic oscillator wavefunction for quantum number n.

        Parameters:
        -----------
        x : array
            Position grid
        n : int
            Quantum number
        """
        # Normalization factor
        norm = 1.0 / np.sqrt(2**n * factorial(n)) * (self.m*self.omega/(np.pi*self.hbar))**(1/4)

        # Dimensionless coordinate
        xi = x / self.x0

        # Hermite polynomial
        H_n = hermite(n)

        # Wavefunction
        psi_n = norm * np.exp(-xi**2 / 2) * H_n(xi)
        return psi_n

    def calculate_eigenstates(self, x_range, n_max=5):
        """
        Calculate first n_max energy eigenstates.

        Returns:
        --------
        x : array
            Position grid
        energies : list
            Energy eigenvalues
        wavefunctions : list
            Corresponding wavefunctions
        """
        x = np.linspace(-x_range, x_range, 1000)

        energies = []
        wavefunctions = []

        for n in range(n_max):
            E_n = self.energy_level(n)
            psi_n = self.wavefunction(x, n)

            energies.append(E_n)
            wavefunctions.append(psi_n)

        return x, energies, wavefunctions

class ParticleInBox:
    """
    Infinite square well (particle in a box) exact solutions.

    Energy eigenvalues: E_n = n²π²ℏ²/(2mL²)
    Wavefunctions: ψ_n(x) = √(2/L) * sin(nπx/L)
    """

    def __init__(self, L=1.0, m=1.0, hbar=1.0):
        """
        Parameters:
        -----------
        L : float
            Box length
        m : float
            Particle mass
        hbar : float
            Reduced Planck constant
        """
        self.L = L
        self.m = m
        self.hbar = hbar

    def energy_level(self, n):
        """Energy eigenvalue for quantum number n (n ≥ 1)."""
        return (n**2 * np.pi**2 * self.hbar**2) / (2 * self.m * self.L**2)

    def wavefunction(self, x, n):
        """
        Particle in box wavefunction for quantum number n.

        Parameters:
        -----------
        x : array
            Position grid (0 ≤ x ≤ L)
        n : int
            Quantum number (n ≥ 1)
        """
        return np.sqrt(2/self.L) * np.sin(n * np.pi * x / self.L)

    def superposition_state(self, x, coefficients, quantum_numbers):
        """
        Create superposition of energy eigenstates.

        ψ(x) = Σ c_n ψ_n(x)

        Parameters:
        -----------
        x : array
            Position grid
        coefficients : array
            Complex amplitudes
        quantum_numbers : array
            Corresponding quantum numbers
        """
        psi = np.zeros_like(x, dtype=complex)

        for c, n in zip(coefficients, quantum_numbers):
            psi += c * self.wavefunction(x, n)

        return psi

    def calculate_eigenstates(self, n_max=5):
        """
        Calculate first n_max energy eigenstates.

        Returns:
        --------
        x : array
            Position grid
        energies : list
            Energy eigenvalues
        wavefunctions : list
            Corresponding wavefunctions
        """
        x = np.linspace(0, self.L, 1000)

        energies = []
        wavefunctions = []

        for n in range(1, n_max + 1):
            E_n = self.energy_level(n)
            psi_n = self.wavefunction(x, n)

            energies.append(E_n)
            wavefunctions.append(psi_n)

        return x, energies, wavefunctions

def demo_quantum_tunneling():
    """
    Demonstrate quantum tunneling through a potential barrier.
    """
    print("=== Quantum Tunneling Demo ===")

    # Setup solver
    solver = SchrodingerSolver(L=10.0, N=1000)

    # Define potential: harmonic well + Gaussian barrier
    omega = 1.0
    V = 0.5 * solver.m * omega**2 * (solver.x - solver.L/2)**2
    V += 2.0 * np.exp(-((solver.x - solver.L/3)/(solver.L/20))**2)

    solver.set_potential(V)

    # Initial Gaussian wave packet
    psi_initial = solver.gaussian_wavepacket(x0=solver.L/4, sigma=solver.L/20, k0=5.0)

    # Time evolution
    times, psi_evolution = solver.evolve(psi_initial, dt=0.001, t_final=5.0)

    print(f"Simulation completed: {len(times)} time steps")
    print(f"Final norm: {np.trapezoid(np.abs(psi_evolution[-1])**2, solver.x):.6f}")

    return solver.x, V, times, psi_evolution

def demo_harmonic_oscillator():
    """
    Demonstrate quantum harmonic oscillator eigenstates.
    """
    print("=== Harmonic Oscillator Demo ===")

    oscillator = HarmonicOscillator(omega=1.0, m=1.0, hbar=1.0)
    x, energies, wavefunctions = oscillator.calculate_eigenstates(x_range=4*oscillator.x0, n_max=6)

    print(f"Calculated {len(energies)} eigenstates")
    print(f"Energy levels (in units of ℏω): {[E/oscillator.hbar/oscillator.omega for E in energies[:4]]}")

    return x, energies, wavefunctions, oscillator

def demo_particle_in_box():
    """
    Demonstrate particle in a box solutions.
    """
    print("=== Particle in Box Demo ===")

    box = ParticleInBox(L=1.0, m=1.0, hbar=1.0)
    x, energies, wavefunctions = box.calculate_eigenstates(n_max=5)

    # Create superposition state
    coefficients = [1/np.sqrt(2), 1/np.sqrt(2)]
    quantum_numbers = [1, 2]
    psi_superposition = box.superposition_state(x, coefficients, quantum_numbers)

    print(f"Calculated {len(energies)} eigenstates")
    print(f"Energy ratios E_n/E_1: {[E/energies[0] for E in energies]}")

    return x, energies, wavefunctions, psi_superposition

if __name__ == "__main__":
    # Run demonstrations
    demo_quantum_tunneling()
    demo_harmonic_oscillator()
    demo_particle_in_box()