{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing Long-Duration Gamma-Ray Bursts from Double Compact Object Populations 🌌"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the GRB TransientPopulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the same BBH population as in the `Analyzing Merging BBH Populations: Rates & Observations` tutorial to compute the long gamma-ray bursts (LGRB) rate associated with the formation of merging BBHs.\n",
    "\n",
    "Make sure the `BBH_contact.h5` file is present to continue with this tutorial.\n",
    "You can find a copy of it in the example dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from posydon.popsyn.synthetic_population import Population\n",
    "\n",
    "file = 'BBH_contact.h5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BBH_population = Population(file)\n",
    "BBH_population.mass_per_metallicity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BBH_population.calculate_formation_channels(mt_history=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We had already selected only binaries that result in a merger in the `BBH_contact.h5` file, so we don't need to perform any selection on the binaries.\n",
    "\n",
    "Thus, our next step is to create a selection function for LGRBs based on the properties in the history and onlines of the binaries.\n",
    "\n",
    "We have already included a basic selection function in `posydon.popsyn.transient_select_func`, which select if a GRB occurs based on the presence of `m_disk_radiated` for either star.\n",
    "It will also output some pre and post supernova properties.\n",
    "\n",
    "As you can see below, the `GRB_selection` requires an additional input parameter of `S1_S2`. \n",
    "The `Population.create_transient_population` function cannot pass these additional arguments to sub-function.\n",
    "Thus, we will be required to wrap the `GRB_selection` functions.\n",
    "\n",
    "This is useful, if you would like to create similar population, but use different model parameters.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from posydon.popsyn.transient_select_funcs import GRB_selection\n",
    "\n",
    "GRB_selection(BBH_population.history[10],\n",
    "              BBH_population.oneline[10],\n",
    "              BBH_population.formation_channels.loc[[10]],\n",
    "              S1_S2='S2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "def GRB_wrapper(transient_chunk, oneline_chunk, formation_channels_chunk):\n",
    "    \"\"\"calculate the GRBs from star 1 and star 2 for a given chunk of the population\"\"\"\n",
    "    # select potential GRBs from star 1\n",
    "    df_1 = GRB_selection(transient_chunk, oneline_chunk, formation_channels_chunk, S1_S2='S1')\n",
    "    # select potential GRBs from star 2\n",
    "    df_2 = GRB_selection(transient_chunk, oneline_chunk, formation_channels_chunk, S1_S2='S2')\n",
    "    \n",
    "    # combine the two dataframes\n",
    "    if df_1 is not None and df_2 is not None:\n",
    "        GRB_df = pd.concat([df_1, df_2])\n",
    "    elif df_1 is not None:\n",
    "        GRB_df = df_1\n",
    "    elif df_2 is not None:\n",
    "        GRB_df = df_2\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "    return GRB_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LGRB = BBH_population.create_transient_population(GRB_wrapper, 'LGRB')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You've now created a LGRB population, where either the first and/or second star has some amount of radiative disk.\n",
    "\n",
    "This is, of course, just an example and the actual amount of disk required to power a LGRB will be higher than some of the values included in our \"LGRB\" rate.\n",
    "Moreover, the number of binaries in this population is insufficient to actually see the all unique events.\n",
    "\n",
    "For now, let's continue and calculate the metallicity bias function for the LGRBs.\n",
    "We don't need to calculate the `underlying_mass` of the population, because we already did this in the [BBH tutorial](bbh_analysis.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LGRB.get_efficiency_over_metallicity()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LGRB.plot_efficiency_over_metallicity(channels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After this selection, we follow similar steps to the BBH analysis for creating and plotting a cosmic star formation history weighted rate.\n",
    "\n",
    "If you've followed the previous tutorial on the BBH analaysis, you will still have access to those populations too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LGRB_rates = LGRB.calculate_cosmic_weights('IllustrisTNG')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LGRB_rates.calculate_intrinsic_rate_density(mt_channels=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from posydon.popsyn.synthetic_population import Rates\n",
    "BBH_rates  = Rates('BBH_contact.h5', 'BBH', 'IllustrisTNG')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BBH_rates.intrinsic_rate_density"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we plot the BBH rate and the un-normalised LGRB rate together, showing how you can access multiple event rates, and transient populations from the same file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.plot(BBH_rates.centers_redshift_bins, BBH_rates.intrinsic_rate_density['total'], label='BBH')\n",
    "plt.plot(LGRB_rates.centers_redshift_bins, LGRB_rates.intrinsic_rate_density['total'], label='LGRB')\n",
    "\n",
    "plt.yscale(\"log\")\n",
    "plt.xlim(0,10)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LGRB_rates.population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1)\n",
    "LGRB_rates.plot_hist_properties('S1_spin_postSN', intrinsic=True, label='S1', color='blue', ax=ax, show=False)\n",
    "LGRB_rates.plot_hist_properties('S2_spin_postSN', intrinsic=True, label='S2', color='orange', ax=ax, show=False)\n",
    "ax.legend()\n",
    "ax.set_xlabel('BH spin post LGRB')\n",
    "ax.set_ylabel('\\# Events in bin')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's possible to calculate more properties for the LGRBs using the `posydon.popsyn.GRB` module.\n",
    "\n",
    "It contains the `get_GRB_properties` that can provide a more detailed GRB properties based on empirical or simulation relations for the LGRB energy and beaming factor.\n",
    "\n",
    "Enjoy exploring this function and creating a more detailed GRB population!\n",
    "Keep in mind that it's not possible to add additional columns to the transient population without overwritting the population.\n",
    "\n",
    "The next tutorial will focus on rerunning specific binaries in a population or setting up an unique binary yourself.\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "posydon_env",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}