{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "12e0a28d",
   "metadata": {},
   "source": [
    "# Modeling X-ray dust scattering halos\n",
    "\n",
    "The `xdust.halos` module computes the surface brightness profile of an X-ray dust scattering halo under a variety of geometric assumptions. This functionality is provided by two subclasses of `xdust.halos.Halo`:\n",
    "\n",
    "- `xdust.halos.UniformGalHalo` assumes the scattering dust is distributed uniformly along the line of sight between the telescope and the X-ray point source.\n",
    "- `xdust.halos.ScreenGalHalo` assumes the scattering dust is confined to an infinitesimally thin screen at position $x = 1 - d/D$, where $d$ is the distance from the telescope to the screen and $D$ is the distance from the telescope to the X-ray point source.\n",
    "\n",
    "See the {ref}`halos` module reference for the full API, including the `UniformGalHaloCP15` and `ScreenGalHaloCP15` shortcuts for the semi-analytic model of Corrales & Paerels (2015).\n",
    "\n",
    "By the end of this tutorial, you will be able to:\n",
    "\n",
    "- initialize a `UniformGalHalo` and a `ScreenGalHalo` over an energy and angle grid\n",
    "- run a scattering halo calculation for a given `GrainPop`\n",
    "- verify a halo calculation by integrating its surface brightness profile\n",
    "- compute the halo intensity produced by an X-ray point-source spectrum\n",
    "- slice a `Halo` object by energy, the same way you would slice a numpy array\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15ffd05a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import astropy.units as u\n",
    "import astropy.constants as c\n",
    "\n",
    "import xdust\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082acd13",
   "metadata": {},
   "source": [
    "## Setting up a scattering halo\n",
    "\n",
    "To initialize a scattering halo object, we need a grid of photon energies and a grid of angular distances (the angles at which the scattering halo is visible around the point source).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "740bf117",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_energy, n_theta = 50, 200\n",
    "energy_grid = np.logspace(-1, 1, n_energy) * u.keV\n",
    "theta_grid = np.logspace(-1, 4, n_theta) * u.arcsec\n",
    "\n",
    "uniform_halo = xdust.halos.UniformGalHalo(energy_grid, theta_grid)\n",
    "screen_halo = xdust.halos.ScreenGalHalo(energy_grid, theta_grid)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4827a267",
   "metadata": {},
   "source": [
    "To run the calculation, each halo also needs to know what kind of dust grains are doing the scattering -- their size distribution and composition, stored in a `xdust.GrainPop` object.\n",
    "\n",
    "Below, we use the Rayleigh-Gans-plus-Drude approximation, which treats a dust grain as a cloud of free electrons and is therefore relatively insensitive to grain composition (unlike Mie scattering). The `xdust.make_MRN_RGDrude` shortcut builds a standard power-law grain size distribution (Mathis, Rumpl & Nordsieck 1977) using the RG-Drude approximation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51df2c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "dust_pop = xdust.make_MRN_RGDrude(md=1.e-6)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab5cae8",
   "metadata": {},
   "source": [
    "Now we run the `calculate` method for each halo, passing in the dust population.\n",
    "\n",
    "For the screen halo, we choose $x = 0.33$, meaning the dust screen sits two-thirds of the way along the line of sight between the telescope and the background X-ray point source. The `%%time` cell magic prints how long each calculation takes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee2a90ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "screen_halo.calculate(dust_pop, x=0.33)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22ff390d",
   "metadata": {},
   "source": [
    "For dust distributed uniformly along the line of sight, the calculation instead integrates the scattering profile over many positions between the telescope and the point source. The `nx` keyword sets how fine that integration grid is -- try changing it to see how it affects both the calculation speed and the shape of the resulting profile (plotted below).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18969f51",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "uniform_halo.calculate(dust_pop, nx=100)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62ea9bc1",
   "metadata": {},
   "source": [
    "## Plotting the surface brightness profile\n",
    "\n",
    "The results are stored in the `norm_int` attribute, with shape `(len(energy_grid), len(theta_grid))` -- the normalized intensity (surface brightness divided by flux) as a function of energy and angle. To see a single profile, pick an energy index.\n",
    "\n",
    "Below we plot an arbitrary energy index; try other values to see how the halo profile shape changes with energy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95fea26",
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 10\n",
    "\n",
    "plt.plot(theta_grid, screen_halo.norm_int[i, :], label='Screen halo')\n",
    "plt.plot(theta_grid, uniform_halo.norm_int[i, :], label='Uniform halo')\n",
    "plt.loglog()\n",
    "plt.legend()\n",
    "\n",
    "plt.title(\"Photon Energy = {:.2f}\".format(energy_grid[i]))\n",
    "plt.xlabel(r\"Angular Distance from Point Source ({})\".format(theta_grid.unit))\n",
    "plt.ylabel(r\"Normalized Intensity ({})\".format(screen_halo.norm_int.unit))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "719c478f",
   "metadata": {},
   "source": [
    "## Verifying the halo calculation\n",
    "\n",
    "As a check, we can integrate the surface brightness profile stored in `norm_int` and confirm it reproduces the scattering optical depth (`taux`) at each energy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca53cab1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A 2D grid of theta values, matching the shape of norm_int\n",
    "theta_grid_2d = np.repeat(theta_grid.reshape(1, n_theta), n_energy, axis=0)\n",
    "\n",
    "# Integrate the surface brightness profile over solid angle: 2*pi*theta*dtheta.\n",
    "# Because the angles are small, sin(theta) -> theta, and since norm_int has\n",
    "# units of arcsec^-2, theta_grid_2d can stay in arcsec.\n",
    "integrated_uniform = np.trapezoid(\n",
    "    uniform_halo.norm_int * 2.0 * np.pi * theta_grid_2d, theta_grid, axis=1)\n",
    "integrated_screen = np.trapezoid(\n",
    "    screen_halo.norm_int * 2.0 * np.pi * theta_grid_2d, theta_grid, axis=1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "258c6101",
   "metadata": {},
   "source": [
    "The plot below shows that integrating each normalized halo profile reproduces its `taux` value -- as expected in the optically thin regime ($\\tau_{\\rm sca} < 1$). Note that the uniform and screen halos share the same `taux` values, even though their profile shapes differ.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f703a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(energy_grid, uniform_halo.taux, 'rs', label='Uniform (taux)')\n",
    "plt.plot(energy_grid, integrated_uniform, color='k', label='Uniform (integrated)')\n",
    "\n",
    "plt.plot(energy_grid, screen_halo.taux, 'bo', label='Screen (taux)')\n",
    "plt.plot(energy_grid, integrated_screen, color='gray', label='Screen (integrated)')\n",
    "\n",
    "plt.loglog()\n",
    "plt.legend()\n",
    "\n",
    "plt.xlabel(str(energy_grid.unit))\n",
    "plt.ylabel(r'X-ray Scattering Opacity ($\\tau_{\\rm sca}$)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cde0a552",
   "metadata": {},
   "source": [
    "We can also plot the fractional difference between the integrated profile and `taux` directly. A value of 0 means the integration exactly reproduces `taux`; in general this calculation is accurate to a few percent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aec475f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(energy_grid, integrated_uniform / uniform_halo.taux - 1.0, label='Uniform')\n",
    "plt.plot(energy_grid, integrated_screen / screen_halo.taux - 1.0, label='Screen')\n",
    "plt.legend()\n",
    "plt.xlabel(str(energy_grid.unit))\n",
    "plt.ylabel('Fractional Difference (integrated / taux - 1.0)')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee97d358",
   "metadata": {},
   "source": [
    "## Computing halo intensity from a point-source spectrum\n",
    "\n",
    "The `Halo.calculate_intensity` method computes the total scattering halo intensity, summed over the energy grid stored in `Halo.lam`.\n",
    "\n",
    "Below, we construct a model spectrum with a shape similar to many Galactic X-ray binaries. `calculate_intensity` integrates using `numpy.sum`, so the input flux must be *bin-integrated* -- each point in the spectrum represents the total counts (or energy) within an energy bin centered on the corresponding `lam` value, not a specific flux density. Bin-integrated flux has units like counts/s/cm$^2$ (as opposed to a specific flux, which would have units of counts/s/cm$^2$/keV).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9dd3933",
   "metadata": {},
   "outputs": [],
   "source": [
    "# An approximation of a typical Galactic X-ray binary spectrum\n",
    "source_flux = (1.0 * np.power(energy_grid.value, -2.5) *\n",
    "               np.exp(-0.1 * np.power(energy_grid.value, -3.5)) *\n",
    "               u.Unit('ct s^-1 cm^-2'))\n",
    "\n",
    "plt.plot(energy_grid, source_flux)\n",
    "plt.loglog()\n",
    "plt.xlabel(str(energy_grid.unit))\n",
    "plt.ylabel(str(source_flux.unit))\n",
    "plt.ylim(1.e-12, 1.e3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "493cce5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "screen_halo.calculate_intensity(source_flux)\n",
    "uniform_halo.calculate_intensity(source_flux)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0361d80d",
   "metadata": {},
   "source": [
    "The result is stored in each halo's `intensity` attribute -- the surface brightness profile summed over all energies, as a 1D array with the same length as `theta_grid`. Note that the units carry through automatically from the input spectrum and `norm_int`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d17c0e46",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(theta_grid, screen_halo.intensity, label='Screen halo')\n",
    "plt.plot(theta_grid, uniform_halo.intensity, label='Uniform halo')\n",
    "plt.loglog()\n",
    "plt.legend()\n",
    "\n",
    "plt.xlabel(r\"Angular Distance from Point Source ({})\".format(theta_grid.unit))\n",
    "plt.ylabel(str(screen_halo.intensity.unit))\n",
    "plt.title(\"Scattering halo surface brightness ({:.2f} - {:.2f} keV)\".format(\n",
    "    energy_grid[0].to('keV').value, energy_grid[-1].to('keV').value))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18b012a0",
   "metadata": {},
   "source": [
    "To get the total flux contained in the scattering halo, integrate `intensity` over solid angle (again, $2\\pi\\theta\\,d\\theta$). Because both halos share the same `GrainPop` and dust mass column, they integrate to the same total flux, despite having different profile shapes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "089200a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_flux_uniform = np.trapezoid(uniform_halo.intensity * 2.0 * np.pi * theta_grid, theta_grid)\n",
    "total_flux_screen = np.trapezoid(screen_halo.intensity * 2.0 * np.pi * theta_grid, theta_grid)\n",
    "\n",
    "print(\"Total flux in uniform scattering halo: {:.3f}\".format(total_flux_uniform))\n",
    "print(\"Total flux in screen scattering halo: {:.3f}\".format(total_flux_screen))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e93cfa6",
   "metadata": {},
   "source": [
    "After `calculate_intensity` has run, the `fhalo` attribute gives the flux spectrum of the scattering halo itself, as a function of energy.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5642bbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(energy_grid, source_flux, color='k', label='Input flux (absorbed)')\n",
    "plt.plot(energy_grid, uniform_halo.fhalo, color='g', label='Uniform halo flux')\n",
    "plt.plot(energy_grid, screen_halo.fhalo, color='b', ls='--', label='Screen halo flux')\n",
    "plt.loglog()\n",
    "plt.legend()\n",
    "plt.ylim(1.e-12, 1e3)\n",
    "plt.xlabel(str(energy_grid.unit))\n",
    "plt.ylabel('Flux ({})'.format(uniform_halo.fhalo.unit))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e65ce83",
   "metadata": {},
   "source": [
    "## Indexing and slicing a Halo\n",
    "\n",
    "A `Halo` object can be indexed or sliced much like a numpy array. Slicing returns a new `Halo` covering a sub-range of the original energy grid, with all of its calculated attributes (including `intensity`) carried over automatically.\n",
    "\n",
    "The slice bounds must be given in the same units used to define `Halo.lam` -- if `lam` was defined in keV, you cannot slice using Angstrom values.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53ad2374",
   "metadata": {},
   "outputs": [],
   "source": [
    "low_energy_halo = uniform_halo[:3.0]   # everything below 3 keV\n",
    "high_energy_halo = uniform_halo[3.0:]   # everything above 3 keV\n",
    "\n",
    "print(\"Low-energy slice: {:.2f} - {:.2f} keV\".format(\n",
    "    low_energy_halo.lam[0].value, low_energy_halo.lam[-1].value))\n",
    "print(\"High-energy slice: {:.2f} - {:.2f} keV\".format(\n",
    "    high_energy_halo.lam[0].value, high_energy_halo.lam[-1].value))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29c95d8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(low_energy_halo.lam, low_energy_halo.taux, 'rs', label='Low-energy slice')\n",
    "plt.plot(high_energy_halo.lam, high_energy_halo.taux, 'gs', label='High-energy slice')\n",
    "plt.plot(uniform_halo.lam, uniform_halo.taux, 'k-', label='Original halo')\n",
    "plt.loglog()\n",
    "plt.xlabel(str(low_energy_halo.lam.unit))\n",
    "plt.ylabel('X-ray Scattering Opacity (taux)')\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7611a87c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(low_energy_halo.theta, low_energy_halo.intensity,\n",
    "         label='{:.2f}-{:.2f} keV'.format(low_energy_halo.lam[0].value, low_energy_halo.lam[-1].value))\n",
    "plt.plot(high_energy_halo.theta, high_energy_halo.intensity,\n",
    "         label='{:.2f}-{:.2f} keV'.format(high_energy_halo.lam[0].value, high_energy_halo.lam[-1].value))\n",
    "plt.loglog()\n",
    "plt.legend()\n",
    "plt.xlabel(str(low_energy_halo.theta.unit))\n",
    "plt.ylabel(str(high_energy_halo.intensity.unit))\n",
    "plt.title(\"Uniform halo surface brightness profile\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}