{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Milky Way CCSN Distributions\n",
    "\n",
    "Load FITS data containing the probability of finding a core-collapse supernova (CCSN) at some radial distance from the Sun. The data come from a variety of papers with more or less \"realistic\" information about the structure of the Milky Way. Since the structure and distribution of CCSN are uncertain these models can be taken to represent the systematic uncertainties in the expected distribution of Galactic CCSNe.\n",
    "\n",
    "Sources for the models include:\n",
    "\n",
    "1. J. Bahcall and T. Piran, ApJL 267:L77, 1983.\n",
    "1. A. Mirizzi, G. Raffelt, P. Serpico, JCAP 0605:012, 2006 (two parameterizations).\n",
    "1. M. Ahlers, P. Mertsch, S. Sarkar, PRD 80:123017, 2009.\n",
    "1. S. Adams, C. Kochanek, J. Beacom, M. Vagins, K. Stanek, ApJ 778:164, 2013.\n",
    "\n",
    "The data are stored as cumulative distributions, making it easy to randomly generate distances with respect to the solar system according to their respective probability density functions. Examples are shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from astropy.io import fits\n",
    "from astropy import units as u\n",
    "from scipy.interpolate import PchipInterpolator\n",
    "\n",
    "from importlib.resources import files\n",
    "\n",
    "from asteria import set_rcparams\n",
    "from asteria.stellardist import FixedDistance, StellarDensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_rcparams()\n",
    "\n",
    "# grid_style =   {     'alpha' : '0.75',\n",
    "#                  'linestyle' : ':' }\n",
    "# legend_style = {  'fontsize' : '12' }\n",
    "# font_syle =    {      'size' : '14' }\n",
    "\n",
    "# mpl.rc(  'font', **font_syle)\n",
    "# mpl.rc(  'grid', **grid_style)\n",
    "# mpl.rc('legend', **legend_style)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test FixedDistance Model\n",
    "\n",
    "The `FixedDistance` class will create a progenitor at a fixed location $d$. By default, $d$ has no uncertainty. It is also possible to introduce a distance uncertainty $\\sigma$; sampling from `FixedDistance` in this case will produce a Gaussian distribution in $d$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fd1 = FixedDistance(10*u.kpc)\n",
    "d1 = fd1.distance(10000)\n",
    "\n",
    "fd2 = FixedDistance(10*u.kpc, sigma=0.5*u.kpc)\n",
    "d2 = fd2.distance(10000)\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize=(12,4), sharex=True)\n",
    "ax1, ax2 = axes\n",
    "\n",
    "bins = np.linspace(7, 13, 61)\n",
    "ax1.hist(d1.value, bins, density=True)\n",
    "ax1.set(ylabel='$p(d)$ [kpc$^{-1}$]',\n",
    "        xlabel='distance [kpc]',\n",
    "        title='$d=${}'.format(fd1.dist))\n",
    "ax1.grid(ls=':')\n",
    "ax2.hist(d2.value, bins, density=True)\n",
    "ax2.set(xlabel='distance [kpc]',\n",
    "        title='$d=${}, $\\sigma=${}'.format(fd2.dist, fd2.sigma))\n",
    "ax2.grid(ls=':')\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Radial Stellar Mass Distribution Models\n",
    "\n",
    "The `StellarDensity` model is based on estimates of the radial distribution of matter (dominated by the Milky Way) with respect to the Sun."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Cumulative Stellar Mass Distributions\n",
    "\n",
    "Internally, `StellarDensity` stores the cumulative distribution as a function of radial distance from the solar system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = ['bahcall', 'mirizzi_fid', 'mirizzi_ben', 'ahlers', 'adams']\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(8,5), tight_layout=True)\n",
    "\n",
    "for model in models:\n",
    "    sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
    "    sd = StellarDensity(sdfile)\n",
    "    ax.plot(sd.dist, sd.cdf, lw=3, alpha=0.7, label=sd.name.replace('&', '\\&'))\n",
    "\n",
    "ax.set(xlabel='distance [kpc]',\n",
    "       ylabel='probability',\n",
    "       ylim=(0,1.05))\n",
    "ax.grid(ls=':')\n",
    "ax.legend(fontsize=12);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cumulative Distributions with LMC and SMC Added\n",
    "\n",
    "Add a simple Gaussian model of the LMC and SMC stellar mass density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = ['bahcall', 'mirizzi_fid', 'mirizzi_ben', 'ahlers', 'adams']\n",
    "\n",
    "fig, ax = plt.subplots(1,1, figsize=(8,5), tight_layout=True)\n",
    "\n",
    "for model in models:\n",
    "    sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
    "    sd = StellarDensity(sdfile, add_LMC=True, add_SMC=True)\n",
    "    ax.plot(sd.dist, sd.cdf, lw=3, alpha=0.7, label=sd.name.replace('&', '\\&'))\n",
    "\n",
    "ax.set(xlabel='distance [kpc]',\n",
    "       ylabel='probability',\n",
    "       ylim=(0,1.05))\n",
    "ax.grid(ls=':')\n",
    "ax.legend(fontsize=12);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Radial Distances\n",
    "\n",
    "Randomly sample radial distances using several available stellar distribution models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = ['bahcall', 'mirizzi_fid', 'mirizzi_ben', 'ahlers', 'adams']\n",
    "\n",
    "fig, axes = plt.subplots(2,3, figsize=(12,5), sharex=True, sharey=True, tight_layout=True)\n",
    "axes = axes.flatten()\n",
    "axes[5].axis('off')\n",
    "\n",
    "for i, (model, ax) in enumerate(zip(models, axes)):\n",
    "    sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
    "    sd = StellarDensity(sdfile)\n",
    "    distances = sd.distance(100000)\n",
    "    \n",
    "    ax.hist(distances.value, bins=np.linspace(0., 30., 61), color='C{}'.format(i),\n",
    "            alpha=0.7,\n",
    "            label=sd.name.replace('&', '\\&'),\n",
    "            density=True)\n",
    "    ax.legend(fontsize=10)\n",
    "    ax.grid()\n",
    "\n",
    "axes[0].set(ylim=(0,0.13),\n",
    "            ylabel='prob. [kpc$^{-1}$]')\n",
    "axes[3].set(xlim=(0,33),\n",
    "            xlabel='distance [kpc]')\n",
    "fig.subplots_adjust(hspace=0, wspace=0);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Radial Distances with LMC and SMC Added"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models = ['bahcall', 'mirizzi_fid', 'mirizzi_ben', 'ahlers', 'adams']\n",
    "\n",
    "fig, axes = plt.subplots(2,3, figsize=(12,5), sharex=True, sharey=True, tight_layout=True)\n",
    "axes = axes.flatten()\n",
    "axes[5].axis('off')\n",
    "\n",
    "for i, (model, ax) in enumerate(zip(models, axes)):\n",
    "    sdfile = files('asteria.data.stellar').joinpath(f'sn_radial_distrib_{model}.fits')\n",
    "    sd = StellarDensity(sdfile, add_LMC=True, add_SMC=True)\n",
    "    distances = sd.distance(100000)\n",
    "    \n",
    "    ax.hist(distances.value, bins=np.linspace(0, 70, 71), color='C{}'.format(i),\n",
    "            alpha=0.7,\n",
    "            label=sd.name.replace('&', '\\&'),\n",
    "            density=True)\n",
    "    ax.legend(fontsize=10)\n",
    "    ax.grid()\n",
    "\n",
    "axes[0].set(ylim=(0,0.13),\n",
    "            ylabel='prob. [kpc$^{-1}$]')\n",
    "axes[3].set(xlim=(0,70),\n",
    "            xlabel='distance [kpc]')\n",
    "fig.subplots_adjust(hspace=0, wspace=0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}