{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Power Law"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The second example shows how to set up a power-law spectral source, as well as how to add two sets of photons together. It also shows how to use the `make_spectrum` functionality of a source model.\n",
    "\n",
    "This example will also briefly show how to set up a mock dataset \"in memory\" using yt. For more details on how to do this, check out [the yt docs on in-memory datasets](http://yt-project.org/doc/examining/generic_array_data.html).\n",
    "\n",
    "Load up the necessary modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "\n",
    "matplotlib.rc(\"font\", size=18, family=\"serif\")\n",
    "import yt\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from yt.units import mp, keV, kpc\n",
    "import pyxsim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The cluster we set up will be a simple isothermal $\\beta$-model system, with a temperature of 4 keV. We'll set it up on a uniform grid of 2 Mpc and 256 cells on a side. The parameters of the model are:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "R = 1000.0  # radius of cluster in kpc\n",
    "r_c = 100.0  # scale radius of cluster in kpc\n",
    "rho_c = 1.673e-26  # scale density in g/cm^3\n",
    "beta = 1.0  # beta parameter\n",
    "kT = 4.0  # cluster temperature in keV\n",
    "nx = 256\n",
    "ddims = (nx, nx, nx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "and we set up the density and temperature arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "x, y, z = np.mgrid[-R : R : nx * 1j, -R : R : nx * 1j, -R : R : nx * 1j]\n",
    "r = np.sqrt(x**2 + y**2 + z**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "dens = np.zeros(ddims)\n",
    "dens[r <= R] = rho_c * (1.0 + (r[r <= R] / r_c) ** 2) ** (-1.5 * beta)\n",
    "dens[r > R] = 0.0\n",
    "temp = (kT * keV).to_value(\"K\", \"thermal\") * np.ones(ddims)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next, we will take the density and temperature arrays and put them into a dictionary with their units, where we will also set up velocity fields set to zero. Then, we'll call the yt function `load_uniform_grid` to set this up as a full-fledged yt dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "data = {}\n",
    "data[\"density\"] = (dens, \"g/cm**3\")\n",
    "data[\"temperature\"] = (temp, \"K\")\n",
    "data[\"velocity_x\"] = (np.zeros(ddims), \"cm/s\")\n",
    "data[\"velocity_y\"] = (np.zeros(ddims), \"cm/s\")\n",
    "data[\"velocity_z\"] = (np.zeros(ddims), \"cm/s\")\n",
    "\n",
    "bbox = np.array(\n",
    "    [[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]]\n",
    ")  # The bounding box of the domain in code units\n",
    "\n",
    "L = (2.0 * R * kpc).to_value(\"cm\")\n",
    "\n",
    "# We have to set default_species_fields=\"ionized\" because\n",
    "# we didn't create any species fields above\n",
    "ds = yt.load_uniform_grid(data, ddims, L, bbox=bbox, default_species_fields=\"ionized\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The next thing we have to do is specify a derived field for the normalization of the power-law emission. This could come from a variety of sources, for example, relativistic cosmic-ray electrons. For simplicity, we're not going to assume a specific model, except that we will only specify that the source of the power law emission is proportional to the gas mass in each cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "norm = yt.YTQuantity(1.0e-19, \"photons/s/keV\")\n",
    "\n",
    "\n",
    "def _power_law_emission(field, data):\n",
    "    return norm * data[\"cell_mass\"] / (1.0 * mp)\n",
    "\n",
    "\n",
    "ds.add_field(\n",
    "    (\"gas\", \"power_law_emission\"),\n",
    "    function=_power_law_emission,\n",
    "    sampling_type=\"local\",\n",
    "    units=\"photons/s/keV\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "where we have normalized the field arbitrarily. Note that the emission field for a power-law model is a bit odd in that it is technically a specific *luminosity* for the cell. This is done primarily for simplicity in designing the underlying algorithm.  \n",
    "\n",
    "Now, let's set up a sphere to collect photons from:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "sp = ds.sphere(\"c\", (0.5, \"Mpc\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "And set the parameters for the initial photon sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "A = yt.YTQuantity(500.0, \"cm**2\")\n",
    "exp_time = yt.YTQuantity(1.0e5, \"s\")\n",
    "redshift = 0.03"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Set up two source models, a thermal model and a power-law model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "therm_model = pyxsim.CIESourceModel(\"apec\", 1, 80.0, 5000, Zmet=0.3)\n",
    "plaw_model = pyxsim.PowerLawSourceModel(1.0, 1, 80.0, \"power_law_emission\", 1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, generate the photons for each source model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "ntp, ntc = pyxsim.make_photons(\"therm_photons\", sp, redshift, A, exp_time, therm_model)\n",
    "npp, npc = pyxsim.make_photons(\"plaw_photons\", sp, redshift, A, exp_time, plaw_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now, we want to project the photons along a line of sight. We'll specify the `\"wabs\"` model for foreground galactic absorption. We'll create events from the total set of photons as well as the power-law only set, to see the difference between the two."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "nte = pyxsim.project_photons(\n",
    "    \"therm_photons\", \"therm_events\", \"x\", (30.0, 45.0), absorb_model=\"wabs\", nH=0.02\n",
    ")\n",
    "npe = pyxsim.project_photons(\n",
    "    \"plaw_photons\", \"plaw_events\", \"x\", (30.0, 45.0), absorb_model=\"wabs\", nH=0.02\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, create energy spectra for both sets of events. We won't bother convolving with instrument responses here, because we just want to see what the spectra look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "et = pyxsim.EventList(\"therm_events.h5\")\n",
    "ep = pyxsim.EventList(\"plaw_events.h5\")\n",
    "et.write_spectrum(\"therm_spec.fits\", 1, 80.0, 5000, overwrite=True)\n",
    "ep.write_spectrum(\"plaw_spec.fits\", 1, 80.0, 5000, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "To visualize the spectra, we'll load them up using [AstroPy's FITS I/O](http://docs.astropy.org/en/stable/io/fits/) and use [Matplotlib](http://matplotlib.org) to plot the spectra: "
   ]
  },
  {
   "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(\"plaw_spec.fits\")\n",
    "plt.figure(figsize=(9, 7))\n",
    "plt.loglog(\n",
    "    f2[\"SPECTRUM\"].data[\"ENERGY\"], f2[\"SPECTRUM\"].data[\"COUNTS\"], label=\"Power-Law\"\n",
    ")\n",
    "plt.loglog(\n",
    "    f2[\"SPECTRUM\"].data[\"ENERGY\"],\n",
    "    f1[\"SPECTRUM\"].data[\"COUNTS\"] + f2[\"SPECTRUM\"].data[\"COUNTS\"],\n",
    "    label=\"Power-Law+Thermal\",\n",
    ")\n",
    "plt.xlim(1, 50)\n",
    "plt.ylim(1, 2.0e4)\n",
    "plt.legend()\n",
    "plt.xlabel(\"E (keV)\")\n",
    "plt.ylabel(\"counts/bin\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As you can see, the orange line shows a sum of a thermal and power-law spectrum, the latter most prominent at high energies. The blue line shows the power-law spectrum. Both spectra are absorbed at low energies, thanks to the Galactic foreground absorption. "
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Finally, we show how to make a spectrum of the entire sphere object using the `make_spectrum` method of the `PowerLawSourceModel`. To do this, we simply supply `make_spectrum` with the sphere object, the minimum and maximum energies in the spectrum, the number of bins in the spectrum, and optionally the redshift. The latter ensures we get a flux spectrum in the observer frame. `make_spectrum` returns a SOXS `Spectrum` object, which you can do other things with, such as apply foreground absorption. "
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "emin = 1.0\n",
    "emax = 80.0\n",
    "nbins = 5000\n",
    "pspec = plaw_model.make_spectrum(sp, emin, emax, nbins, redshift=redshift)\n",
    "pspec.apply_foreground_absorption(0.02)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "To show this is the same spectrum that we obtained above, we compare it to the above spectrum we derived from the mock photons we generated before (converting it to a flux first)."
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "emid = f2[\"SPECTRUM\"].data[\"ENERGY\"]\n",
    "de = np.diff(emid)[0]\n",
    "# We have to convert counts/bin to counts/s/cm**2/keV here\n",
    "flux = f2[\"SPECTRUM\"].data[\"COUNTS\"] / (A.v * exp_time.v * de)\n",
    "plt.figure(figsize=(9, 7))\n",
    "plt.loglog(emid, flux)\n",
    "plt.loglog(pspec.emid, pspec.flux, lw=2)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "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
}