{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d5f95feb",
   "metadata": {},
   "source": [
    "# Convergence Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a9733ff",
   "metadata": {},
   "source": [
    "These convergence plots help diagnose the quality of TARDIS simulation results:\n",
    "\n",
    "The **plasma plots** show radiation temperature and dilution factor across velocity space, indicating when the radiation field reaches equilibrium throughout the ejecta. When lines stop changing between iterations, physical conditions have converged.\n",
    "\n",
    "The **luminosity plots** track:\n",
    "- Inner boundary temperature stabilization\n",
    "- Whether emitted luminosity matches requested luminosity (energy conservation)\n",
    "- Residual luminosity percentage (ideally <5%)\n",
    "\n",
    "Properly converged plots indicate your synthetic spectrum is based on a physically self-consistent model, making it reliable for comparison with observed supernova spectra. Poor convergence suggests you may need to adjust model parameters or extend the number of iterations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc1a0c1f",
   "metadata": {},
   "source": [
    "The Convergence Plots consist of two Plotly FigureWidget Subplots, the `plasma_plot` and the `t_inner_luminosities_plot`. The plots can be displayed by setting the `show_convergence_plots` option in the `run_tardis` function to `True`. The plots are stored in the `convergence_plots` attribute of the simulation object `sim` and can be accessed using `sim.convergence_plots.plasma_plot` and `sim.convergence_plots.t_inner_luminosities_plot`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6db4edf2",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "Note\n",
    "    \n",
    "You only need to include `export_convergence_plots=True` in the `run_tardis` function when you want to share the notebook. The function shows the plot using the Plotly `notebook_connected` renderer, which helps display the plot online. You don't need to do it when running the notebook locally.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e625935a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import plotly.graph_objects as go\n",
    "from astropy import units as u\n",
    "\n",
    "import tardis.visualization.plot_util as pu\n",
    "from tardis import run_tardis\n",
    "from tardis.io.atom_data import download_atom_data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fb6e7a0",
   "metadata": {},
   "source": [
    "Every simulation run requires [atomic data](io/configuration/components/atomic/atomic_data.rst) and a [configuration file](io/configuration/index.rst). \n",
    "\n",
    "## Atomic Data\n",
    "\n",
    "We recommend using the [kurucz_cd23_chianti_H_He_latest.h5](https://github.com/tardis-sn/tardis-regression-data/raw/main/atom_data/kurucz_cd23_chianti_H_He_latest.h5) dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba7cc7d2",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# We download the atomic data needed to run the simulation\n",
    "download_atom_data(\"kurucz_cd23_chianti_H_He_latest\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31b27062",
   "metadata": {},
   "source": [
    "## Example Configuration File\n",
    "\n",
    "The configuration file [tardis_example.yml](https://github.com/tardis-sn/tardis/tree/master/docs/tardis_example.yml) is used throughout this Quickstart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f410a818",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -q -nc https://raw.githubusercontent.com/tardis-sn/tardis/master/docs/tardis_example.yml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3c242e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat tardis_example.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e1f4732",
   "metadata": {},
   "source": [
    "## Callback Function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56385894",
   "metadata": {},
   "source": [
    "Callback function to extract luminosity, radiation temperature, and dilution factor at each simulation iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f418ea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "emitted_luminosity = []\n",
    "absorbed_luminosity = []\n",
    "t_rad = []\n",
    "w = []\n",
    "\n",
    "\n",
    "def fetch_simulation_properties(simulation):\n",
    "    emitted_luminosity.append(np.copy(simulation.emitted_luminosity.value))\n",
    "    absorbed_luminosity.append(np.copy(simulation.reabsorbed_luminosity.value))\n",
    "    t_rad.append(np.copy(simulation.simulation_state.t_radiative.value))\n",
    "    w.append(np.copy(simulation.simulation_state.dilution_factor))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acec6975",
   "metadata": {},
   "source": [
    "## Running the Simulation\n",
    "\n",
    "To run the simulation, import the `run_tardis` function and create the `sim` object."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87c3bc65",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**Note:**\n",
    "\n",
    "Get more information about the [progress bars](io/output/progress_bars.rst), [logging configuration](io/optional/tutorial_logging_configuration.ipynb), and [convergence plots](io/visualization/tutorial_convergence_plot.ipynb).\n",
    "\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a6dcfdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = run_tardis(\n",
    "    \"tardis_example.yml\",\n",
    "    log_level=\"ERROR\",\n",
    "    simulation_callbacks=[(fetch_simulation_properties,)],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63fb0984",
   "metadata": {},
   "source": [
    "## Preparing Luminosity and Temperature Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "562b723b",
   "metadata": {},
   "outputs": [],
   "source": [
    "luminosity_requested = sim.luminosity_requested\n",
    "velocity_km_s = np.copy(sim.simulation_state.velocity.to(u.km / u.s).value)\n",
    "luminosities = [\"Emitted\", \"Absorbed\", \"Requested\"]\n",
    "\n",
    "value_data = {\n",
    "    \"Emitted\": emitted_luminosity,\n",
    "    \"Absorbed\": absorbed_luminosity,\n",
    "    \"Requested\": [luminosity_requested.value] * sim.iterations,\n",
    "    \"t_inner\": np.copy(sim.iterations_t_inner.value),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85b96c88",
   "metadata": {},
   "source": [
    "## Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e74a778",
   "metadata": {},
   "source": [
    "### Plasma Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bde1728",
   "metadata": {},
   "outputs": [],
   "source": [
    "plasma_colorscale = pu.get_hex_color_strings(sim.iterations)\n",
    "fig = go.FigureWidget().set_subplots(rows=1, cols=2, shared_xaxes=True)\n",
    "\n",
    "# empty traces to build figure\n",
    "fig.add_scatter(row=1, col=1)\n",
    "fig.add_scatter(row=1, col=2)\n",
    "\n",
    "# 2 y axes and 2 x axes correspond to the 2 subplots in the plasma plot\n",
    "fig = fig.update_layout(\n",
    "    xaxis={\n",
    "        \"tickformat\": \"g\",\n",
    "        \"title\": pu.axis_label_in_latex(\"Velocity\", u.km / u.s),\n",
    "    },\n",
    "    xaxis2={\n",
    "        \"tickformat\": \"g\",\n",
    "        \"title\": pu.axis_label_in_latex(\"Velocity\", u.km / u.s),\n",
    "        \"matches\": \"x\",\n",
    "    },\n",
    "    yaxis={\n",
    "        \"tickformat\": \"g\",\n",
    "        \"title\": pu.axis_label_in_latex(\"T_{rad}\", u.K, only_text=False),\n",
    "        \"nticks\": 15,\n",
    "    },\n",
    "    yaxis2={\n",
    "        \"tickformat\": \"g\",\n",
    "        \"title\": r\"$W$\",\n",
    "        \"nticks\": 15,\n",
    "    },\n",
    "    height=450,\n",
    "    legend_title_text=\"Iterations\",\n",
    "    legend_traceorder=\"reversed\",\n",
    "    margin=dict(\n",
    "        l=10, r=135, b=25, t=25, pad=0\n",
    "    ),  # reduce whitespace surrounding the plot and increase right indentation to align with the t_inner and luminosity plot\n",
    ")\n",
    "\n",
    "plasma_plot = fig\n",
    "\n",
    "for iter_num in range(sim.iterations):\n",
    "    # add luminosity data in hover data in plasma plots\n",
    "    customdata = len(velocity_km_s) * [\n",
    "        (\n",
    "            \"<br>\"\n",
    "            \"Emitted Luminosity: \"\n",
    "            f\"{emitted_luminosity[iter_num]:.4g}\"\n",
    "            \"<br>\"\n",
    "            \"Requested Luminosity: \"\n",
    "            f\"{luminosity_requested:.4g}\"\n",
    "            \"<br>\"\n",
    "            \"Absorbed Luminosity: \"\n",
    "            f\"{absorbed_luminosity[iter_num]:.4g}\"\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    # add a radiation temperature vs shell velocity trace to the plasma plot\n",
    "    plasma_plot.add_scatter(\n",
    "        x=velocity_km_s,\n",
    "        y=np.append(t_rad[iter_num], t_rad[iter_num][-1:]),\n",
    "        line_color=plasma_colorscale[iter_num],\n",
    "        line_shape=\"hv\",\n",
    "        row=1,\n",
    "        col=1,\n",
    "        name=iter_num + 1,\n",
    "        legendgroup=f\"group-{iter_num}\",\n",
    "        showlegend=False,\n",
    "        customdata=customdata,\n",
    "        hovertemplate=\"<b>Y</b>: %{y:.3f} at <b>X</b> = %{x:,.0f}%{customdata}\",\n",
    "    )\n",
    "\n",
    "    # add a dilution factor vs shell velocity trace to the plasma plot\n",
    "    plasma_plot.add_scatter(\n",
    "        x=velocity_km_s,\n",
    "        y=np.append(w[iter_num], w[iter_num][-1:]),\n",
    "        line_color=plasma_colorscale[iter_num],\n",
    "        line_shape=\"hv\",\n",
    "        row=1,\n",
    "        col=2,\n",
    "        legendgroup=f\"group-{iter_num}\",\n",
    "        name=iter_num + 1,\n",
    "        customdata=customdata,\n",
    "        hovertemplate=\"<b>Y</b>: %{y:.3f} at <b>X</b> = %{x:,.0f}%{customdata}\",\n",
    "    )\n",
    "\n",
    "plasma_plot.show(renderer=\"notebook_connected\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5837fe9e",
   "metadata": {},
   "source": [
    "### Luminosity Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4a63616",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_inner_luminosities_colors = pu.get_hex_color_strings(5)\n",
    "fig = go.FigureWidget().set_subplots(\n",
    "    rows=3,\n",
    "    cols=1,\n",
    "    shared_xaxes=True,\n",
    "    vertical_spacing=0.08,\n",
    "    row_heights=[0.25, 0.5, 0.25],\n",
    ")\n",
    "\n",
    "# add inner boundary temperature vs iterations plot\n",
    "fig.add_scatter(\n",
    "    name=\"Inner<br>Boundary<br>Temperature\",\n",
    "    row=1,\n",
    "    col=1,\n",
    "    hovertext=\"text\",\n",
    "    marker_color=t_inner_luminosities_colors[0],\n",
    "    mode=\"lines+markers\",\n",
    ")\n",
    "\n",
    "# add luminosity vs iterations plot\n",
    "# has three traces for emitted, requested and absorbed luminosities\n",
    "for luminosity, line_color in zip(\n",
    "    luminosities, t_inner_luminosities_colors[1:4]\n",
    "):\n",
    "    fig.add_scatter(\n",
    "        name=luminosity + \"<br>Luminosity\",\n",
    "        mode=\"lines+markers\",\n",
    "        row=2,\n",
    "        col=1,\n",
    "        marker_color=line_color,\n",
    "    )\n",
    "\n",
    "# add residual luminosity vs iterations plot\n",
    "fig.add_scatter(\n",
    "    name=\"Residual<br>Luminosity\",\n",
    "    row=3,\n",
    "    col=1,\n",
    "    marker_color=t_inner_luminosities_colors[4],\n",
    "    mode=\"lines+markers\",\n",
    ")\n",
    "\n",
    "# 3 y axes and 3 x axes correspond to the 3 subplots in the t_inner and luminosity convergence plot\n",
    "fig = fig.update_layout(\n",
    "    xaxis=dict(range=[0, sim.iterations + 1], dtick=2),\n",
    "    xaxis2=dict(\n",
    "        matches=\"x\",\n",
    "        range=[0, sim.iterations + 1],\n",
    "        dtick=2,\n",
    "    ),\n",
    "    xaxis3=dict(\n",
    "        title=r\"$\\mbox{Iteration Number}$\",\n",
    "        dtick=2,\n",
    "    ),\n",
    "    yaxis=dict(\n",
    "        title=pu.axis_label_in_latex(\"T_{inner}\", u.K, only_text=False),\n",
    "        automargin=True,\n",
    "        tickformat=\"g\",\n",
    "        exponentformat=\"e\",\n",
    "        nticks=4,\n",
    "    ),\n",
    "    yaxis2=dict(\n",
    "        exponentformat=\"e\",\n",
    "        title=pu.axis_label_in_latex(\"Luminosity\", u.erg / u.s),\n",
    "        title_font_size=13,\n",
    "        automargin=True,\n",
    "        nticks=7,\n",
    "    ),\n",
    "    yaxis3=dict(\n",
    "        title=r\"$~~\\text{Residual}\\\\\\text{Luminosity[%]}$\",\n",
    "        title_font_size=12,\n",
    "        automargin=True,\n",
    "        nticks=4,\n",
    "    ),\n",
    "    height=630,\n",
    "    hoverlabel_align=\"right\",\n",
    "    margin=dict(b=25, t=25, pad=0),  # reduces whitespace surrounding the plot\n",
    ")\n",
    "\n",
    "t_inner_luminosities_plot = fig\n",
    "\n",
    "x = list(range(1, sim.iterations + 1))\n",
    "\n",
    "with t_inner_luminosities_plot.batch_update():\n",
    "    # traces are updated according to the order they were added\n",
    "    # the first trace is of the inner boundary temperature plot\n",
    "    t_inner_luminosities_plot.data[0].x = x\n",
    "    t_inner_luminosities_plot.data[0].y = value_data[\"t_inner\"]\n",
    "    t_inner_luminosities_plot.data[\n",
    "        0\n",
    "    ].hovertemplate = \"<b>%{y:.3f}</b> at X = %{x:,.0f}<extra>Inner Boundary Temperature</extra>\"  # trace name in extra tag to avoid new lines in hoverdata\n",
    "\n",
    "    # the next three for emitted, absorbed and requested luminosities\n",
    "    for index, luminosity in zip(range(1, 4), luminosities):\n",
    "        t_inner_luminosities_plot.data[index].x = x\n",
    "        t_inner_luminosities_plot.data[index].y = value_data[luminosity]\n",
    "        t_inner_luminosities_plot.data[index].hovertemplate = (\n",
    "            \"<b>%{y:.4g}</b>\" + \"<br>at X = %{x}<br>\"\n",
    "        )\n",
    "    # last is for the residual luminosity\n",
    "    y = [\n",
    "        ((emitted - luminosity_requested.value) * 100)\n",
    "        / luminosity_requested.value\n",
    "        for emitted in value_data[\"Emitted\"]\n",
    "    ]\n",
    "    t_inner_luminosities_plot.data[4].x = x\n",
    "    t_inner_luminosities_plot.data[4].y = y\n",
    "    t_inner_luminosities_plot.data[\n",
    "        4\n",
    "    ].hovertemplate = \"<b>%{y:.2f}%</b> at X = %{x:,.0f}\"\n",
    "\n",
    "t_inner_luminosities_plot.show(renderer=\"notebook_connected\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tardis",
   "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
}