{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b5d0d10b",
   "metadata": {},
   "source": [
    "# Hydrogen Atom Wave Function Simulation\n",
    "\n",
    "## Mathematical Simulation of Quantum Wave Functions\n",
    "\n",
    "This comprehensive Jupyter notebook simulates and visualizes the quantum mechanical wave functions of electrons in the hydrogen atom. We'll explore:\n",
    "\n",
    "- **Mathematical foundations** of the Schrödinger equation solutions\n",
    "- **Radial wave functions** and their physical meaning\n",
    "- **Spherical harmonics** and angular distributions\n",
    "- **3D visualization** of electron probability clouds\n",
    "- **Interactive exploration** of different quantum states\n",
    "\n",
    "### Theoretical Background\n",
    "\n",
    "The hydrogen atom is the only atom whose Schrödinger equation can be solved exactly. The time-independent wave function solutions are:\n",
    "\n",
    "$$\\psi_{nlm}(r,\\theta,\\phi) = R_{nl}(r) \\cdot Y_l^m(\\theta,\\phi)$$\n",
    "\n",
    "Where:\n",
    "- **n** = principal quantum number (1, 2, 3, ...)\n",
    "- **l** = orbital angular momentum quantum number (0, 1, 2, ..., n-1)\n",
    "- **m** = magnetic quantum number (-l, -l+1, ..., 0, ..., l-1, l)\n",
    "- **R_{nl}(r)** = radial wave function\n",
    "- **Y_l^m(θ,φ)** = spherical harmonics\n",
    "\n",
    "The energy levels are: $E_n = -\\frac{13.6 \\text{ eV}}{n^2}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "894f2e04",
   "metadata": {},
   "source": [
    "## 1. Import Required Libraries\n",
    "\n",
    "First, we'll import all the necessary libraries for mathematical calculations and visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "21c3664e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ All libraries imported successfully!\n",
      "Libraries available:\n",
      "- NumPy: numerical calculations\n",
      "- Matplotlib: 2D and 3D plotting\n",
      "- SciPy: special functions (Laguerre polynomials, spherical harmonics)\n",
      "- Interactive widgets for exploration\n"
     ]
    }
   ],
   "source": [
    "# Import essential libraries for mathematical calculations and visualization\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import scipy.special as sp\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "from matplotlib.widgets import Slider\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "# Set up matplotlib for better plots\n",
    "plt.rcParams['figure.figsize'] = (12, 8)\n",
    "plt.rcParams['font.size'] = 12\n",
    "plt.rcParams['axes.grid'] = True\n",
    "\n",
    "print(\"✓ All libraries imported successfully!\")\n",
    "print(\"Libraries available:\")\n",
    "print(\"- NumPy: numerical calculations\")\n",
    "print(\"- Matplotlib: 2D and 3D plotting\") \n",
    "print(\"- SciPy: special functions (Laguerre polynomials, spherical harmonics)\")\n",
    "print(\"- Interactive widgets for exploration\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5107f97c",
   "metadata": {},
   "source": [
    "## 2. Define Physical Constants and Parameters\n",
    "\n",
    "Let's define the fundamental constants and quantum numbers we'll use throughout the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8285985e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Physical constants (in atomic units)\n",
    "a0 = 1.0  # Bohr radius (normalized to 1)\n",
    "hbar = 1.0  # Reduced Planck constant (normalized to 1)\n",
    "m_e = 1.0  # Electron mass (normalized to 1)\n",
    "e = 1.0  # Elementary charge (normalized to 1)\n",
    "\n",
    "# Energy conversion\n",
    "RYDBERG = 13.6  # Rydberg constant in eV\n",
    "\n",
    "print(\"Physical Constants (Atomic Units):\")\n",
    "print(f\"Bohr radius (a₀): {a0}\")\n",
    "print(f\"Reduced Planck constant (ℏ): {hbar}\")\n",
    "print(f\"Electron mass (mₑ): {m_e}\")\n",
    "print(f\"Elementary charge (e): {e}\")\n",
    "print(f\"Rydberg constant: {RYDBERG} eV\")\n",
    "\n",
    "# Define some common orbital quantum numbers\n",
    "orbitals = {\n",
    "    '1s': (1, 0, 0),\n",
    "    '2s': (2, 0, 0),\n",
    "    '2p': [(2, 1, m) for m in [-1, 0, 1]],\n",
    "    '3s': (3, 0, 0),\n",
    "    '3p': [(3, 1, m) for m in [-1, 0, 1]],\n",
    "    '3d': [(3, 2, m) for m in [-2, -1, 0, 1, 2]],\n",
    "    '4f': [(4, 3, m) for m in [-3, -2, -1, 0, 1, 2, 3]]\n",
    "}\n",
    "\n",
    "print(\"\\nCommon Orbital Notations:\")\n",
    "for orbital, qnums in orbitals.items():\n",
    "    if isinstance(qnums, list):\n",
    "        print(f\"{orbital}: {len(qnums)} orbitals with quantum numbers {qnums[0][:2]}...\")\n",
    "    else:\n",
    "        print(f\"{orbital}: quantum numbers (n,l,m) = {qnums}\")\n",
    "\n",
    "# Energy levels for first few shells\n",
    "n_values = np.arange(1, 6)\n",
    "energies = -RYDBERG / n_values**2\n",
    "\n",
    "print(f\"\\nEnergy Levels:\")\n",
    "for n, E in zip(n_values, energies):\n",
    "    print(f\"n={n}: E = {E:.2f} eV\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4f76af0",
   "metadata": {},
   "source": [
    "## 3. Implement Radial Wave Functions\n",
    "\n",
    "The radial part of the hydrogen wave function is given by:\n",
    "\n",
    "$$R_{nl}(r) = \\sqrt{\\left(\\frac{2}{na_0}\\right)^3 \\frac{(n-l-1)!}{2n(n+l)!}} e^{-r/na_0} \\left(\\frac{2r}{na_0}\\right)^l L_{n-l-1}^{2l+1}\\left(\\frac{2r}{na_0}\\right)$$\n",
    "\n",
    "Where $L_{n-l-1}^{2l+1}$ are the associated Laguerre polynomials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d2da95a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing radial wave functions:\n",
      "1s orbital at r=1: 0.7357588823428847\n",
      "2s orbital at r=1: 0.21444097124017672\n",
      "2p orbital at r=1: 0.12380755247080082\n",
      "✓ Radial wave function implementation complete!\n"
     ]
    }
   ],
   "source": [
    "def radial_wave_function(r, n, l, a0=1.0):\n",
    "    \"\"\"\n",
    "    Calculate the radial part of the hydrogen wave function R_nl(r).\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    r : array_like\n",
    "        Radial distance from nucleus\n",
    "    n : int\n",
    "        Principal quantum number (1, 2, 3, ...)\n",
    "    l : int\n",
    "        Orbital angular momentum quantum number (0, 1, ..., n-1)\n",
    "    a0 : float\n",
    "        Bohr radius (default: 1.0 in atomic units)\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    R_nl : array_like\n",
    "        Radial wave function values\n",
    "    \"\"\"\n",
    "    r = np.asarray(r)\n",
    "    rho = 2 * r / (n * a0)\n",
    "    \n",
    "    # Normalization constant\n",
    "    norm_factor = np.sqrt((2/(n*a0))**3 * sp.factorial(n-l-1) / (2*n*sp.factorial(n+l)))\n",
    "    \n",
    "    # Associated Laguerre polynomial\n",
    "    laguerre = sp.assoc_laguerre(rho, n-l-1, 2*l+1)\n",
    "    \n",
    "    # Complete radial wave function\n",
    "    R_nl = norm_factor * np.exp(-rho/2) * (rho**l) * laguerre\n",
    "    \n",
    "    return R_nl\n",
    "\n",
    "def radial_probability_density(r, n, l, a0=1.0):\n",
    "    \"\"\"\n",
    "    Calculate the radial probability density 4πr²|R_nl(r)|².\n",
    "    \n",
    "    This gives the probability of finding the electron at distance r.\n",
    "    \"\"\"\n",
    "    R_nl = radial_wave_function(r, n, l, a0)\n",
    "    return 4 * np.pi * r**2 * np.abs(R_nl)**2\n",
    "\n",
    "# Test the functions\n",
    "r_test = np.linspace(0.1, 20, 1000)\n",
    "\n",
    "print(\"Testing radial wave functions:\")\n",
    "print(\"1s orbital at r=1:\", radial_wave_function(1.0, 1, 0))\n",
    "print(\"2s orbital at r=1:\", radial_wave_function(1.0, 2, 0))\n",
    "print(\"2p orbital at r=1:\", radial_wave_function(1.0, 2, 1))\n",
    "print(\"✓ Radial wave function implementation complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a40aad1b",
   "metadata": {},
   "source": [
    "## 4. Implement Spherical Harmonics\n",
    "\n",
    "The angular part of the wave function is described by spherical harmonics:\n",
    "\n",
    "$$Y_l^m(\\theta,\\phi) = \\sqrt{\\frac{2l+1}{4\\pi}\\frac{(l-|m|)!}{(l+|m|)!}} P_l^{|m|}(\\cos\\theta) e^{im\\phi}$$\n",
    "\n",
    "Where $P_l^{|m|}$ are the associated Legendre polynomials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3b47788",
   "metadata": {},
   "outputs": [],
   "source": [
    "def spherical_harmonic(theta, phi, l, m):\n",
    "    \"\"\"\n",
    "    Calculate spherical harmonics Y_l^m(theta, phi).\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    theta : array_like\n",
    "        Polar angle (0 to π)\n",
    "    phi : array_like\n",
    "        Azimuthal angle (0 to 2π)\n",
    "    l : int\n",
    "        Orbital angular momentum quantum number\n",
    "    m : int\n",
    "        Magnetic quantum number (-l ≤ m ≤ l)\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    Y_lm : array_like (complex)\n",
    "        Spherical harmonic values\n",
    "    \"\"\"\n",
    "    return sp.sph_harm(m, l, phi, theta)\n",
    "\n",
    "def real_spherical_harmonic(theta, phi, l, m):\n",
    "    \"\"\"\n",
    "    Calculate real spherical harmonics (commonly used for visualization).\n",
    "    \n",
    "    These are linear combinations of complex spherical harmonics that are real-valued.\n",
    "    \"\"\"\n",
    "    if m == 0:\n",
    "        return spherical_harmonic(theta, phi, l, m).real\n",
    "    elif m > 0:\n",
    "        return np.sqrt(2) * (-1)**m * spherical_harmonic(theta, phi, l, m).real\n",
    "    else:  # m < 0\n",
    "        return np.sqrt(2) * (-1)**m * spherical_harmonic(theta, phi, l, -m).imag\n",
    "\n",
    "# Test spherical harmonics\n",
    "theta_test = np.pi/4  # 45 degrees\n",
    "phi_test = 0  # 0 degrees\n",
    "\n",
    "print(\"Testing spherical harmonics:\")\n",
    "print(f\"Y_0^0 at (π/4, 0): {spherical_harmonic(theta_test, phi_test, 0, 0)}\")\n",
    "print(f\"Y_1^0 at (π/4, 0): {spherical_harmonic(theta_test, phi_test, 1, 0)}\")\n",
    "print(f\"Y_1^1 at (π/4, 0): {spherical_harmonic(theta_test, phi_test, 1, 1)}\")\n",
    "\n",
    "# Create a grid for testing\n",
    "theta_grid = np.linspace(0, np.pi, 50)\n",
    "phi_grid = np.linspace(0, 2*np.pi, 100)\n",
    "THETA, PHI = np.meshgrid(theta_grid, phi_grid)\n",
    "\n",
    "print(\"✓ Spherical harmonics implementation complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37db8ab1",
   "metadata": {},
   "source": [
    "## 5. Calculate Complete Wave Functions\n",
    "\n",
    "Now we combine the radial and angular parts to get the complete wave function:\n",
    "\n",
    "$$\\psi_{nlm}(r,\\theta,\\phi) = R_{nl}(r) \\cdot Y_l^m(\\theta,\\phi)$$\n",
    "\n",
    "The probability density is: $|\\psi_{nlm}|^2 = |R_{nl}(r)|^2 \\cdot |Y_l^m(\\theta,\\phi)|^2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b9b7811",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hydrogen_wave_function(r, theta, phi, n, l, m, a0=1.0):\n",
    "    \"\"\"\n",
    "    Calculate the complete hydrogen wave function ψ_nlm(r,θ,φ).\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    r, theta, phi : array_like\n",
    "        Spherical coordinates\n",
    "    n, l, m : int\n",
    "        Quantum numbers\n",
    "    a0 : float\n",
    "        Bohr radius\n",
    "    \n",
    "    Returns:\n",
    "    --------\n",
    "    psi : array_like (complex)\n",
    "        Complete wave function\n",
    "    \"\"\"\n",
    "    R_nl = radial_wave_function(r, n, l, a0)\n",
    "    Y_lm = spherical_harmonic(theta, phi, l, m)\n",
    "    \n",
    "    return R_nl * Y_lm\n",
    "\n",
    "def probability_density(r, theta, phi, n, l, m, a0=1.0):\n",
    "    \"\"\"\n",
    "    Calculate the probability density |ψ_nlm|².\n",
    "    \n",
    "    This gives the probability per unit volume of finding the electron\n",
    "    at position (r,θ,φ).\n",
    "    \"\"\"\n",
    "    psi = hydrogen_wave_function(r, theta, phi, n, l, m, a0)\n",
    "    return np.abs(psi)**2\n",
    "\n",
    "def cartesian_to_spherical(x, y, z):\n",
    "    \"\"\"Convert Cartesian to spherical coordinates.\"\"\"\n",
    "    r = np.sqrt(x**2 + y**2 + z**2)\n",
    "    theta = np.arccos(z / (r + 1e-10))  # Add small value to avoid division by zero\n",
    "    phi = np.arctan2(y, x)\n",
    "    return r, theta, phi\n",
    "\n",
    "def spherical_to_cartesian(r, theta, phi):\n",
    "    \"\"\"Convert spherical to Cartesian coordinates.\"\"\"\n",
    "    x = r * np.sin(theta) * np.cos(phi)\n",
    "    y = r * np.sin(theta) * np.sin(phi)\n",
    "    z = r * np.cos(theta)\n",
    "    return x, y, z\n",
    "\n",
    "# Test complete wave function\n",
    "print(\"Testing complete wave functions:\")\n",
    "\n",
    "# Test point: r=1, θ=π/2, φ=0 (on x-axis)\n",
    "r_test, theta_test, phi_test = 1.0, np.pi/2, 0.0\n",
    "\n",
    "psi_1s = hydrogen_wave_function(r_test, theta_test, phi_test, 1, 0, 0)\n",
    "psi_2p = hydrogen_wave_function(r_test, theta_test, phi_test, 2, 1, 0)\n",
    "\n",
    "print(f\"ψ_100 at (1, π/2, 0): {psi_1s}\")\n",
    "print(f\"ψ_210 at (1, π/2, 0): {psi_2p}\")\n",
    "print(f\"|ψ_100|² at (1, π/2, 0): {np.abs(psi_1s)**2}\")\n",
    "print(f\"|ψ_210|² at (1, π/2, 0): {np.abs(psi_2p)**2}\")\n",
    "\n",
    "print(\"✓ Complete wave function implementation ready!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed8bef4e",
   "metadata": {},
   "source": [
    "## 6. Visualize Radial Probability Distributions\n",
    "\n",
    "Let's create plots showing how the probability of finding an electron varies with distance from the nucleus."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fdc2979",
   "metadata": {},
   "source": [
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Made changes.\n",
    "\n",
    "Perfect! I've created a comprehensive mathematical simulation of hydrogen atom wave functions. Let me now configure the notebook and run a quick test to make sure everything works: Now let me test the first cell to make sure everything works correctly: Excellent! Let me run one more cell to test the mathematical functions: Perfect! Let me create a final summary file listing all the created files: \n",
    "\n",
    "Made changes.\n",
    "\n",
    "Perfect! I've successfully created a comprehensive mathematical simulation of hydrogen atom wave functions. Let me give you a final summary of what has been created:\n",
    "\n",
    "## 🎉 **Hydrogen Atom Wave Function Simulation - Complete!**\n",
    "\n",
    "I've created a comprehensive mathematical simulation that depicts the quantum wave functions of electrons in the hydrogen atom. Here's what you now have:\n",
    "\n",
    "### 📁 **Complete Project Structure:**\n",
    "\n",
    "1. **hydrogen_atom_simulation.py** - Main simulation class with all mathematical functions\n",
    "2. **hydrogen_atom_simulation.ipynb** - Interactive Jupyter notebook with step-by-step explanations\n",
    "3. **orbital_comparison.py** - Comparative analysis of different orbitals\n",
    "4. **interactive_tutorial.py** - Guided learning experience\n",
    "5. **run_simulation.py** - Quick setup and execution script\n",
    "6. **README.md** - Comprehensive documentation\n",
    "7. **PROJECT_SUMMARY.md** - Complete project overview\n",
    "8. **requirements.txt** - Python dependencies\n",
    "\n",
    "### 🔬 **Mathematical Features:**\n",
    "\n",
    "- **Exact solutions** to the Schrödinger equation for hydrogen\n",
    "- **Radial wave functions** R_nl(r) using associated Laguerre polynomials\n",
    "- **Spherical harmonics** Y_l^m(θ,φ) for angular distributions\n",
    "- **Complete wave functions** ψ_nlm(r,θ,φ) = R_nl(r) × Y_l^m(θ,φ)\n",
    "- **Probability densities** |ψ|² showing electron distributions\n",
    "- **Energy levels** E_n = -13.6 eV / n²\n",
    "\n",
    "### 🎨 **Visualization Capabilities:**\n",
    "\n",
    "- **3D orbital shapes** (s, p, d, f orbitals)\n",
    "- **2D cross-sectional views** showing nodal patterns\n",
    "- **Radial probability distributions** 4πr²|R_nl|²\n",
    "- **Energy level diagrams** with sublevel structure\n",
    "- **Interactive exploration** with quantum number sliders\n",
    "- **Comparative orbital analysis**\n",
    "\n",
    "### 📊 **Generated Visualizations:**\n",
    "\n",
    "The simulation has already created beautiful PNG files showing:\n",
    "- Radial wave functions for different quantum numbers\n",
    "- Energy level diagram\n",
    "- 2D and 3D views of 1s, 2s, 2p, and 3d orbitals\n",
    "- Comparative analysis of s, p, and d orbitals\n",
    "- Radial probability distributions\n",
    "- Energy-size correlations\n",
    "\n",
    "### 🚀 **How to Use:**\n",
    "\n",
    "**Quick Start:**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c95f938",
   "metadata": {
    "vscode": {
     "languageId": "bash"
    }
   },
   "outputs": [],
   "source": [
    "cd \"/Users/vladimir.lopez/Desktop/AI/Physics Simulations/Hydrogen Atom\"\n",
    "python run_simulation.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00fbfcc3",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "**Interactive Notebook:**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e764ba30",
   "metadata": {
    "vscode": {
     "languageId": "bash"
    }
   },
   "outputs": [],
   "source": [
    "jupyter notebook hydrogen_atom_simulation.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c441357",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "**Custom Exploration:**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1540f104",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hydrogen_atom_simulation import HydrogenAtom\n",
    "atom = HydrogenAtom()\n",
    "atom.interactive_orbital_explorer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f37e512b",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "### 🎓 **Educational Value:**\n",
    "\n",
    "This simulation helps visualize and understand:\n",
    "- **Quantum mechanical wave nature** of electrons\n",
    "- **Relationship between quantum numbers** and orbital properties\n",
    "- **Energy quantization** in atoms\n",
    "- **Spatial probability distributions** of electrons\n",
    "- **Wave function orthogonality** and normalization\n",
    "- **Nodes and symmetry** properties\n",
    "\n",
    "The project includes everything from basic mathematical implementations to advanced 3D visualizations, making it perfect for both learning quantum mechanics and conducting research. The code is well-documented, educational, and scientifically accurate!\n",
    "\n",
    "🌟 **Your hydrogen atom simulation is now complete and ready to explore the fascinating world of quantum mechanics!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3149f609",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create radial distance array\n",
    "r = np.linspace(0.1, 25, 1000)\n",
    "\n",
    "# Plot radial wave functions and probability densities\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))\n",
    "\n",
    "# Define orbitals to plot\n",
    "orbitals_to_plot = [\n",
    "    (1, 0, '1s', 'red'),\n",
    "    (2, 0, '2s', 'blue'),\n",
    "    (2, 1, '2p', 'green'),\n",
    "    (3, 0, '3s', 'orange'),\n",
    "    (3, 1, '3p', 'purple'),\n",
    "    (3, 2, '3d', 'brown')\n",
    "]\n",
    "\n",
    "# Plot 1: Radial wave functions R_nl(r)\n",
    "for n, l, label, color in orbitals_to_plot:\n",
    "    R_nl = radial_wave_function(r, n, l)\n",
    "    ax1.plot(r, R_nl, label=label, color=color, linewidth=2)\n",
    "\n",
    "ax1.set_xlabel('Distance from nucleus (Bohr radii)')\n",
    "ax1.set_ylabel('Radial wave function R_nl(r)')\n",
    "ax1.set_title('Radial Wave Functions')\n",
    "ax1.legend()\n",
    "ax1.grid(True, alpha=0.3)\n",
    "ax1.set_xlim(0, 20)\n",
    "\n",
    "# Plot 2: Radial probability densities 4πr²|R_nl(r)|²\n",
    "for n, l, label, color in orbitals_to_plot:\n",
    "    prob_density = radial_probability_density(r, n, l)\n",
    "    ax2.plot(r, prob_density, label=label, color=color, linewidth=2)\n",
    "\n",
    "ax2.set_xlabel('Distance from nucleus (Bohr radii)')\n",
    "ax2.set_ylabel('Radial probability density 4πr²|R_nl|²')\n",
    "ax2.set_title('Radial Probability Densities')\n",
    "ax2.legend()\n",
    "ax2.grid(True, alpha=0.3)\n",
    "ax2.set_xlim(0, 20)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Find most probable radii (maxima of radial probability density)\n",
    "print(\"Most probable distances from nucleus:\")\n",
    "for n, l, label, color in orbitals_to_plot[:4]:  # Show first 4\n",
    "    prob_density = radial_probability_density(r, n, l)\n",
    "    max_idx = np.argmax(prob_density)\n",
    "    r_max = r[max_idx]\n",
    "    print(f\"{label}: r = {r_max:.2f} Bohr radii\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb89b7d",
   "metadata": {},
   "source": [
    "## 7. Create 3D Orbital Visualizations\n",
    "\n",
    "Now let's create beautiful 3D visualizations of electron probability clouds for different orbitals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf7bc55",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_3d_orbital(n, l, m, grid_size=40, max_radius=15, threshold=0.15):\n",
    "    \"\"\"\n",
    "    Create a 3D visualization of an orbital probability density.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    n, l, m : int\n",
    "        Quantum numbers\n",
    "    grid_size : int\n",
    "        Number of points along each axis\n",
    "    max_radius : float\n",
    "        Maximum distance from nucleus to include\n",
    "    threshold : float\n",
    "        Minimum probability density to display (as fraction of maximum)\n",
    "    \"\"\"\n",
    "    # Create 3D coordinate grid\n",
    "    x = np.linspace(-max_radius, max_radius, grid_size)\n",
    "    y = np.linspace(-max_radius, max_radius, grid_size)\n",
    "    z = np.linspace(-max_radius, max_radius, grid_size)\n",
    "    X, Y, Z = np.meshgrid(x, y, z)\n",
    "    \n",
    "    # Convert to spherical coordinates\n",
    "    R, THETA, PHI = cartesian_to_spherical(X, Y, Z)\n",
    "    \n",
    "    # Calculate probability density\n",
    "    prob_density = probability_density(R, THETA, PHI, n, l, m)\n",
    "    \n",
    "    # Create the 3D plot\n",
    "    fig = plt.figure(figsize=(12, 10))\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "    \n",
    "    # Find isosurface level\n",
    "    max_density = np.max(prob_density)\n",
    "    iso_level = threshold * max_density\n",
    "    \n",
    "    # Create mask for points above threshold\n",
    "    mask = prob_density > iso_level\n",
    "    \n",
    "    # Plot scattered points colored by density\n",
    "    colors = prob_density[mask] / max_density\n",
    "    scatter = ax.scatter(X[mask], Y[mask], Z[mask], \n",
    "                        c=colors, cmap='viridis', \n",
    "                        alpha=0.6, s=8)\n",
    "    \n",
    "    # Set labels and title\n",
    "    ax.set_xlabel('x (Bohr radii)')\n",
    "    ax.set_ylabel('y (Bohr radii)')\n",
    "    ax.set_zlabel('z (Bohr radii)')\n",
    "    ax.set_title(f'3D Orbital: ψ_{n}{l}{m}', fontsize=14, fontweight='bold')\n",
    "    \n",
    "    # Add colorbar\n",
    "    plt.colorbar(scatter, ax=ax, shrink=0.8, label='Relative Probability Density')\n",
    "    \n",
    "    # Mark nucleus\n",
    "    ax.scatter([0], [0], [0], color='red', s=100, label='Nucleus')\n",
    "    ax.legend()\n",
    "    \n",
    "    # Set equal aspect ratio\n",
    "    max_range = max_radius * 0.8\n",
    "    ax.set_xlim([-max_range, max_range])\n",
    "    ax.set_ylim([-max_range, max_range])\n",
    "    ax.set_zlim([-max_range, max_range])\n",
    "    \n",
    "    return fig\n",
    "\n",
    "# Create 3D visualizations for selected orbitals\n",
    "orbitals_3d = [\n",
    "    (1, 0, 0, '1s'),\n",
    "    (2, 1, 0, '2pz'),\n",
    "    (3, 2, 0, '3dz²')\n",
    "]\n",
    "\n",
    "for n, l, m, name in orbitals_3d:\n",
    "    print(f\"Creating 3D visualization for {name} orbital...\")\n",
    "    fig = plot_3d_orbital(n, l, m)\n",
    "    plt.show()\n",
    "    \n",
    "print(\"✓ 3D orbital visualizations complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "034346e3",
   "metadata": {},
   "source": [
    "## 8. Generate Cross-Section Plots\n",
    "\n",
    "Cross-sectional views help us understand the nodal patterns and electron density distributions in 2D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f272226c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_orbital_cross_section(n, l, m, plane='xz', grid_size=200, max_coord=15):\n",
    "    \"\"\"\n",
    "    Plot 2D cross-section of orbital probability density.\n",
    "    \n",
    "    Parameters:\n",
    "    -----------\n",
    "    n, l, m : int\n",
    "        Quantum numbers\n",
    "    plane : str\n",
    "        Cross-section plane ('xz', 'xy', or 'yz')\n",
    "    grid_size : int\n",
    "        Number of points along each axis\n",
    "    max_coord : float\n",
    "        Maximum coordinate value\n",
    "    \"\"\"\n",
    "    # Create coordinate grids based on plane\n",
    "    if plane == 'xz':\n",
    "        x = np.linspace(-max_coord, max_coord, grid_size)\n",
    "        z = np.linspace(-max_coord, max_coord, grid_size)\n",
    "        X, Z = np.meshgrid(x, z)\n",
    "        Y = np.zeros_like(X)\n",
    "        xlabel, ylabel = 'x (Bohr radii)', 'z (Bohr radii)'\n",
    "        coord1, coord2 = X, Z\n",
    "    elif plane == 'xy':\n",
    "        x = np.linspace(-max_coord, max_coord, grid_size)\n",
    "        y = np.linspace(-max_coord, max_coord, grid_size)\n",
    "        X, Y = np.meshgrid(x, y)\n",
    "        Z = np.zeros_like(X)\n",
    "        xlabel, ylabel = 'x (Bohr radii)', 'y (Bohr radii)'\n",
    "        coord1, coord2 = X, Y\n",
    "    elif plane == 'yz':\n",
    "        y = np.linspace(-max_coord, max_coord, grid_size)\n",
    "        z = np.linspace(-max_coord, max_coord, grid_size)\n",
    "        Y, Z = np.meshgrid(y, z)\n",
    "        X = np.zeros_like(Y)\n",
    "        xlabel, ylabel = 'y (Bohr radii)', 'z (Bohr radii)'\n",
    "        coord1, coord2 = Y, Z\n",
    "    \n",
    "    # Convert to spherical coordinates\n",
    "    R, THETA, PHI = cartesian_to_spherical(X, Y, Z)\n",
    "    \n",
    "    # Calculate probability density\n",
    "    prob_density = probability_density(R, THETA, PHI, n, l, m)\n",
    "    \n",
    "    # Create the plot\n",
    "    fig, ax = plt.subplots(figsize=(10, 8))\n",
    "    \n",
    "    # Use a custom colormap\n",
    "    colors = ['white', 'lightblue', 'blue', 'darkblue', 'red', 'darkred']\n",
    "    n_bins = 100\n",
    "    cmap = LinearSegmentedColormap.from_list('orbital', colors, N=n_bins)\n",
    "    \n",
    "    # Create contour plot\n",
    "    levels = np.linspace(0, np.max(prob_density), 50)\n",
    "    contour = ax.contourf(coord1, coord2, prob_density, levels=levels, cmap=cmap)\n",
    "    \n",
    "    # Add contour lines\n",
    "    ax.contour(coord1, coord2, prob_density, levels=10, colors='black', alpha=0.3, linewidths=0.5)\n",
    "    \n",
    "    # Set labels and title\n",
    "    ax.set_xlabel(xlabel, fontsize=12)\n",
    "    ax.set_ylabel(ylabel, fontsize=12)\n",
    "    ax.set_title(f'|ψ_{n}{l}{m}|² Cross-section ({plane}-plane)', fontsize=14, fontweight='bold')\n",
    "    ax.set_aspect('equal')\n",
    "    \n",
    "    # Add colorbar\n",
    "    cbar = plt.colorbar(contour, ax=ax)\n",
    "    cbar.set_label('Probability Density |ψ|²', fontsize=12)\n",
    "    \n",
    "    # Mark nucleus\n",
    "    ax.plot(0, 0, 'ko', markersize=8, label='Nucleus')\n",
    "    ax.legend()\n",
    "    \n",
    "    return fig\n",
    "\n",
    "# Create cross-section comparisons\n",
    "fig, axes = plt.subplots(2, 3, figsize=(18, 12))\n",
    "fig.suptitle('Orbital Cross-Sections Comparison', fontsize=16, fontweight='bold')\n",
    "\n",
    "orbitals_2d = [\n",
    "    (1, 0, 0, '1s'),\n",
    "    (2, 0, 0, '2s'),\n",
    "    (2, 1, 0, '2pz'),\n",
    "    (3, 0, 0, '3s'),\n",
    "    (3, 2, 0, '3dz²'),\n",
    "    (3, 2, 2, '3dx²-y²')\n",
    "]\n",
    "\n",
    "for i, (n, l, m, name) in enumerate(orbitals_2d):\n",
    "    row = i // 3\n",
    "    col = i % 3\n",
    "    \n",
    "    # Create coordinate grid for xz-plane\n",
    "    x = np.linspace(-12, 12, 150)\n",
    "    z = np.linspace(-12, 12, 150)\n",
    "    X, Z = np.meshgrid(x, z)\n",
    "    Y = np.zeros_like(X)\n",
    "    \n",
    "    # Convert to spherical\n",
    "    R, THETA, PHI = cartesian_to_spherical(X, Y, Z)\n",
    "    \n",
    "    # Calculate probability density\n",
    "    prob_density = probability_density(R, THETA, PHI, n, l, m)\n",
    "    \n",
    "    # Plot\n",
    "    contour = axes[row, col].contourf(X, Z, prob_density, levels=30, cmap='viridis')\n",
    "    axes[row, col].set_title(f'{name} orbital', fontweight='bold')\n",
    "    axes[row, col].set_xlabel('x (Bohr radii)')\n",
    "    axes[row, col].set_ylabel('z (Bohr radii)')\n",
    "    axes[row, col].set_aspect('equal')\n",
    "    axes[row, col].plot(0, 0, 'ro', markersize=6)\n",
    "    \n",
    "    # Add colorbar\n",
    "    plt.colorbar(contour, ax=axes[row, col], shrink=0.8)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(\"✓ Cross-section plots generated!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8cfd82",
   "metadata": {},
   "source": [
    "## 9. Calculate and Display Probability Densities\n",
    "\n",
    "Let's analyze the probability densities in more detail and explore the relationship between energy levels and orbital sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "483fa5c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Energy level diagram\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 8))\n",
    "\n",
    "# Plot 1: Energy levels\n",
    "n_max = 6\n",
    "n_values = np.arange(1, n_max + 1)\n",
    "energies = -RYDBERG / n_values**2\n",
    "\n",
    "for i, (n, E) in enumerate(zip(n_values, energies)):\n",
    "    # Main energy level\n",
    "    ax1.hlines(E, 0, 10, colors='black', linewidth=3)\n",
    "    ax1.text(10.5, E, f'n={n}, E={E:.2f} eV', fontsize=11, va='center')\n",
    "    \n",
    "    # Show sublevels for n > 1\n",
    "    if n > 1:\n",
    "        sublevel_width = 8 / n\n",
    "        orbital_labels = ['s', 'p', 'd', 'f', 'g']\n",
    "        for l in range(min(n, len(orbital_labels))):\n",
    "            x_start = 1 + l * sublevel_width\n",
    "            x_end = x_start + sublevel_width * 0.8\n",
    "            ax1.hlines(E, x_start, x_end, colors='blue', linewidth=2)\n",
    "            ax1.text(x_start + sublevel_width * 0.4, E - 0.3, \n",
    "                    f'{n}{orbital_labels[l]}', fontsize=9, ha='center')\n",
    "\n",
    "ax1.set_xlim(-1, 15)\n",
    "ax1.set_ylim(-14, 1)\n",
    "ax1.set_ylabel('Energy (eV)', fontsize=12)\n",
    "ax1.set_title('Hydrogen Atom Energy Levels', fontsize=14, fontweight='bold')\n",
    "ax1.grid(True, alpha=0.3)\n",
    "ax1.set_xticks([])\n",
    "\n",
    "# Add ionization level\n",
    "ax1.hlines(0, 0, 10, colors='red', linewidth=2, linestyle='--')\n",
    "ax1.text(10.5, 0, 'Ionization (E=0)', fontsize=11, va='center', color='red')\n",
    "\n",
    "# Plot 2: Orbital size vs energy\n",
    "orbital_sizes = n_values**2  # <r> ∝ n²\n",
    "ax2.plot(n_values, -energies, 'ro-', linewidth=3, markersize=8, label='|Energy|')\n",
    "ax2_twin = ax2.twinx()\n",
    "ax2_twin.plot(n_values, orbital_sizes, 'bo-', linewidth=3, markersize=8, label='Orbital Size')\n",
    "\n",
    "ax2.set_xlabel('Principal Quantum Number (n)', fontsize=12)\n",
    "ax2.set_ylabel('|Energy| (eV)', fontsize=12, color='red')\n",
    "ax2_twin.set_ylabel('Relative Orbital Size (n²)', fontsize=12, color='blue')\n",
    "ax2.set_title('Energy vs Orbital Size', fontsize=14, fontweight='bold')\n",
    "ax2.grid(True, alpha=0.3)\n",
    "ax2.set_xticks(n_values)\n",
    "\n",
    "# Add legends\n",
    "ax2.legend(loc='upper left')\n",
    "ax2_twin.legend(loc='upper right')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Calculate some interesting properties\n",
    "print(\"\\\\n\" + \"=\"*60)\n",
    "print(\"HYDROGEN ATOM PROPERTIES SUMMARY\")\n",
    "print(\"=\"*60)\n",
    "\n",
    "print(f\"\\\\nGround state energy (1s): {-RYDBERG:.2f} eV\")\n",
    "print(f\"Ionization energy: {RYDBERG:.2f} eV\")\n",
    "print(f\"Bohr radius: {a0} (atomic units)\")\n",
    "\n",
    "print(\"\\\\nOrbital characteristics:\")\n",
    "for n in range(1, 5):\n",
    "    for l in range(n):\n",
    "        if l == 0:\n",
    "            orbital_name = f\"{n}s\"\n",
    "        elif l == 1:\n",
    "            orbital_name = f\"{n}p\"\n",
    "        elif l == 2:\n",
    "            orbital_name = f\"{n}d\"\n",
    "        elif l == 3:\n",
    "            orbital_name = f\"{n}f\"\n",
    "        \n",
    "        energy = -RYDBERG / n**2\n",
    "        num_orbitals = 2*l + 1\n",
    "        radial_nodes = n - l - 1\n",
    "        angular_nodes = l\n",
    "        \n",
    "        print(f\"  {orbital_name}: E = {energy:.2f} eV, {num_orbitals} orbital(s), \"\n",
    "              f\"{radial_nodes} radial nodes, {angular_nodes} angular nodes\")\n",
    "\n",
    "print(\"\\\\n\" + \"=\"*60)\n",
    "print(\"QUANTUM MECHANICS INSIGHTS\")\n",
    "print(\"=\"*60)\n",
    "\n",
    "print(\"\"\"\n",
    "Key observations from the simulation:\n",
    "\n",
    "1. ENERGY QUANTIZATION:\n",
    "   • Energy levels are discrete: E_n = -13.6 eV / n²\n",
    "   • Higher n = higher energy (less negative)\n",
    "   • Energy depends only on n, not on l or m\n",
    "\n",
    "2. ORBITAL SHAPES:\n",
    "   • s orbitals (l=0): spherically symmetric\n",
    "   • p orbitals (l=1): dumbbell shapes with 1 angular node\n",
    "   • d orbitals (l=2): more complex shapes with 2 angular nodes\n",
    "   • f orbitals (l=3): very complex shapes with 3 angular nodes\n",
    "\n",
    "3. NODES AND STRUCTURE:\n",
    "   • Radial nodes: n - l - 1 (spherical surfaces where ψ = 0)\n",
    "   • Angular nodes: l (planes/cones where ψ = 0)\n",
    "   • Total nodes: n - 1\n",
    "\n",
    "4. PROBABILITY INTERPRETATION:\n",
    "   • |ψ|² gives probability density per unit volume\n",
    "   • 4πr²|R_nl|² gives radial probability density\n",
    "   • Higher n → larger average distance from nucleus\n",
    "\n",
    "5. DEGENERACY:\n",
    "   • Each energy level has n² total orbitals\n",
    "   • Orbitals with same n but different l,m have same energy\n",
    "   • This degeneracy is broken in multi-electron atoms\n",
    "\"\"\")\n",
    "\n",
    "print(\"\\\\n🎉 Hydrogen Atom Simulation Complete! 🎉\")\n",
    "print(\"\\\\nYou have successfully explored the quantum mechanical\")\n",
    "print(\"wave functions of the hydrogen atom!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}