{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Counter-Propagating Fields\n",
    "\n",
    "Many atomic physics experiments use fields that propagate in opposite directions through\n",
    "the medium. A common example is Rydberg EIT, where a near-infrared probe and a blue\n",
    "coupling field enter the vapour cell from opposite ends to reduce (or eliminate) the\n",
    "Doppler broadening of the two-photon resonance.\n",
    "\n",
    "## Why a counter-propagating geometry changes the Doppler shift\n",
    "\n",
    "An atom moving with velocity $v$ along the propagation axis sees each field at a\n",
    "Doppler-shifted frequency. For a field with wave-vector $\\mathbf{k}$:\n",
    "\n",
    "- **Co-propagating** ($\\mathbf{k} \\parallel \\hat{z}$): the atom sees the field\n",
    "  blue-shifted by $kv$, so its effective one-photon detuning shifts by $+kv$.\n",
    "- **Counter-propagating** ($\\mathbf{k} \\antiparallel \\hat{z}$): the shift is\n",
    "  $-kv$.\n",
    "\n",
    "In a two-photon (Ξ or Λ) scheme, the _two-photon_ detuning is the sum of the\n",
    "individual detunings. When both fields propagate in the same direction, the\n",
    "two-photon detuning picks up $+kv + kv = 2kv$, and the resonance is\n",
    "Doppler-broadened. When the fields counter-propagate with similar wavelengths, the\n",
    "shifts partially or fully cancel, giving a narrower (Doppler-reduced or\n",
    "Doppler-free) two-photon resonance.\n",
    "\n",
    "## How MaxwellBloch handles this\n",
    "\n",
    "MaxwellBloch propagates all fields in the $+z$ direction (it solves an\n",
    "initial-value problem marching from $z_\\mathrm{min}$ to $z_\\mathrm{max}$). A\n",
    "genuinely counter-propagating field would require boundary data at $z_\\mathrm{max}$\n",
    "and backward integration, turning the problem into a boundary-value problem.\n",
    "\n",
    "However, when the counter-propagating field **does not deplete significantly**\n",
    "across the cell—a valid approximation for many coupling fields in EIT\n",
    "experiments—the forward-propagating solver is equivalent to the true\n",
    "counter-propagating geometry. The only physical difference is the **sign of the\n",
    "Doppler shift** each velocity class contributes to that field. MaxwellBloch exposes\n",
    "this sign via the `counter_propagating` field attribute.\n",
    "\n",
    "> **Note on phrasing.** This is *not* an 'undepleted-coupling approximation' in the\n",
    "> sense of neglecting absorption. The solver computes absorption honestly for every\n",
    "> velocity class. The approximation is purely geometric: it exploits the spatial\n",
    "> symmetry that holds when the counter-propagating field's intensity profile is\n",
    "> essentially flat across the cell."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic usage: `counter_propagating`\n",
    "\n",
    "Set `\"counter_propagating\": true` on any field in the JSON configuration. This\n",
    "tells MaxwellBloch to apply the Doppler shift with the opposite sign for that field,\n",
    "while leaving the forward-propagation geometry unchanged."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"num_states\": 3,\n",
    "    \"decays\": [{\"channels\": [[0, 1], [1, 2]], \"rate\": 1.0}],\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"label\": \"probe\",\n",
    "        \"coupled_levels\": [[0, 1]],\n",
    "        \"rabi_freq\": 1.0e-3,\n",
    "        \"rabi_freq_t_func\": \"gaussian\",\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"centre\": 0.0, \"fwhm\": 1.0}\n",
    "      },\n",
    "      {\n",
    "        \"label\": \"coupling\",\n",
    "        \"coupled_levels\": [[1, 2]],\n",
    "        \"rabi_freq\": 5.0,\n",
    "        \"rabi_freq_t_func\": \"ramp_onoff\",\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"fwhm\": 0.2, \"on\": -2.0, \"off\": 8.0},\n",
    "        \"counter_propagating\": true\n",
    "      }\n",
    "    ]\n",
    "  },\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 8.0,\n",
    "  \"t_steps\": 120,\n",
    "  \"z_min\": 0.0,\n",
    "  \"z_max\": 1.0,\n",
    "  \"z_steps\": 20,\n",
    "  \"z_steps_inner\": 2,\n",
    "  \"interaction_strengths\": [1.0, 1.0],\n",
    "  \"velocity_classes\": {\n",
    "    \"thermal_width\": 5.0,\n",
    "    \"thermal_delta_min\": -15.0,\n",
    "    \"thermal_delta_max\": 15.0,\n",
    "    \"thermal_delta_steps\": 10\n",
    "  },\n",
    "  \"savefile\": \"counter-propagating-eit\"\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from maxwellbloch import mb_solve\n",
    "\n",
    "mbs = mb_solve.MBSolve().from_json_str(mb_solve_json)\n",
    "print(f\"Coupling field counter_propagating: {mbs.atom.fields[1].counter_propagating}\")\n",
    "print(f\"Coupling factor_doppler_shift:      {mbs.atom.fields[1].factor_doppler_shift}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Omegas_zt, states_zt = mbs.mbsolve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Depletion check\n",
    "\n",
    "After solving, MaxwellBloch automatically checks whether the counter-propagating\n",
    "field depleted significantly across the cell. The check computes\n",
    "\n",
    "$$\n",
    "\\text{depletion} = 1 - \\frac{\\max_t |\\Omega(z_\\mathrm{max}, t)|}{\\max_t |\\Omega(z_\\mathrm{min}, t)|}\n",
    "$$\n",
    "\n",
    "- **< 1 %** — check passes silently. The forward-propagation result is reliable.\n",
    "- **1 – 10 %** — a `UserWarning` is emitted. Results may still be usable but\n",
    "  caution is advised.\n",
    "- **> 10 %** — a `CounterPropagatingDepletionError` is raised. The approximation\n",
    "  has broken down.\n",
    "\n",
    "You can suppress the check with `check_counter_prop_depletion=False` if you know\n",
    "the depletion is acceptable for your purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "coupling_idx = 1  # second field\n",
    "Omega_coupling = Omegas_zt[coupling_idx]\n",
    "peak_in  = np.max(np.abs(Omega_coupling[0]))\n",
    "peak_out = np.max(np.abs(Omega_coupling[-1]))\n",
    "depletion = 1.0 - peak_out / peak_in\n",
    "print(f\"Coupling field depletion: {depletion:.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced: `factor_doppler_shift`\n",
    "\n",
    "For schemes where the two counter-propagating fields have **different wavelengths**\n",
    "(e.g. a three-photon Doppler-free ladder where $k_1 \\neq k_2$), the Doppler\n",
    "cancellation is only partial. Use `factor_doppler_shift` directly to set the\n",
    "fractional weight:\n",
    "\n",
    "```json\n",
    "{\n",
    "  \"coupled_levels\": [[2, 3]],\n",
    "  \"rabi_freq\": 1.5,\n",
    "  \"factor_doppler_shift\": -0.62\n",
    "}\n",
    "```\n",
    "\n",
    "The sign determines the direction; the magnitude scales the Doppler shift relative\n",
    "to the thermal width unit. `factor_doppler_shift = 1.0` (the default) gives\n",
    "standard co-propagating behaviour; `-1.0` is equivalent to\n",
    "`counter_propagating = true`.\n",
    "\n",
    "> **Note on `detuning_positive`.** The `detuning_positive` flag on a field is a\n",
    "> **separate and orthogonal** concept. It controls the sign convention for the RWA\n",
    "> Hamiltonian based on the level ordering (used when the second coupled level sits\n",
    "> below the first in energy, e.g. the lower leg of a Λ system). It has nothing to\n",
    "> do with the laser k-vector direction. Both flags can be set independently."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pulse-area check\n",
    "\n",
    "As a sanity check, verify that the probe pulse area at the input is correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "probe_area_in = np.trapezoid(np.abs(Omegas_zt[0, 0, :]), mbs.tlist) / np.pi\nprint(f\"Probe area at z=z_min: {probe_area_in:.4f} π\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "1. Mohapatra, Jackson, Adams — *Coherent Optical Detection of Highly Excited\n",
    "   Rydberg States Using Electromagnetically Induced Transparency*,\n",
    "   PRL **98**, 113003 (2007). Example of the counter-propagating Rydberg EIT\n",
    "   geometry this feature supports.\n",
    "2. Adams, Pritchard, Shaffer — *Rydberg atom physics*,\n",
    "   J. Phys. B **53**, 012002 (2020). Broader review including Doppler-free\n",
    "   multi-photon spectroscopy."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}