{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "04bf2ca6",
   "metadata": {},
   "source": [
    "# Takata 2025 2D Models with Axionlike Production\n",
    "\n",
    "Neutrino data for 1-D core-collapse supernova models in the presence of axion-like particles.\n",
    "\n",
    "The reference article is [Progenitor dependence of neutrino-driven supernova explosions with the aid of heavy axionlike particles](https://doi.org/10.1103/PhysRevD.111.103028) by T. Takata et al., Phys. Rev. D 111:103028, 2025.\n",
    "\n",
    "The following models are supported:\n",
    "\n",
    "## ALP parameters and progenitor masses simulated\n",
    "\n",
    "| Progenitor mass [Msun] | Axion mass [MeV] | Axion-photon coupling [1e-10/GeV] |\n",
    "| ---------------------- | ---------------- | --------------------------------- |\n",
    "| 11.2 | 0, 40, 100, 150, 200, 300, 400, 600, 800 | 0, 2, 4, 6, 8, 10 |\n",
    "| 20 | 0, 40, 100, 150, 200, 300, 400, 600, 800 | 0, 2, 4, 6, 8, 10 |\n",
    "| 25 | 0, 40, 100, 150, 200, 300, 400, 600, 800 | 0, 2, 4, 6, 8, 10 |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29ceae1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snewpy.neutrino import Flavor\n",
    "from snewpy.models.ccsn import Takata_2025\n",
    "\n",
    "from astropy import units as u\n",
    "\n",
    "from scipy.interpolate import PchipInterpolator\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b7566ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rc('font', size=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66d13225",
   "metadata": {},
   "source": [
    "## Initialize the 2D models\n",
    "\n",
    "Use the `param` property of the model class to see the available parameters. Models are initialized using the `axion_mass` and `axion_coupling` parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f62e690f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Takata_2025.param"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55479166",
   "metadata": {},
   "source": [
    "The model with `axion_mass=0` and `axion_coupling=0` is a standard simulation with no ALP production.\n",
    "\n",
    "We'll use the model with a 25 M$_\\odot$ progenitor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649f5ae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "pmass = 25*u.Msun\n",
    "model_std = Takata_2025(axion_mass=0, axion_coupling=0, progenitor_mass=pmass)\n",
    "model_std.metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f15de7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a handful of axion models.\n",
    "models = {}\n",
    "for (am, ac) in ((40*u.MeV, 4e-10/u.GeV), (100*u.MeV, 6e-10/u.GeV), (200*u.MeV, 6e-10/u.GeV), (300*u.MeV, 8e-10/u.GeV), (400*u.MeV, 10e-10/u.GeV), (600*u.MeV, 6e-10/u.GeV)):\n",
    "    models[(am,ac)] = Takata_2025(axion_mass=am, axion_coupling=ac, progenitor_mass=pmass)\n",
    "\n",
    "models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cab5a9c",
   "metadata": {},
   "source": [
    "## Plot Model Luminosities\n",
    "\n",
    "Compare axion model luminosity to the standard 2D simulation.\n",
    "\n",
    "Higher mass models with stronger coupling constants should produce a decrease in neutrino luminosity at all flavors relative to the reference simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a828245a-719a-4f56-a6ad-42d34b9bc18e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_std.metadata[\"Progenitor mass\"].to_value(\"Msun\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a398333",
   "metadata": {},
   "outputs": [],
   "source": [
    "for (m,c), model in models.items():\n",
    "    \n",
    "    fig, axes = plt.subplots(2, 6, figsize=(40, 8), sharex=True,\n",
    "                         gridspec_kw={'height_ratios':[3.2,1], 'hspace':0, 'wspace':0.05})\n",
    "    \n",
    "    Lmin,  Lmax  = 1e99, -1e99\n",
    "    dLmin, dLmax = 1e99, -1e99\n",
    "    \n",
    "    for j, (flavor) in enumerate(Flavor):\n",
    "        ax = axes[0][j]\n",
    "        \n",
    "        ax.plot(model_std.time, model_std.luminosity[flavor]/1e51, 'k', label=rf'${{{model_std.metadata[\"Progenitor mass\"].to_value(\"Msun\"):g}}}M_\\odot$ reference')  # Report luminosity in [foe/s]\n",
    "        Lmin = np.minimum(Lmin, np.min(model_std.luminosity[flavor].to_value('1e51 erg/s')))\n",
    "        Lmax = np.maximum(Lmax, np.max(model_std.luminosity[flavor].to_value('1e51 erg/s')))\n",
    "        \n",
    "        modlabel = rf\"{flavor.to_tex()}: $m_a=${m.to_string(format='latex_inline')}\" + \"\\n\" + rf\"     $g_{{a\\gamma}}=${c.to_string(format='latex_inline')}\"\n",
    "        ax.plot(model.time, model.luminosity[flavor]/1e51,  # Report luminosity in units foe/s\n",
    "                label=modlabel,\n",
    "                color='C0' if flavor.is_electron else 'C1',\n",
    "                ls='-' if flavor.is_neutrino else ':',\n",
    "                lw=2)\n",
    "        if j==0:\n",
    "            ax.set(ylabel=r'luminosity [$10^{51}$ erg s$^{-1}$]')\n",
    "        \n",
    "        ax.legend(fontsize=12)\n",
    "        ax.set(xlim=(model_std.time[0].to_value('s'), model_std.time[-1].to_value('s')))\n",
    "        \n",
    "        ax = axes[1][j]\n",
    "        tmin = np.maximum(model.time[0], model_std.time[0]).to_value('s')\n",
    "        tmax = np.minimum(model.time[-1], model_std.time[-1]).to_value('s')\n",
    "        times = np.arange(tmin, tmax, 0.001)*u.s\n",
    "\n",
    "        Lstd = PchipInterpolator(model_std.time, model_std.luminosity[flavor].to_value('1e51 erg/s'))\n",
    "        Lstd_t = Lstd(times)\n",
    "        select = Lstd_t != 0\n",
    "        \n",
    "        Lmod = PchipInterpolator(model.time, model.luminosity[flavor].to_value('1e51 erg/s'))\n",
    "        Lmod_t = Lmod(times)\n",
    "        dL = (Lmod_t[select] - Lstd_t[select]) / Lstd_t[select]\n",
    "        \n",
    "        dLmin = np.minimum(dLmin, np.min(dL))\n",
    "        dLmax = np.maximum(dLmax, np.max(dL))\n",
    "\n",
    "        ax.plot(times[select], dL)\n",
    "        if j==0:\n",
    "            ax.set(xlabel='time [s]',\n",
    "                   ylabel=r'$\\Delta L_\\nu/L_\\nu$')\n",
    "            \n",
    "    for j in range(6):\n",
    "        axes[0][j].set(ylim=(Lmin, 1.1*Lmax))\n",
    "        axes[1][j].set(ylim=(dLmin, dLmax))\n",
    "        if j > 0:\n",
    "            axes[0][j].set_yticklabels([])\n",
    "            axes[1][j].set_yticklabels([])\n",
    "        \n",
    "    fig.suptitle(rf\"Axionlike model: $m_a=${m.to_string(format='latex_inline')}, $g_{{a\\gamma}}=${c.to_string(format='latex_inline')}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eabc4b8e-3b43-4c22-bd23-17ff73707a08",
   "metadata": {},
   "source": [
    "## Plot Model Spectra\n",
    "\n",
    "Plot model spectra for the baseline model and a few of the alternative models for several times. Assume the progenitor is at a distance of 1 kpc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cd83985-a0ed-42b5-84ac-378a346595ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.arange(0.05, 0.35, 0.05) * u.s\n",
    "E = np.arange(0, 50.1, 0.1) * u.MeV\n",
    "d = 1*u.kpc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d74a1fc-3e16-49df-920b-a4ca3f263d9f",
   "metadata": {},
   "source": [
    "### Compare Baseline Model to 100 MeV ALP\n",
    "\n",
    "Show the spectra from the ALP model at several times. The baseline model (no ALP) spectra are drawn as dashed lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc54866f-bbcc-47e2-bba6-6234e0455fa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = t.shape[0]\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize=(10,4.5), sharex=True, tight_layout=True)\n",
    "\n",
    "fmin, fmax = 0, -1e99\n",
    "\n",
    "for fl in (Flavor.NU_E, Flavor.NU_E_BAR):\n",
    "    ax = axes[fl]\n",
    "\n",
    "    for j in np.arange(nt):\n",
    "        color = None\n",
    "        for i, model in enumerate([model_std, models[(100*u.MeV, 6e-10/u.GeV)]]):\n",
    "            d2fdEdt = model.get_flux(t, E, d)\n",
    "            dfdE = d2fdEdt.array[fl, j, :].to_value('1e12/(MeV s cm2)')\n",
    "            fmax = np.maximum(fmax, np.max(dfdE))\n",
    "            ls = '--' if i == 0 else '-'\n",
    "            label = None if i == 0 else rf'$t_\\text{{pb}}={t[j].to_value(\"ms\"):g}$ ms'\n",
    "            width = 1 if i==0 else 1.5\n",
    "            line, = ax.plot(E, dfdE, label=label, ls=ls, lw=width, color=color)\n",
    "            color = line.get_color()\n",
    "\n",
    "    ax.set(xlim=E[[0,-1]].to_value(),\n",
    "           xlabel='energy [MeV]',\n",
    "           ylim=(fmin, 1.1*fmax),\n",
    "           ylabel=rf'$\\Phi_{{{fl.to_tex()[1:-1]}}}$ [$10^{{12}}$ MeV$^{{-1}}$ s$^{{-1}}$ m$^{{-2}}$]',\n",
    "          )\n",
    "    ax.set_title(rf'$m_a=\\text{{{model.metadata[\"Axion mass\"]}}}$, $g_{{a\\gamma}}={model.metadata[\"Axion coupling\"].to_value(\"1e-10/GeV\")}\\times10^{{-10}}$ GeV$^{{-1}}$', fontsize=12)\n",
    "    \n",
    "    ax.legend(loc='upper right', fontsize=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf64f3c5-f1b6-4a46-a4ad-499bdb3da64c",
   "metadata": {},
   "source": [
    "### Compare Baseline Model to 400 MeV ALP\n",
    "\n",
    "Show the spectra from the ALP model at several times. The baseline model (no ALP) spectra are drawn as dashed lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84c54e19-8051-48f3-ba08-606013fd44b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "nt = t.shape[0]\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize=(10,4.5), sharex=True, tight_layout=True)\n",
    "\n",
    "fmin, fmax = 0, -1e99\n",
    "\n",
    "for fl in (Flavor.NU_E, Flavor.NU_E_BAR):\n",
    "    ax = axes[fl]\n",
    "\n",
    "    for j in np.arange(nt):\n",
    "        color = None\n",
    "        for i, model in enumerate([model_std, models[(400*u.MeV, 10e-10/u.GeV)]]):\n",
    "            d2fdEdt = model.get_flux(t, E, d)\n",
    "            dfdE = d2fdEdt.array[fl, j, :].to_value('1e12/(MeV s cm2)')\n",
    "            fmax = np.maximum(fmax, np.max(dfdE))\n",
    "            ls = '--' if i == 0 else '-'\n",
    "            label = None if i == 0 else rf'$t_\\text{{pb}}={t[j].to_value(\"ms\"):g}$ ms'\n",
    "            width = 1 if i==0 else 1.5\n",
    "            line, = ax.plot(E, dfdE, label=label, ls=ls, lw=width, color=color)\n",
    "            color = line.get_color()\n",
    "\n",
    "    ax.set(xlim=E[[0,-1]].to_value(),\n",
    "           xlabel='energy [MeV]',\n",
    "           ylim=(fmin, 1.1*fmax),\n",
    "           ylabel=rf'$\\Phi_{{{fl.to_tex()[1:-1]}}}$ [$10^{{12}}$ MeV$^{{-1}}$ s$^{{-1}}$ m$^{{-2}}$]',\n",
    "          )\n",
    "    ax.set_title(rf'$m_a=\\text{{{model.metadata[\"Axion mass\"]}}}$, $g_{{a\\gamma}}={model.metadata[\"Axion coupling\"].to_value(\"1e-10/GeV\")}\\times10^{{-10}}$ GeV$^{{-1}}$', fontsize=12)\n",
    "    \n",
    "    ax.legend(loc='upper right', fontsize=10)"
   ]
  }
 ],
 "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.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}