{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro",
   "metadata": {},
   "source": [
    "# Λ-Type Three-Level: EIT, Slow Light and Storage\n",
    "\n",
    "In the **Λ-type three-level system** a weak probe field couples the ground\n",
    "state |0⟩ to the excited state |1⟩, and a strong coupling field couples\n",
    "|1⟩ to a second ground state |2⟩. When the two-photon resonance condition\n",
    "is met, quantum interference between the two excitation pathways suppresses\n",
    "absorption of the probe — **electromagnetically induced transparency** (EIT).\n",
    "\n",
    "Inside the EIT window the probe propagates with a greatly reduced group\n",
    "velocity $v_g \\ll c$, compressing the pulse spatially by the same factor.\n",
    "Adiabatically switching the coupling field off traps the probe as a spin-wave\n",
    "coherence; switching it back on retrieves the pulse — **light storage**.\n",
    "\n",
    "This notebook demonstrates:\n",
    "\n",
    "1. **No coupling** — probe absorbed as in a two-level medium.\n",
    "2. **With coupling** — EIT transparency and slow light.\n",
    "3. **Gaussian atom cloud** — slow light and spatial pulse compression.\n",
    "4. **Light storage and retrieval** — coupling switched off then on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "\n",
    "from maxwellbloch import mb_solve\n",
    "\n",
    "sns.set_style(\"darkgrid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sec-no-coupling-intro",
   "metadata": {},
   "source": [
    "## No coupling — resonant absorption\n",
    "\n",
    "With the coupling field amplitude set to zero the Λ system reduces to a\n",
    "two-level atom. The weak Gaussian probe ($\\Omega_0 = 10^{-3}\\,\\Gamma$) is\n",
    "resonantly absorbed as it propagates through the dense medium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-no-coupling",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json_no_coupling = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"coupled_levels\": [[0, 1]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": true,\n",
    "        \"label\": \"probe\",\n",
    "        \"rabi_freq\": 1.0e-3,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"centre\": 0.0, \"fwhm\": 1.0},\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      },\n",
    "      {\n",
    "        \"coupled_levels\": [[1, 2]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": false,\n",
    "        \"label\": \"coupling\",\n",
    "        \"rabi_freq\": 0.0,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"fwhm\": 0.2, \"on\": -1.0, \"off\": 9.0},\n",
    "        \"rabi_freq_t_func\": \"ramp_onoff\"\n",
    "      }\n",
    "    ],\n",
    "    \"num_states\": 3\n",
    "  },\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 10.0,\n",
    "  \"t_steps\": 120,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 100,\n",
    "  \"z_steps_inner\": 2,\n",
    "  \"interaction_strengths\": [10.0, 10.0],\n",
    "  \"savefile\": \"mbs-lambda-weak-pulse-more-atoms-no-coupling\"\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "mbs_no_coupling = mb_solve.MBSolve().from_json_str(mb_solve_json_no_coupling)\n",
    "mbs_no_coupling.mbsolve(recalc=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-no-coupling",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 12))\n",
    "\n",
    "ax = fig.add_subplot(211)\n",
    "cf = ax.contourf(mbs_no_coupling.tlist, mbs_no_coupling.zlist,\n",
    "                 np.abs(mbs_no_coupling.Omegas_zt[0] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 1.0e-3, 11), cmap=plt.cm.Blues)\n",
    "ax.set_title(r\"Rabi Frequency ($\\Gamma / 2\\pi$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Probe\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c=\"grey\", lw=1.0, ls=\"dotted\")\n",
    "\n",
    "ax = fig.add_subplot(212)\n",
    "cf = ax.contourf(mbs_no_coupling.tlist, mbs_no_coupling.zlist,\n",
    "                 np.abs(mbs_no_coupling.Omegas_zt[1] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 10.0, 11), cmap=plt.cm.Greens)\n",
    "ax.set_xlabel(r\"Time ($1/\\Gamma$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Coupling\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c=\"grey\", lw=1.0, ls=\"dotted\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sec-with-coupling-intro",
   "metadata": {},
   "source": [
    "## With coupling — EIT transparency and slow light\n",
    "\n",
    "Switching the coupling on ($\\Omega_c = 5\\,\\Gamma$) opens a narrow\n",
    "transparency window at two-photon resonance. The probe now passes through\n",
    "the medium with only a small reduction in amplitude, but its group\n",
    "velocity is greatly reduced — it arrives significantly later than a\n",
    "pulse propagating in vacuum would."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-with-coupling",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json_with_coupling = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"coupled_levels\": [[0, 1]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": true,\n",
    "        \"label\": \"probe\",\n",
    "        \"rabi_freq\": 1.0e-3,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"centre\": 0.0, \"fwhm\": 1.0},\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      },\n",
    "      {\n",
    "        \"coupled_levels\": [[1, 2]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": false,\n",
    "        \"label\": \"coupling\",\n",
    "        \"rabi_freq\": 5.0,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"fwhm\": 0.2, \"on\": -1.0, \"off\": 9.0},\n",
    "        \"rabi_freq_t_func\": \"ramp_onoff\"\n",
    "      }\n",
    "    ],\n",
    "    \"num_states\": 3\n",
    "  },\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 10.0,\n",
    "  \"t_steps\": 120,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 100,\n",
    "  \"z_steps_inner\": 2,\n",
    "  \"interaction_strengths\": [10.0, 10.0],\n",
    "  \"savefile\": \"mbs-lambda-weak-pulse-more-atoms-with-coupling\"\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "mbs_with_coupling = mb_solve.MBSolve().from_json_str(mb_solve_json_with_coupling)\n",
    "mbs_with_coupling.mbsolve(recalc=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-with-coupling",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 12))\n",
    "\n",
    "ax = fig.add_subplot(211)\n",
    "cf = ax.contourf(mbs_with_coupling.tlist, mbs_with_coupling.zlist,\n",
    "                 np.abs(mbs_with_coupling.Omegas_zt[0] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 1.0e-3, 11), cmap=plt.cm.Blues)\n",
    "ax.set_title(r\"Rabi Frequency ($\\Gamma / 2\\pi$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Probe\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c=\"grey\", lw=1.0, ls=\"dotted\")\n",
    "\n",
    "ax = fig.add_subplot(212)\n",
    "cf = ax.contourf(mbs_with_coupling.tlist, mbs_with_coupling.zlist,\n",
    "                 np.abs(mbs_with_coupling.Omegas_zt[1] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 10.0, 11), cmap=plt.cm.Greens)\n",
    "ax.set_xlabel(r\"Time ($1/\\Gamma$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Coupling\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "for y in [0.0, 1.0]:\n",
    "    ax.axhline(y, c=\"grey\", lw=1.0, ls=\"dotted\")\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sec-cloud-intro",
   "metadata": {},
   "source": [
    "## Gaussian atom cloud — slow light and pulse compression\n",
    "\n",
    "Replacing the uniform medium with a Gaussian density profile makes the\n",
    "slow-light effect spatially inhomogeneous: the probe decelerates as it\n",
    "enters the dense region and reaccelerates as it exits, arriving well\n",
    "after the vacuum transit time.\n",
    "\n",
    "Because the leading edge of the pulse slows before the trailing edge, the\n",
    "pulse is spatially compressed as it enters the medium — compressed by the\n",
    "same factor $v_g / c$ by which it is slowed. In EIT experiments with\n",
    "ultracold gases, this compression reduces kilometre-scale pulses to\n",
    "sub-millimetre spatial extent inside the cloud."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-cloud",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json_cloud = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"coupled_levels\": [[0, 1]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": true,\n",
    "        \"label\": \"probe\",\n",
    "        \"rabi_freq\": 1.0e-3,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"centre\": 0.0, \"fwhm\": 1.0},\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      },\n",
    "      {\n",
    "        \"coupled_levels\": [[1, 2]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": false,\n",
    "        \"label\": \"coupling\",\n",
    "        \"rabi_freq\": 5.0,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"fwhm\": 0.2, \"on\": -1.0, \"off\": 9.0},\n",
    "        \"rabi_freq_t_func\": \"ramp_onoff\"\n",
    "      }\n",
    "    ],\n",
    "    \"num_states\": 3\n",
    "  },\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 10.0,\n",
    "  \"t_steps\": 120,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 70,\n",
    "  \"z_steps_inner\": 100,\n",
    "  \"num_density_z_func\": \"gaussian\",\n",
    "  \"num_density_z_args\": {\"ampl\": 1.0, \"fwhm\": 0.5, \"centre\": 0.5},\n",
    "  \"interaction_strengths\": [1.0e3, 1.0e3],\n",
    "  \"savefile\": \"mbs-lambda-weak-pulse-cloud-atoms-some-coupling\"\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "mbs_cloud = mb_solve.MBSolve().from_json_str(mb_solve_json_cloud)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-cloud-density",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(16, 3))\n",
    "ax.plot(mbs_cloud.zlist, mbs_cloud.num_density_z_func(mbs_cloud.zlist, mbs_cloud.num_density_z_args))\n",
    "ax.set_xlabel(\"Distance ($L$)\")\n",
    "ax.set_ylabel(\"Number density\")\n",
    "for y_val in [0.0, 1.0]:\n",
    "    ax.axvline(y_val, c=\"grey\", lw=1.0, ls=\"dotted\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "solve-cloud",
   "metadata": {},
   "outputs": [],
   "source": [
    "mbs_cloud.mbsolve(recalc=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-cloud",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 12))\n",
    "\n",
    "ax = fig.add_subplot(211)\n",
    "cf = ax.contourf(mbs_cloud.tlist, mbs_cloud.zlist,\n",
    "                 np.abs(mbs_cloud.Omegas_zt[0] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 1.0e-3, 11), cmap=plt.cm.Blues)\n",
    "ax.set_title(r\"Rabi Frequency ($\\Gamma / 2\\pi$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Probe\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "\n",
    "ax = fig.add_subplot(212)\n",
    "cf = ax.contourf(mbs_cloud.tlist, mbs_cloud.zlist,\n",
    "                 np.abs(mbs_cloud.Omegas_zt[1] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 8.0, 11), cmap=plt.cm.Greens)\n",
    "ax.set_xlabel(r\"Time ($1/\\Gamma$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Coupling\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "\n",
    "for ax in fig.axes:\n",
    "    for y in [0.0, 1.0]:\n",
    "        ax.axhline(y, c=\"grey\", lw=1.0, ls=\"dotted\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sec-store-intro",
   "metadata": {},
   "source": [
    "## Light storage and retrieval\n",
    "\n",
    "The coupling field is now switched **off** at $t = 4$ (once the probe is\n",
    "inside the cloud) and switched back **on** at $t = 6$. Ramping the coupling\n",
    "to zero rotates the dark-state mixing angle to $\\theta \\to \\pi/2$, mapping\n",
    "the probe coherently onto a long-lived spin-wave excitation of the medium.\n",
    "Restoring the coupling reverses the rotation and the probe re-emerges with\n",
    "its original profile.\n",
    "\n",
    "> *Note: this is the most computationally expensive case in this notebook\n",
    "> (140 × 50 inner steps). On older hardware it can take several hours.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup-store",
   "metadata": {},
   "outputs": [],
   "source": [
    "mb_solve_json_store = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"coupled_levels\": [[0, 1]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"label\": \"probe\",\n",
    "        \"rabi_freq\": 1.0e-3,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"centre\": 0.0, \"fwhm\": 1.0},\n",
    "        \"rabi_freq_t_func\": \"gaussian\"\n",
    "      },\n",
    "      {\n",
    "        \"coupled_levels\": [[1, 2]],\n",
    "        \"detuning\": 0.0,\n",
    "        \"detuning_positive\": false,\n",
    "        \"label\": \"coupling\",\n",
    "        \"rabi_freq\": 5.0,\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"fwhm\": 0.2, \"off\": 4.0, \"on\": 6.0},\n",
    "        \"rabi_freq_t_func\": \"ramp_offon\"\n",
    "      }\n",
    "    ],\n",
    "    \"num_states\": 3\n",
    "  },\n",
    "  \"t_min\": -2.0,\n",
    "  \"t_max\": 12.0,\n",
    "  \"t_steps\": 140,\n",
    "  \"z_min\": -0.2,\n",
    "  \"z_max\": 1.2,\n",
    "  \"z_steps\": 140,\n",
    "  \"z_steps_inner\": 50,\n",
    "  \"num_density_z_func\": \"gaussian\",\n",
    "  \"num_density_z_args\": {\"ampl\": 1.0, \"fwhm\": 0.5, \"centre\": 0.5},\n",
    "  \"interaction_strengths\": [1.0e3, 1.0e3],\n",
    "  \"savefile\": \"mbs-lambda-weak-pulse-cloud-atoms-some-coupling-store\"\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "mbs_store = mb_solve.MBSolve().from_json_str(mb_solve_json_store)\n",
    "mbs_store.mbsolve(recalc=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plot-store",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(16, 12))\n",
    "\n",
    "ax = fig.add_subplot(211)\n",
    "cf = ax.contourf(mbs_store.tlist, mbs_store.zlist,\n",
    "                 np.abs(mbs_store.Omegas_zt[0] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 1.0e-3, 11), cmap=plt.cm.Blues)\n",
    "ax.set_title(r\"Rabi Frequency ($\\Gamma / 2\\pi$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Probe\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "\n",
    "ax = fig.add_subplot(212)\n",
    "cf = ax.contourf(mbs_store.tlist, mbs_store.zlist,\n",
    "                 np.abs(mbs_store.Omegas_zt[1] / (2 * np.pi)),\n",
    "                 np.linspace(0.0, 8.0, 11), cmap=plt.cm.Greens)\n",
    "ax.set_xlabel(r\"Time ($1/\\Gamma$)\")\n",
    "ax.set_ylabel(\"Distance ($L$)\")\n",
    "ax.text(0.02, 0.95, \"Coupling\", va=\"top\", ha=\"left\", transform=ax.transAxes, color=\"grey\", fontsize=16)\n",
    "plt.colorbar(cf)\n",
    "\n",
    "for ax in fig.axes:\n",
    "    for y in [0.0, 1.0]:\n",
    "        ax.axhline(y, c=\"grey\", lw=1.0, ls=\"dotted\")\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}