{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Custom Callbacks When Running TARDIS"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function `run_tardis` allows users to provide a set of callbacks to the simulation. These callbacks are functions that will be run at the end of each iteration, and can do a variety of things, such as printing information about the simulation, storing data to a table, or even changing simulation parameters between iterations. This tutorial will show three examples of callbacks and how they can be used in TARDIS. One important thing to note is that the first argument of the callback must be the `Simulation` object being run.\n",
    "\n",
    "Our first callback example will compute the (volume-weighted) average radiative temperature in the supernova ejecta (outside of the photosphere) and will print its value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_temp(sim):\n",
    "    t_rads = sim.simulation_state.t_radiative\n",
    "    volumes = sim.simulation_state.volume\n",
    "    avg = sum(t_rads*volumes) / sum(volumes)\n",
    "    print(f\"Average temperature for iteration {sim.iterations_executed}: {avg}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we give the callback to `run_tardis`.  `run_tardis` offers the keyword argument `simulation_callbacks` which takes a list of lists containing the callback as well as any optional arguments you wish to include with your callback.  For this example our function requires no extra arguments and we only have a single callback, so we give `run_tardis` a 2D list containing the callback as its only element:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We filter out warnings in this notebook\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from tardis import run_tardis\n",
    "from tardis.io.atom_data.util import download_atom_data\n",
    "\n",
    "# We download the atomic data needed to run the simulation\n",
    "download_atom_data('kurucz_cd23_chianti_H_He')\n",
    "\n",
    "# We run the simulation with our callback\n",
    "sim = run_tardis('tardis_example.yml', \n",
    "                 simulation_callbacks=[[average_temp]])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running Callbacks with Extra Arguments\n",
    "\n",
    "The callbacks provided to `run_tardis` can also take extra arguments. As an example, we'll make a callback that appends the number of monte carlo packets emitted by the supernova for each iteration to a list so we can plot the number of emitted packets for each iteration. We will also specify that we want this information for all but the last iteration, as more packets are used in the last iteration than are used in the other iterations. The callback will take a list we want to append to as an argument. We'll send both this new callback and our original `average_temp` callback to `run_tardis` as an example of using multiple callbacks at once:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def append_num_emitted_to_list(sim, lst):\n",
    "    if sim.iterations_executed < sim.iterations:\n",
    "        num_emitted_packets = len(sim.transport.transport_state.emitted_packet_nu)\n",
    "        lst.append(num_emitted_packets)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to add our new callback, we just create another entry in our list of callbacks. Since `append_num_emitted_to_list` takes an extra argument, we will provide that argument in the inner list containing the callback:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a list to store the number of emitted packets\n",
    "num_emitted_list = []\n",
    "\n",
    "# Make our list of callbacks\n",
    "callbacks = [[average_temp], \n",
    "             [append_num_emitted_to_list, num_emitted_list]]\n",
    "\n",
    "# Run the simulation with both of our callbacks\n",
    "sim = run_tardis('tardis_example.yml', \n",
    "                 simulation_callbacks=callbacks)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can look at how many packets are emitted after each iteration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Generate a list of each iteration number for the x-axis\n",
    "iterations = list(range(1, len(num_emitted_list)+1))\n",
    "\n",
    "# Plot the number of emitted packets\n",
    "plt.plot(iterations, num_emitted_list)\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Number of emitted packets\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Callbacks to Add New Functionality\n",
    "\n",
    "Callbacks can also add new functionality to the code. For example, we introduce one final callback `inc_packets` that will increase the number of packets in the following iteration by a number $N$ (which is an argument to the callback, in our example we shall use $N=1000$):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inc_packets(sim, N):\n",
    "    sim.no_of_packets += N"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a new list to store the number of emitted packets\n",
    "num_emitted_list_new = []\n",
    "\n",
    "# Make our new list of callbacks\n",
    "callbacks = [[average_temp],\n",
    "             [append_num_emitted_to_list, num_emitted_list_new], \n",
    "             [inc_packets, 1000]]\n",
    "\n",
    "# Run the simulation with all three of our callbacks\n",
    "sim = run_tardis('tardis_example.yml', \n",
    "                 simulation_callbacks=callbacks)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's see how this affected our plot for packets emitted in each iteration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(iterations, num_emitted_list_new)\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"Number of emitted packets\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the number of packets emitted will keep on increasing as 1000 more packets are run each iteration."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}