{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro",
   "metadata": {},
   "source": [
    "# Plotting\n",
    "\n",
    "`maxwellbloch.plot` provides Plotly-based visualisation primitives for\n",
    "MBSolve and OBSolve results. Figures are interactive in the browser —\n",
    "hover to read values, drag to zoom, click the legend to toggle traces.\n",
    "\n",
    "## Installation\n",
    "\n",
    "The plotting module requires the optional `[plot]` dependencies:\n",
    "\n",
    "```bash\n",
    "pip install maxwellbloch[plot]\n",
    "```\n",
    "\n",
    "or with uv:\n",
    "\n",
    "```bash\n",
    "uv pip install maxwellbloch[plot]\n",
    "```\n",
    "\n",
    "`kaleido` is included in `[plot]` and is required for static PNG export\n",
    "(`fig.write_image()`). If you only need interactive figures in a notebook\n",
    "you can install `plotly` alone.\n",
    "\n",
    "## Usage pattern\n",
    "\n",
    "All plot functions take a solved `MBSolve` instance as their first argument\n",
    "and return a `plotly.graph_objects.Figure`. They never call `.show()`\n",
    "internally — that is always your call:\n",
    "\n",
    "```python\n",
    "from maxwellbloch import plot\n",
    "\n",
    "fig = plot.field_spacetime(mbs)         # returns a Figure\n",
    "fig.show(renderer='notebook_connected') # interactive in Jupyter\n",
    "fig.write_image('output.png')           # static PNG via kaleido\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "setup-header",
   "metadata": {},
   "source": [
    "## Setup: solve a two-level absorber\n",
    "\n",
    "We use a weak Gaussian probe propagating through a two-level absorber.\n",
    "This is enough to demonstrate every primitive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup",
   "metadata": {},
   "outputs": [],
   "source": [
    "from maxwellbloch import mb_solve, plot\n",
    "\n",
    "mb_solve_json = \"\"\"\n",
    "{\n",
    "  \"atom\": {\n",
    "    \"num_states\": 2,\n",
    "    \"decays\": [{\"channels\": [[0, 1]], \"rate\": 1.0}],\n",
    "    \"fields\": [\n",
    "      {\n",
    "        \"label\": \"probe\",\n",
    "        \"coupled_levels\": [[0, 1]],\n",
    "        \"rabi_freq\": 0.01,\n",
    "        \"rabi_freq_t_func\": \"gaussian\",\n",
    "        \"rabi_freq_t_args\": {\"ampl\": 1.0, \"centre\": 0.0, \"fwhm\": 1.0}\n",
    "      }\n",
    "    ]\n",
    "  },\n",
    "  \"t_min\": -3.0,\n",
    "  \"t_max\": 6.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\": [2.0],\n",
    "  \"savefile\": \"plotting-two-level\"\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "mbs = mb_solve.MBSolve.from_json_str(mb_solve_json)\n",
    "Omegas_zt, states_zt = mbs.mbsolve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spacetime-header",
   "metadata": {},
   "source": [
    "## `field_spacetime` — space-time heatmap of |Ω(z, t)|\n",
    "\n",
    "The most common view: the full propagation in one figure.\n",
    "Hover over any pixel to read the exact (z, t, |Ω|) values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "spacetime",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.field_spacetime(mbs)\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "envelope-header",
   "metadata": {},
   "source": [
    "## `field_envelope` — pulse shape at fixed z\n",
    "\n",
    "Compare the input pulse (z = z_min) to the output pulse (z = z_max).\n",
    "The probe is absorbed as it propagates through the medium."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "envelope",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.field_envelope(mbs, z_indices=[0, -1])\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b26c63a",
   "source": "## `field_z_profile_anim` — animated spatial profile\n\nWatch the pulse propagate through the medium. Each frame shows\nthe field envelope $|\\Omega(z)|$ at a single instant. Use the slider to scrub\nto a specific time, or press **▶ Play** to animate.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "id": "d8633d11",
   "source": "fig = plot.field_z_profile_anim(mbs)\nfig.show(renderer='notebook_connected')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "pulse-area-header",
   "metadata": {},
   "source": [
    "## `pulse_area` — area theorem\n",
    "\n",
    "The pulse area $\\mathcal{A} = \\int |\\Omega(z,t)|\\, dt / \\pi$ decreases\n",
    "from the input value as the probe is absorbed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pulse-area",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.pulse_area(mbs)\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spectrum-header",
   "metadata": {},
   "source": "## `spectrum` — frequency-domain absorption\n\nComputes the absorption spectrum via Fourier transform of the time-domain\nfield at `z_idx` (default: exit face). This works in the linear (weak-field)\nregime; for strong fields the result is the nonlinear transmission, not the\nsusceptibility.\n\nKey options:\n\n- `freq_range` — clip the displayed frequency axis to `|f| ≤ freq_range` (γ),\n  hiding the noisy high-frequency wings.\n- `show_dispersion=True` — add the dispersive component on a secondary axis.\n- `freq_scale=\"arcsinh\"` — compress the frequency axis nonlinearly (linear\n  near zero, log-like at large |f|) so narrow and broad features coexist.\n  `arcsinh_scale` sets the width of the linear region in γ.\n- `window` — apply a SciPy window (e.g. `\"hann\"`) to the time-domain field\n  before the FFT to reduce spectral leakage from truncated free-induction decay."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "spectrum",
   "metadata": {},
   "outputs": [],
   "source": "# Narrow window around the resonance with dispersion on secondary axis\nfig = plot.spectrum(mbs, freq_range=3.0, show_dispersion=True)\nfig.show(renderer='notebook_connected')"
  },
  {
   "cell_type": "code",
   "id": "0a9f87dd",
   "source": "# Arcsinh frequency scale (linear ±1 γ, log-like beyond) with Hann windowing\n# to reduce leakage from the free-induction decay tail\nfig = plot.spectrum(mbs, freq_scale=\"arcsinh\", arcsinh_scale=1.0, window=\"hann\")\nfig.show(renderer='notebook_connected')",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "id": "population-header",
   "metadata": {},
   "source": [
    "## `population` — state populations vs time\n",
    "\n",
    "The diagonal elements $\\rho_{ii}(t)$ of the density matrix at a fixed\n",
    "z position. The probe drives population from |0⟩ to |1⟩ and back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "population",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.population(mbs, z_idx=0)\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pop-spacetime-header",
   "metadata": {},
   "source": [
    "## `population_spacetime` — population heatmap\n",
    "\n",
    "Like `field_spacetime` but for a density-matrix diagonal element.\n",
    "Useful for visualising population inversion, storage, and retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pop-spacetime",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.population_spacetime(mbs, state_idx=1)\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "coherence-header",
   "metadata": {},
   "source": [
    "## `coherence` — off-diagonal density matrix elements\n",
    "\n",
    "Plot $|\\rho_{01}(t)|$, $\\mathrm{Re}(\\rho_{01})$, or $\\mathrm{Im}(\\rho_{01})$\n",
    "at a fixed z. The coherence is proportional to the macroscopic polarisation\n",
    "of the medium and drives the field evolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "coherence",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plot.coherence(mbs, 0, 1, z_idx=0, component='abs')\n",
    "fig.show(renderer='notebook_connected')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "static-export-header",
   "metadata": {},
   "source": [
    "## Static PNG export\n",
    "\n",
    "Any figure can be saved as a static PNG using `fig.write_image()`. This\n",
    "requires `kaleido` (included in `maxwellbloch[plot]`).\n",
    "\n",
    "```python\n",
    "fig = plot.field_spacetime(mbs)\n",
    "fig.write_image('field_spacetime.png', scale=2)  # scale=2 for high-DPI\n",
    "```\n",
    "\n",
    "SVG and PDF are also supported:\n",
    "\n",
    "```python\n",
    "fig.write_image('field_spacetime.svg')\n",
    "fig.write_image('field_spacetime.pdf')\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.14.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}