{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# More Advanced Thermal Emission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this example, we'll look at the emission from a disk galaxy from the Illustris TNG simulations.\n",
    "This dataset has metallicity information for several species in it. We'll make a cut in phase space like we did in the previous example. The dataset we want to use for this example is available for download [here](https://hea-www.cfa.harvard.edu/~jzuhone/cutout_31_rotated.hdf5). \n",
    "\n",
    "First, import our necessary modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import yt\n",
    "import pyxsim\n",
    "import soxs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We will make phase space cuts on the gas cells using density, temperature, and star formation rate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Note that the units of all numbers in this function are CGS\n",
    "# define hot gas filter\n",
    "def hot_gas(pfilter, data):\n",
    "    pfilter1 = data[pfilter.filtered_type, \"temperature\"] > 3.0e5\n",
    "    pfilter2 = data[\"PartType0\", \"StarFormationRate\"] == 0.0\n",
    "    pfilter3 = data[pfilter.filtered_type, \"density\"] < 5e-25\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=[\"temperature\", \"density\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, we `load` the dataset with yt, and add the `\"hot_gas\"` filter to the dataset: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ds = yt.load(\n",
    "    \"cutout_31_rotated.hdf5\",\n",
    "    bounding_box=[[-1000.0, 1000], [-1000.0, 1000], [-1000.0, 1000]],\n",
    ")\n",
    "ds.add_particle_filter(\"hot_gas\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We also need to tell pyXSIM which elements have fields in the dataset that \n",
    "should be used. To do this we create a `var_elem` dictionary of (key, value) \n",
    "pairs corresponding to the element name and the yt field name (assuming the \n",
    "`\"hot_gas\"` type)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# metal fields to use\n",
    "metals = [\n",
    "    \"C_fraction\",\n",
    "    \"N_fraction\",\n",
    "    \"O_fraction\",\n",
    "    \"Ne_fraction\",\n",
    "    \"Mg_fraction\",\n",
    "    \"Si_fraction\",\n",
    "    \"Fe_fraction\",\n",
    "]\n",
    "var_elem = {elem.split(\"_\")[0]: (\"hot_gas\", elem) for elem in metals}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now that we have everything we need, we'll set up the `IGMSourceModel`, which is based on Cloudy and includes resonant scattering off of the CXB (see [here](https://hea-www.cfa.harvard.edu/~jzuhone/pyxsim/source_models/thermal_sources.html#igm-source-model) for more details). Because we created a hot gas filter, we will use the `\"hot_gas\"` field type for the emission measure, temperature, and metallicity fields. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "source_model = pyxsim.IGMSourceModel(\n",
    "    0.1,\n",
    "    4.0,\n",
    "    5000,\n",
    "    (\"hot_gas\", \"metallicity\"),\n",
    "    binscale=\"log\",\n",
    "    resonant_scattering=True,\n",
    "    temperature_field=(\"hot_gas\", \"temperature\"),\n",
    "    emission_measure_field=(\"hot_gas\", \"emission_measure\"),\n",
    "    nh_field=(\"hot_gas\", \"H_nuclei_density\"),\n",
    "    var_elem=var_elem,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As in other examples, we choose big numbers for the collecting area and exposure time, and a redshift:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "exp_time = (1.0, \"Ms\")  # exposure time\n",
    "area = (5000.0, \"cm**2\")  # collecting area\n",
    "redshift = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, we'll create a box object to serve as a source for the photons. The dataset consists of only\n",
    "the galaxy at a specific location, which we use below, and pick a width of 1 Mpc:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "c = ds.arr([0.0, 0.0, 0.0], \"code_length\")\n",
    "width = ds.quan(1.0, \"Mpc\")\n",
    "le = c - 0.5 * width\n",
    "re = c + 0.5 * width\n",
    "box = ds.box(le, re)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "So, that's everything--let's create the photons! We use the `make_photons` function for this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "n_photons, n_cells = pyxsim.make_photons(\n",
    "    \"cutout_31_photons\", box, redshift, area, exp_time, source_model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "And now we create events using the `project_photons` function. Let's project along the `\"z\"` axis. We'll use the `\"tbabs\"` foreground absorption model this time, with a neutral hydrogen column of $N_H = 2 \\times 10^{20}~{\\rm cm}^{-2}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "n_events = pyxsim.project_photons(\n",
    "    \"cutout_31_photons\",\n",
    "    \"cutout_31_events\",\n",
    "    \"x\",\n",
    "    (30.0, 45.0),\n",
    "    absorb_model=\"tbabs\",\n",
    "    nH=0.02,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now that we have a set of \"events\" on the sky, we can use them as an input to the instrument simulator in SOXS. We'll observe it with the 2eV LEM model for 1 Ms. First, we'll create a background file that we'll use for the background:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.make_background_file(\n",
    "    \"bkgnd_evt_31.fits\", (1000.0, \"ks\"), \"lem_2eV\", [30.0, 45.0], overwrite=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we simulate the source itself, adding in the background:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.instrument_simulator(\n",
    "    \"cutout_31_events.h5\",\n",
    "    \"evt_31.fits\",\n",
    "    (1000.0, \"ks\"),\n",
    "    \"lem_2eV\",\n",
    "    [30.0, 45.0],\n",
    "    overwrite=True,\n",
    "    bkgnd_file=\"bkgnd_evt_31.fits\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can use the `write_image()` function in SOXS to bin the events into an image and write them \n",
    "to a file, restricting the energies between 0.644 and 0.65 keV, which focuses on the redshifted OVIII line:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.write_image(\"evt_31.fits\", \"img_31.fits\", emin=0.644, emax=0.65, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we can take a quick look at the image: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.plot_image(\"img_31.fits\", stretch=\"log\", cmap=\"arbre\", width=0.4, vmin=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now we will make spectra to look at. First, filter the events of both the combined source and background files and the background-only files within 0.15 degree of the center:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.filter_events(\n",
    "    \"evt_31.fits\",\n",
    "    \"evt_31_filter.fits\",\n",
    "    overwrite=True,\n",
    "    region='fk5\\ncircle(30.0000000,45.0000000,540.000\")',\n",
    ")\n",
    "soxs.filter_events(\n",
    "    \"bkgnd_evt_31.fits\",\n",
    "    \"bkgnd_evt_31_filter.fits\",\n",
    "    overwrite=True,\n",
    "    region='fk5\\ncircle(30.0000000,45.0000000,540.000\")',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now bin up spectra for these new event files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.write_spectrum(\"evt_31_filter.fits\", \"evt_31.pi\", overwrite=True)\n",
    "soxs.write_spectrum(\"bkgnd_evt_31_filter.fits\", \"bkgnd_evt_31.pi\", overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, we can plot the spectra. Below, the total spectrum is in blue and the background/foreground spectrum is in orange. The lines from the emission of the distant galaxy are redshifted away from the foreground Milky Way lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "fig, ax = soxs.plot_spectrum(\"evt_31.pi\", xmin=0.5, xmax=0.7, xscale=\"linear\", ymin=0.5)\n",
    "soxs.plot_spectrum(\"bkgnd_evt_31.pi\", xmin=0.5, xmax=0.7, fig=fig, ax=ax, ymin=0.5)"
   ]
  }
 ],
 "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}