{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Run TARDIS with a Custom Packet Source\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, TARDIS generates energy packets using its interface `BasePacketSource` class, which has a derived class `BlackBodySimpleSource`, which models the photosphere of the supernova as a perfect blackbody (see [Energy Packet Initialization](../../physics/montecarlo/initialization.ipynb)). However, users may create their own packet source, as will be shown in this notebook. In order to do this, a user must create a class (that inherits from `BasePacketSource`) and implement the following abstract functions:\n",
    "\n",
    "- create_packet_radii (returns array of packet radii)\n",
    "- create_packet_nus (returns array of packet frequencies)\n",
    "- create_packet_mus (returns array of packet directions)\n",
    "- create_packet_energies (returns array of packet energies. See [Energy Packet Initialization](../../physics/montecarlo/initialization.ipynb) for more information)\n",
    "- create_packets (wrapper which calls the above 4 functions, and is the function used by external code)\n",
    "- set_state_from_model (set the state of the source from a model object)\n",
    "\n",
    "[Note: In this notebook, we have extended the `BlackBodySimpleSource` class because it already implements some of the above functions]\n",
    "\n",
    "To use your packet source in a run of TARDIS, you must pass an instance of your class into the `run_tardis` function under the `packet_source` keyword argument.\n",
    "\n",
    ".. note:: In both the `BlackBodySimpleSource` class and in the example here, all packets are generated at the same radius. This need not be true in general (though one call of the `create_packets` method will pick the same radius from the packet source state).\n",
    "\n",
    "We show an example of how a custom packet source is used:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary packages\n",
    "import numpy as np\n",
    "from tardis import constants as const\n",
    "from astropy import units as u\n",
    "from tardis.transport.montecarlo.packet_source import BlackBodySimpleSource\n",
    "from tardis.transport.montecarlo.packet_collections import (\n",
    "    PacketCollection,\n",
    ")\n",
    "from tardis import run_tardis\n",
    "import matplotlib.pyplot as plt\n",
    "from tardis.io.atom_data import download_atom_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the atomic data used for a run of TARDIS\n",
    "download_atom_data(\"kurucz_cd23_chianti_H_He\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a packet source class that contains a create_packets method\n",
    "class TruncBlackbodySource(BlackBodySimpleSource):\n",
    "    \"\"\"\n",
    "    Custom inner boundary source class to replace the Blackbody source\n",
    "    with a truncated Blackbody source.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    truncation_wavelength : float\n",
    "        truncation wavelength in Angstrom.\n",
    "        Only wavelengths higher than the truncation wavelength\n",
    "        will be sampled.\n",
    "    radius : float64\n",
    "        Initial packet radius\n",
    "    temperature : float\n",
    "        Absolute Temperature.\n",
    "    base_seed : int\n",
    "        Base Seed for random number generator\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, truncation_wavelength=None, **kwargs):\n",
    "        self.truncation_wavelength = truncation_wavelength\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def create_packets(self, no_of_packets, drawing_sample_size=None, seed_offset=0, *args, **kwargs):\n",
    "        \"\"\"\n",
    "        Packet source that generates a truncated Blackbody source.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        no_of_packets : int\n",
    "            number of packets to be created\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        array\n",
    "            Packet radii\n",
    "        array\n",
    "            Packet frequencies\n",
    "        array\n",
    "            Packet directions\n",
    "        array\n",
    "            Packet energies\n",
    "        \"\"\"\n",
    "\n",
    "        self._reseed(self.base_seed + seed_offset)\n",
    "        packet_seeds = self.rng.choice(\n",
    "            self.MAX_SEED_VAL, no_of_packets, replace=True\n",
    "        )\n",
    "\n",
    "        # Makes uniform array of packet radii from blackbody source\n",
    "        radii = self.create_packet_radii(no_of_packets, *args, **kwargs)\n",
    "\n",
    "        # Use mus and energies from normal blackbody source.\n",
    "        mus = self.create_packet_mus(no_of_packets, *args, **kwargs)\n",
    "        energies = self.create_packet_energies(no_of_packets, *args, **kwargs)\n",
    "\n",
    "        # If not specified, draw 2 times as many packets and reject any beyond no_of_packets.\n",
    "        if drawing_sample_size is None:\n",
    "            drawing_sample_size = 2 * no_of_packets\n",
    "\n",
    "        # Blackbody will be truncated below truncation_wavelength / above truncation_frequency.\n",
    "        truncation_frequency = (\n",
    "            u.Quantity(self.truncation_wavelength, u.Angstrom)\n",
    "            .to(u.Hz, equivalencies=u.spectral())\n",
    "            .value\n",
    "        )\n",
    "\n",
    "        # Draw nus from blackbody distribution and reject based on truncation_frequency.\n",
    "        # If more nus.shape[0] > no_of_packets use only the first no_of_packets.\n",
    "        nus = self.create_packet_nus(drawing_sample_size, *args, **kwargs)\n",
    "        nus = nus[nus < truncation_frequency][:no_of_packets]\n",
    "\n",
    "        # Only required if the truncation wavelength is too big compared to the maximum\n",
    "        # of the blackbody distribution. Keep sampling until nus.shape[0] > no_of_packets.\n",
    "        while nus.shape[0] < no_of_packets:\n",
    "            additional_nus = self.create_packet_nus(drawing_sample_size, *args, **kwargs)\n",
    "            mask = additional_nus < truncation_frequency\n",
    "            additional_nus = additional_nus[mask][:no_of_packets]\n",
    "            nus = np.hstack([nus, additional_nus])[:no_of_packets]\n",
    "\n",
    "        radiation_field_luminosity = (\n",
    "            self.calculate_radfield_luminosity().to(u.erg / u.s).value\n",
    "        )\n",
    "\n",
    "        return PacketCollection(radii, nus, mus, energies, packet_seeds, radiation_field_luminosity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call an instance of the packet source class\n",
    "packet_source = TruncBlackbodySource(\n",
    "    truncation_wavelength=2000, base_seed=53253\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now run TARDIS both with and without our custom packet source, and we compare the results:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mdl = run_tardis(\"tardis_example.yml\", packet_source=packet_source)\n",
    "mdl_norm = run_tardis(\"tardis_example.yml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "plt.plot(mdl.transport.transport_state.spectrum_virtual.wavelength,\n",
    "         mdl.transport.transport_state.spectrum_virtual.luminosity_density_lambda,\n",
    "         color='red', label='truncated blackbody (custom packet source)')\n",
    "plt.plot(mdl_norm.transport.transport_state.spectrum_virtual.wavelength,\n",
    "         mdl_norm.transport.transport_state.spectrum_virtual.luminosity_density_lambda,\n",
    "         color='blue', label='normal blackbody (default packet source)')\n",
    "plt.xlabel('$\\lambda [\\AA]$')\n",
    "plt.ylabel('$L_\\lambda$ [erg/s/$\\AA$]')\n",
    "plt.xlim(500, 10000)\n",
    "plt.legend()"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}