{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro",
   "metadata": {},
   "source": [
    "# Rb 87 D2: Cycling Transition, σ⁺ Optical Pumping\n",
    "\n",
    "The D2 line of $^{87}$Rb connects the $5S_{1/2}$ ground state to the\n",
    "$5P_{3/2}$ excited state at 780 nm. With circularly polarised ($\\sigma^+$,\n",
    "$q=+1$) light tuned to the $F=2 \\to F'=3$ transition, the atoms are\n",
    "optically pumped into the $m_F=+2$ **cycling state** — the sub-transition\n",
    "$|F=2,\\,m_F=+2\\rangle \\to |F'=3,\\,m_{F'}=+3\\rangle$ from which the only\n",
    "allowed decay is back to $|F=2,\\,m_F=+2\\rangle$.\n",
    "\n",
    "## Level structure\n",
    "\n",
    "```\n",
    "5P₃/₂  F'=3  mF = -3 -2 -1  0 +1 +2 +3\n",
    "               ─────────────────────────\n",
    "                    σ⁺ (q=+1) drive\n",
    "5S₁/₂  F=2   mF = -2 -1  0 +1 +2\n",
    "               ─────────────────\n",
    "```\n",
    "\n",
    "The $F=2 \\to F'=3$ transition is **closed** under $\\sigma^+$: spontaneous\n",
    "decay from any $m_{F'}$ can only reach $m_F \\leq m_{F'}$, so population\n",
    "cascades toward $m_F=+2$ over many scattering cycles.\n",
    "\n",
    "## Hyperfine structure construction\n",
    "\n",
    "The `hyperfine.Atom1e` helper computes the Clebsch-Gordan-weighted coupling\n",
    "and decay factors for the full $m_F$ sublevel structure. This eliminates\n",
    "manual look-up of angular-momentum coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from maxwellbloch import hyperfine, mb_solve, plot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hfs-section",
   "metadata": {},
   "source": [
    "## Build the hyperfine structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "build-hfs",
   "metadata": {},
   "outputs": [],
   "source": [
    "# D2 line: 5S₁/₂ (J=1/2) → 5P₃/₂ (J=3/2), I = 3/2 for ⁸⁷Rb\n",
    "Rb87_5s12_F2 = hyperfine.LevelF(I=1.5, J=0.5, F=2)\n",
    "Rb87_5p32_F3 = hyperfine.LevelF(I=1.5, J=1.5, F=3)\n",
    "\n",
    "atom1e = hyperfine.Atom1e(element=\"Rb\", isotope=\"87\")\n",
    "atom1e.add_F_level(Rb87_5s12_F2)   # index 0: F=2, 5 mF sublevels\n",
    "atom1e.add_F_level(Rb87_5p32_F3)   # index 1: F'=3, 7 mF sublevels\n",
    "\n",
    "NUM_STATES = atom1e.get_num_mF_levels()\n",
    "print(f\"Total mF sublevels: {NUM_STATES}\")  # 5 + 7 = 12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "build-channels",
   "metadata": {},
   "outputs": [],
   "source": [
    "# σ⁺ field: couples F=2 → F'=3 with q=+1 (Δm_F = +1)\n",
    "q = 1\n",
    "FIELD_CHANNELS = atom1e.get_coupled_levels(F_level_idxs_a=(0,), F_level_idxs_b=(1,))\n",
    "FIELD_FACTORS  = atom1e.get_clebsch_hf_factors(F_level_idxs_a=(0,), F_level_idxs_b=(1,), q=q)\n",
    "DECAY_CHANNELS = atom1e.get_coupled_levels(F_level_idxs_a=(0,), F_level_idxs_b=(1,))\n",
    "DECAY_FACTORS  = atom1e.get_decay_factors(F_level_idxs_a=(0,), F_level_idxs_b=(1,))\n",
    "ENERGIES       = atom1e.get_energies()\n",
    "\n",
    "# Initial state: uniform across F=2 (mF = -2, -1, 0, +1, +2)\n",
    "INITIAL_STATE = [1.0/5.0]*5 + [0.0]*7\n",
    "\n",
    "# Detuning: on resonance with F=2 → F'=3\n",
    "DETUNING = 0.0\n",
    "\n",
    "print(f\"Field channels: {len(FIELD_CHANNELS)}\")\n",
    "print(f\"sum(|d_i|²) field factors: {np.sum(FIELD_FACTORS**2):.4f}\")\n",
    "print(f\"sum(|d_i|²) decay factors: {np.sum(DECAY_FACTORS**2):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "solve-section",
   "metadata": {},
   "source": [
    "## Solve: weak σ⁺ probe, optical pumping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "solve",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json = \"\"\"\n",
    "{{\n",
    "  \"atom\": {{\n",
    "    \"decays\": [\n",
    "      {{\n",
    "        \"channels\": {decay_channels},\n",
    "        \"rate\": 1.0,\n",
    "        \"factors\": {decay_factors}\n",
    "      }}\n",
    "    ],\n",
    "    \"energies\": {energies},\n",
    "    \"fields\": [\n",
    "      {{\n",
    "        \"coupled_levels\": {field_channels},\n",
    "        \"factors\": {field_factors},\n",
    "        \"detuning\": {detuning},\n",
    "        \"detuning_positive\": true,\n",
    "        \"label\": \"probe\",\n",
    "        \"rabi_freq\": 1e-3,\n",
    "        \"rabi_freq_t_args\": {{\n",
    "          \"ampl\": 1.0,\n",
    "          \"centre\": 0.0,\n",
    "          \"fwhm\": 1.0\n",
    "        }},\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      }}\n",
    "    ],\n",
    "    \"num_states\": {num_states},\n",
    "    \"initial_state\": {initial_state}\n",
    "  }},\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 10.0,\n",
    "  \"t_steps\": 100,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 100,\n",
    "  \"interaction_strengths\": [1.0e2],\n",
    "  \"savefile\": \"mbs-Rb87_5s12_5p32_F23_q1-weak-pulse-decay\"\n",
    "}}\n",
    "\"\"\".format(\n",
    "    num_states=NUM_STATES,\n",
    "    energies=ENERGIES,\n",
    "    initial_state=INITIAL_STATE,\n",
    "    detuning=DETUNING,\n",
    "    field_channels=FIELD_CHANNELS,\n",
    "    field_factors=FIELD_FACTORS.tolist(),\n",
    "    decay_channels=DECAY_CHANNELS,\n",
    "    decay_factors=DECAY_FACTORS.tolist(),\n",
    ")\n",
    "\n",
    "mbs = mb_solve.MBSolve().from_json_str(mb_solve_json)\n",
    "mbs.mbsolve(recalc=False)\n",
    "print(\"Done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "field-section",
   "metadata": {},
   "source": [
    "## Probe propagation\n",
    "\n",
    "The space-time plot shows the weak probe pulse propagating through the\n",
    "Rb vapour. On resonance, the pulse is absorbed as it passes through the\n",
    "dense medium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-field",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.field_spacetime(mbs, field_idx=0)\n",
    "fig.update_layout(title=\"Probe |Ω(z, t)| — Rb 87 D2, F=2→F'=3, σ⁺\")\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "population-section",
   "metadata": {},
   "source": [
    "## Optical pumping: mF sublevel populations\n",
    "\n",
    "The plot below shows the 12 sublevel populations at the input face as a\n",
    "function of time. States 0–4 are the $F=2$ ground sublevels\n",
    "($m_F = -2, -1, 0, +1, +2$) and states 5–11 are the $F'=3$ excited\n",
    "sublevels ($m_{F'} = -3, \\ldots, +3$).\n",
    "\n",
    "Starting from a uniform distribution across $F=2$, $\\sigma^+$ excitation\n",
    "preferentially drives $\\Delta m_F = +1$ transitions. After several\n",
    "absorption/emission cycles, population accumulates in $m_F = +2$ — the\n",
    "cycling state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-population",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.population(mbs, state_indices=list(range(NUM_STATES)), z_idx=0)\n",
    "fig.update_layout(title=\"mF sublevel populations at z = z_min — D2 σ⁺ optical pumping\")\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ground-section",
   "metadata": {},
   "source": [
    "## Ground-state populations only\n",
    "\n",
    "Zooming in on just the $F=2$ ground sublevels makes the optical pumping\n",
    "dynamics clearer: $m_F = +2$ grows at the expense of lower $m_F$ states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-ground",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.population(mbs, state_indices=list(range(5)), z_idx=0)\n",
    "fig.update_layout(\n",
    "    title=\"F=2 ground sublevel populations at z = z_min — D2 σ⁺ optical pumping\",\n",
    ")\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "summary",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "- The $F=2 \\to F'=3$ transition on the Rb 87 D2 line is a **closed cycling\n",
    "  transition** under $\\sigma^+$ excitation: spontaneous emission from\n",
    "  $F'=3$ can only reach $F=2$, not $F=1$.\n",
    "- Starting from uniform $m_F$ distribution, $\\sigma^+$ optical pumping\n",
    "  drives population toward $m_F=+2$ over multiple scattering cycles.\n",
    "- The $|m_F=+2\\rangle \\to |m_{F'}=+3\\rangle$ sub-transition is the\n",
    "  **true cycling transition**: $m_{F'}=+3$ can only decay to $m_F=+2$\n",
    "  (the only allowed $\\Delta m_F \\in \\{-1, 0, +1\\}$ channel from $m_{F'}=+3$\n",
    "  to $F=2$), making atoms permanently trapped in the cycling loop.\n",
    "- The `hyperfine.Atom1e` module constructs the full $m_F$ sublevel\n",
    "  structure with correct Clebsch-Gordan-weighted coupling and decay\n",
    "  factors, valid for any alkali $nS_{1/2} \\to nP_J$ transition.\n",
    "\n",
    "## References\n",
    "\n",
    "1. D. A. Steck, *Rubidium 87 D Line Data* (2001, revision 2.2.2, 2021).\n",
    "   Comprehensive Rb 87 spectroscopic data including D2 hyperfine constants,\n",
    "   transition dipole elements, and natural linewidth. Available at\n",
    "   https://steck.us/alkalidata.\n",
    "2. C. J. Foot, *Atomic Physics*, Oxford University Press (2004).\n",
    "   Chapter 7: laser cooling and optical pumping. Cycling transitions and\n",
    "   the MOT.\n",
    "3. H. J. Metcalf and P. van der Straten, *Laser Cooling and Trapping*,\n",
    "   Springer (1999). Optical pumping and rate equations, Chapter 4."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}