{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Line Emission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "This example shows how to create a line emission model. It uses a galaxy cluster from a Gadget SPH cosmological dataset, and will create a thermal model out of the gas particles and will use the dark matter particles to add line emission to the spectrum, assuming that the emission comes from some decay process of the dark matter.\n",
    "\n",
    "First, we load the modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib\n",
    "\n",
    "matplotlib.rc(\"font\", size=18, family=\"serif\")\n",
    "import yt\n",
    "import matplotlib.pyplot as plt\n",
    "from yt.units import mp\n",
    "import pyxsim"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "As in the [Advanced Thermal Emission](Advanced_Thermal_Emission.ipynb) example, we'll set \n",
    "up a hot gas filter:"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# Note that the units of all numbers in this function are CGS\n",
    "def hot_gas(pfilter, data):\n",
    "    pfilter1 = data[pfilter.filtered_type, \"density\"] < 5e-25\n",
    "    pfilter2 = data[pfilter.filtered_type, \"temperature\"] > 3481355.78432401\n",
    "    pfilter3 = data[pfilter.filtered_type, \"temperature\"] < 4.5e8\n",
    "    return (pfilter1) & (pfilter2) & (pfilter3)\n",
    "\n",
    "\n",
    "yt.add_particle_filter(\n",
    "    \"hot_gas\",\n",
    "    function=hot_gas,\n",
    "    filtered_type=\"gas\",\n",
    "    requires=[\"density\", \"temperature\"],\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "The dataset used here does not have a field for the electron number\n",
    "density, which is required to construct the emission measure field. Because we'll\n",
    "only be using the hot gas, we can create a `(\"gas\",\"El_number_density\")` field\n",
    "which assumes complete ionization (while taking into account the H and He mass\n",
    "fractions vary from particle to particle). This is not strictly true for all of the\n",
    "`\"gas\"` type particles, but since we'll be using the `\"hot_gas\"` type it should\n",
    "be sufficiently accurate for our purposes. We'll define the field here and add it."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def _El_number_density(field, data):\n",
    "    mueinv = data[\"gas\", \"H_fraction\"] + 0.5 * (1.0 - data[\"gas\", \"He_fraction\"])\n",
    "    return data[\"gas\", \"density\"] * mueinv / (1.0 * mp)\n",
    "\n",
    "\n",
    "yt.add_field(\n",
    "    (\"gas\", \"El_number_density\"),\n",
    "    _El_number_density,\n",
    "    units=\"cm**-3\",\n",
    "    sampling_type=\"local\",\n",
    ")"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "and we load the dataset in yt:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ds = yt.load(\"snapshot_033/snap_033.0.hdf5\")\n",
    "ds.add_particle_filter(\"hot_gas\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To get a sense of what the cluster looks like, let's take a slice through the density and temperature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "slc = yt.SlicePlot(\n",
    "    ds,\n",
    "    \"z\",\n",
    "    [(\"gas\", \"density\"), (\"gas\", \"temperature\")],\n",
    "    center=\"max\",\n",
    "    width=(3.0, \"Mpc\"),\n",
    ")\n",
    "slc.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now set up a sphere centered on the maximum density in the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "sp = ds.sphere(\"max\", (1.0, \"Mpc\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "and create a thermal source model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "thermal_model = pyxsim.CIESourceModel(\"apec\", 0.2, 11.0, 10000, (\"gas\", \"metallicity\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we'll set up a line emission field for the dark matter particles. We won't try to replicate a specific model, but will simply assume that the emission is proportional to the dark matter mass. Note that this field is a particle field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "norm = yt.YTQuantity(100.0, \"g**-1*s**-1\")\n",
    "\n",
    "\n",
    "def _dm_emission(field, data):\n",
    "    return data[\"PartType1\", \"particle_mass\"] * norm\n",
    "\n",
    "\n",
    "ds.add_field(\n",
    "    (\"PartType1\", \"dm_emission\"),\n",
    "    function=_dm_emission,\n",
    "    sampling_type=\"particle\",\n",
    "    units=\"photons/s\",\n",
    "    force_override=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we can set up the `LineSourceModel` object. The first argument is the line center energy in keV, and the second is the field we just created, that sets up the line amplitude. There is another parameter, `sigma`, for adding in broadening of the line, but in this case we'll rely on the velocities of the dark matter particles themselves to produce the line broadening."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "line_model = pyxsim.LineSourceModel(3.5, (\"PartType1\", \"dm_emission\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now set up the parameters for generating the photons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "exp_time = (300.0, \"ks\")  # exposure time\n",
    "area = (1000.0, \"cm**2\")  # collecting area\n",
    "redshift = 0.05"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "and actually generate the photons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ntp, ntc = pyxsim.make_photons(\n",
    "    \"therm_photons.h5\", sp, redshift, area, exp_time, thermal_model\n",
    ")\n",
    "nlp, nlc = pyxsim.make_photons(\n",
    "    \"line_photons.h5\", sp, redshift, area, exp_time, line_model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, project the photons for the total set and the line set by itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "nte = pyxsim.project_photons(\n",
    "    \"therm_photons.h5\",\n",
    "    \"therm_events.h5\",\n",
    "    \"y\",\n",
    "    (30.0, 45.0),\n",
    "    absorb_model=\"wabs\",\n",
    "    nH=0.02,\n",
    ")\n",
    "nle = pyxsim.project_photons(\n",
    "    \"line_photons.h5\", \"line_events.h5\", \"y\", (30.0, 45.0), absorb_model=\"wabs\", nH=0.02\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Write the raw, unconvolved spectra to disk:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "et = pyxsim.EventList(\"therm_events.h5\")\n",
    "el = pyxsim.EventList(\"line_events.h5\")\n",
    "et.write_spectrum(\"therm_spec.fits\", 0.1, 10.0, 5000, overwrite=True)\n",
    "el.write_spectrum(\"line_spec.fits\", 0.1, 10.0, 5000, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now let's plot up both spectra. We see that we have a thermal spectrum with the addition of a line at 3.5 keV (in real life such a line would not be so prominent, but it makes the example easier to see):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import astropy.io.fits as pyfits\n",
    "\n",
    "f1 = pyfits.open(\"therm_spec.fits\")\n",
    "f2 = pyfits.open(\"line_spec.fits\")\n",
    "fig = plt.figure(figsize=(9, 7))\n",
    "ax = fig.add_subplot(111)\n",
    "ax.loglog(\n",
    "    f1[\"SPECTRUM\"].data[\"ENERGY\"],\n",
    "    f1[\"SPECTRUM\"].data[\"COUNTS\"] + f2[\"SPECTRUM\"].data[\"COUNTS\"],\n",
    ")\n",
    "ax.loglog(f2[\"SPECTRUM\"].data[\"ENERGY\"], f2[\"SPECTRUM\"].data[\"COUNTS\"])\n",
    "ax.set_xlim(0.2, 10)\n",
    "ax.set_ylim(1, 3.0e4)\n",
    "ax.set_xlabel(\"E (keV)\")\n",
    "ax.set_ylabel(\"counts/bin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's zoom into the region surrounding the line, seeing that it has some broadening due to the random velocities of the dark matter particles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ax.set_xscale(\"linear\")\n",
    "ax.set_xlim(3, 3.7)\n",
    "ax.set_ylim(1.0, 3.0e2)\n",
    "fig"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}