{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Advanced Thermal Emission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "In this example, we'll look at another galaxy cluster, but this time the dataset will\n",
    "have metallicity information for several species in it. In contrast to the thermal emission\n",
    "example, which used a grid-based dataset, the dataset we'll use here is SPH, taken from the \n",
    "[Magneticum](http://www.magneticum.org) suite of simulations. Finally, there are phases of \n",
    "gas in this dataset that  will not emit in X-rays, so we also show how to make cuts in phase \n",
    "space and focus only on the X-ray emitting gas. The dataset we want to use for this example \n",
    "is available for download from the [yt Project](http://yt-project.org) at \n",
    "[this link](https://yt-project.org/data/MagneticumCluster.tar.gz). \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\n",
    "from yt.units import mp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We mentioned above that we wanted to make phase spaces cuts on the gas. Because this is an \n",
    "SPH dataset, the best way to do that is with a \"particle filter\" in yt. "
   ]
  },
  {
   "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",
    "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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The Magneticum 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def _El_number_density(field, data):\n",
    "    mueinv = data[\"gas\", \"H_fraction\"] + 0.5 * 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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As mentioned above, a number of elements are tracked in the SPH particles in \n",
    "the dataset (10 to be precise, along with a trace field for the remaining, \n",
    "unspecified metals). We will deal with these elements later, since we want to \n",
    "use the specific mass fractions of these elements to determine their emission \n",
    "line strengths in the mock observation. However, we also need to specify the \n",
    "metallicity for the rest of the (non-hydrogen) elements. The best way to do\n",
    "this is to sum the masses for the metals over every particle and divide by\n",
    "the mass of the particle to get the metallicity for that particle. That will\n",
    "be assumed to be the metallicity for all non-tracked metals in this pyXSIM\n",
    "run. The field in the Magneticum dataset to do this is called \n",
    "`(\"Gas\", \"ElevenMetalMasses\")`, which has a shape of (number of SPH particles, 11). \n",
    "We'll define this field here and add it to the dataset specifically later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def _metallicity(field, data):\n",
    "    # We index the array starting with 1 here because the first element is\n",
    "    # helium (thus not a metal)\n",
    "    return data[\"Gas\", \"ElevenMetalMasses\"][:, 1:].sum(axis=1) / data[\"Gas\", \"Mass\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, we `load` the dataset with yt: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ds = yt.load(\n",
    "    \"MagneticumCluster/snap_132\", long_ids=True, field_spec=\"magneticum_box2_hr\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "and now we add the derived fields and the `\"hot_gas\"` particle filter to this dataset. \n",
    "Note that for the derived fields to be picked up by the filter, they must be specified first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% \n"
    }
   },
   "outputs": [],
   "source": [
    "ds.add_field((\"gas\", \"metallicity\"), _metallicity, units=\"\", sampling_type=\"local\")\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": [
    "var_elem = {\n",
    "    elem: (\"hot_gas\", f\"{elem}_fraction\")\n",
    "    for elem in [\"He\", \"C\", \"Ca\", \"O\", \"N\", \"Ne\", \"Mg\", \"S\", \"Si\", \"Fe\"]\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now that we have everything we need, we'll set up the `CIESourceModel`. \n",
    "Because we created a hot gas filter, we will use the `\"hot_gas\"` field type \n",
    "for the emission measure, temperature, and metallicity fields. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "source_model = pyxsim.CIESourceModel(\n",
    "    \"apec\",\n",
    "    0.1,\n",
    "    10.0,\n",
    "    1000,\n",
    "    (\"hot_gas\", \"metallicity\"),\n",
    "    temperature_field=(\"hot_gas\", \"temperature\"),\n",
    "    emission_measure_field=(\"hot_gas\", \"emission_measure\"),\n",
    "    var_elem=var_elem,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As before, we choose big numbers for the collecting area and exposure time, but \n",
    "the redshift should be taken from the cluster itself, since this dataset has a redshift:"
   ]
  },
  {
   "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 = ds.current_redshift"
   ]
  },
  {
   "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 cluster at a specific location, which we use below, and pick a width of 3 Mpc:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "c = ds.arr([310306.53, 340613.47, 265758.47], \"code_length\")\n",
    "width = ds.quan(3.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",
    "    \"snap_132_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 pick an off-axis normal vector, \n",
    "and a `north_vector` to decide which way is \"up.\" We'll use the `\"wabs\"` foreground absorption model\n",
    "this time, with a neutral hydrogen column of $N_H = 10^{20}~{\\rm cm}^{-2}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "L = [0.1, 0.2, -0.3]  # normal vector\n",
    "N = [0.0, 1.0, 0.0]  # north vector\n",
    "n_events = pyxsim.project_photons(\n",
    "    \"snap_132_photons\",\n",
    "    \"snap_132_events\",\n",
    "    L,\n",
    "    (45.0, 30.0),\n",
    "    absorb_model=\"wabs\",\n",
    "    nH=0.01,\n",
    "    north_vector=N,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now that we have a set of \"events\" on the sky, we can read them in and write them to a SIMPUT file:"
   ]
  },
  {
   "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 use a small \n",
    "exposure time (100 ks instead of 300 ks), and observe it with the as-launched ACIS-I model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.instrument_simulator(\n",
    "    \"snap_132_events.h5\",\n",
    "    \"evt.fits\",\n",
    "    (100.0, \"ks\"),\n",
    "    \"chandra_acisi_cy0\",\n",
    "    [45.0, 30.0],\n",
    "    overwrite=True,\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.5 and 2.0 keV:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "soxs.write_image(\"evt.fits\", \"img.fits\", emin=0.5, emax=2.0, 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.fits\", stretch=\"sqrt\", cmap=\"arbre\", vmin=0.0, vmax=6.0, width=0.2)"
   ]
  }
 ],
 "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.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}