{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4a3eda48",
   "metadata": {},
   "source": [
    "# Extinction curves from custom grain populations\n",
    "\n",
    "## Initial parameters\n",
    "\n",
    "Extinction calculations need a grid of photon energies or wavelengths, and (optionally) a dust mass column so that the output can be expressed as an optical depth rather than an efficiency. Here we set up an X-ray energy grid and an optical/UV wavelength grid, and compute a dust mass column from a gas column density and a dust-to-gas mass ratio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9711a11d",
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "446eb8ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "xray_grid = np.logspace(-1, 1, 30) * u.keV\n",
    "UV_IR_grid = np.logspace(-1, 1.5, 100) * u.micron\n",
    "\n",
    "N_H = 1.e22  # hydrogen column density, cm^-2\n",
    "dust_to_gas = 0.009  # dust-to-gas mass ratio\n",
    "\n",
    "# Dust mass column, in g cm^-2\n",
    "# The factor 1.4 accounts for the presence of helium\n",
    "dust_mass_column = N_H * 1.4 * c.m_p.to('g').value * dust_to_gas\n",
    "\n",
    "print(\"Dust mass column: {:.1e} g cm^-2\".format(dust_mass_column))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e1e2581",
   "metadata": {},
   "source": [
    "## Grain Populations\n",
    "\n",
    "An `xdust` grain population ties together three ingredients: a grain **size distribution**, a grain **composition** (the optical constants of the material), and a **scattering model** (the physics used to compute how light interacts with the grains).\n",
    "\n",
    "By the end of this tutorial, you will be able to:\n",
    "\n",
    "- construct a `SingleGrainPop`, describing one grain composition and size distribution\n",
    "- compute and plot its extinction, scattering, and absorption efficiencies\n",
    "- combine several `SingleGrainPop` objects into a `GrainPop`, and work with it like a dictionary\n",
    "- use the built-in helper functions `xdust.make_MRN` and `xdust.make_MRN_RGDrude` for common dust models\n",
    "\n",
    "### Constructing a SingleGrainPop\n",
    "\n",
    "`SingleGrainPop` accepts short string names for the size distribution, composition, and scattering model, so you can build a grain population in one line.\n",
    "\n",
    "**Size distributions:**\n",
    "\n",
    "- `'Grain'` -- a single grain size\n",
    "- `'Powerlaw'`\n",
    "- `'ExpCutoff'`\n",
    "- `'Astrodust'` -- Astrodust size distribution from Hensley & Draine (2022)\n",
    "- `'WD01'` -- Weingartner & Draine (2001) size distributions\n",
    "\n",
    "Any additional keyword arguments passed to `SingleGrainPop` are forwarded to the size distribution class.\n",
    "\n",
    "**Compositions:**\n",
    "\n",
    "- `'Drude'` -- the Drude approximation for the complex index of refraction\n",
    "- `'Silicate'` -- silicate properties from Draine (2003)\n",
    "- `'Graphite'` -- perpendicular graphite properties from Draine (2003)\n",
    "\n",
    "**Scattering models:**\n",
    "\n",
    "- `'RG'` -- the Rayleigh-Gans approximation (relevant for X-rays)\n",
    "- `'Mie'` -- the Mie scattering algorithm of Bohren & Huffman"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daf01df7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A power-law size distribution of silicate grains, using Mie scattering\n",
    "silicate_pop = xdust.SingleGrainPop('Powerlaw', 'Silicate', 'Mie', md=dust_mass_column)\n",
    "\n",
    "# Run the extinction calculation for the UV-optical wavelength grid\n",
    "silicate_pop.calculate_ext(UV_IR_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ece558d",
   "metadata": {},
   "source": [
    "The built-in `plot_ext` method plots the extinction properties of the grain population, integrated over its grain size distribution. It takes two required arguments: a matplotlib axis, and a string specifying which property to plot:\n",
    "\n",
    "- `'all'` -- scattering, absorption, and extinction\n",
    "- `'sca'` -- scattering only\n",
    "- `'abs'` -- absorption only\n",
    "- `'ext'` -- extinction only (extinction = absorption + scattering)\n",
    "\n",
    "The `unit` keyword changes the units on the x-axis. Any other keyword arguments are passed to `matplotlib.pyplot.legend` when plotting `'all'`, or to `matplotlib.pyplot.plot` otherwise.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "003b938a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "silicate_pop.plot_ext(ax, 'all', unit='nm', frameon=False)\n",
    "plt.loglog()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c6399a3",
   "metadata": {},
   "source": [
    "Keyword arguments can be used to customize the appearance of each curve individually.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a13b88a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "silicate_pop.plot_ext(ax, 'ext', color='g', lw=3, label='Extinction')\n",
    "silicate_pop.plot_ext(ax, 'sca', color='b', lw=2, label='Scattering')\n",
    "silicate_pop.plot_ext(ax, 'abs', color='r', lw=1, label='Absorption')\n",
    "plt.ylabel(r'$\\tau$', size=16)\n",
    "plt.loglog()\n",
    "plt.legend(loc='upper right', frameon=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daaa02f6",
   "metadata": {},
   "source": [
    "### Customizing a SingleGrainPop\n",
    "\n",
    "The string shortcuts above initialize the composition and scattering model with their default parameters. To change a setting -- for example, a non-default value for a composition's `rho` -- construct the composition object yourself and pass it in directly.\n",
    "\n",
    "In fact, `SingleGrainPop` accepts either a string shortcut *or* an actual size distribution ({ref}`graindist`), composition ({ref}`composition`), or scattering model ({ref}`scatmodels`), for each of its three arguments.\n",
    "\n",
    "The example below creates separate populations for the perpendicular and parallel orientations of graphitic grains (see Draine 2003 for details), which requires constructing `CmGraphite` objects directly since there is no separate string shortcut for each orientation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f81956",
   "metadata": {},
   "outputs": [],
   "source": [
    "from xdust.graindist.composition import CmGraphite\n",
    "\n",
    "graphite_perp = xdust.SingleGrainPop('Powerlaw', CmGraphite(orient='perp'), 'Mie')\n",
    "graphite_para = xdust.SingleGrainPop('Powerlaw', CmGraphite(orient='para'), 'Mie')\n",
    "\n",
    "graphite_perp.calculate_ext(UV_IR_grid)\n",
    "graphite_para.calculate_ext(UV_IR_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75ae1699",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "graphite_perp.plot_ext(ax, 'ext', unit='nm', color='k', ls='-', label='Perpendicular')\n",
    "graphite_para.plot_ext(ax, 'ext', unit='nm', color='k', ls='--', label='Parallel')\n",
    "plt.legend()\n",
    "plt.loglog()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "356ef264",
   "metadata": {},
   "source": [
    "## Combining multiple SingleGrainPop objects into a GrainPop\n",
    "\n",
    "A `GrainPop` combines several `SingleGrainPop` objects and runs extinction calculations on all of them with a single method call. It behaves like a dictionary: each `SingleGrainPop` it contains is accessed by a key you choose when constructing it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d3513f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine the two graphitic grain populations initialized above into a GrainPop dictionary\n",
    "gra_dust_mix = xdust.GrainPop([graphite_perp, graphite_para], keys=['gra_perp', 'gra_para'])\n",
    "\n",
    "# Calculate the extinction for every SingleGrainPop in the mixture, in one call\n",
    "gra_dust_mix.calculate_ext(UV_IR_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d683b68a",
   "metadata": {},
   "source": [
    "The `GrainPop.info()` method prints a short summary of the `SingleGrainPop` objects it contains.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14bb32c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "gra_dust_mix.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0387f861",
   "metadata": {},
   "source": [
    "The `GrainPop` extinction curves can be accessed individually or plotted in sum using the `plot_ext` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e9387ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "\n",
    "# Each SingleGrainPop is accessed using the key assigned above\n",
    "gra_dust_mix['gra_perp'].plot_ext(ax, 'ext', color='g', label='Graphite (perpendicular)')\n",
    "gra_dust_mix['gra_para'].plot_ext(ax, 'ext', color='b', label='Graphite (parallel)')\n",
    "\n",
    "# The GrainPop itself can be plotted the same way as a SingleGrainPop --\n",
    "# this shows the sum of the silicate and graphite extinction\n",
    "gra_dust_mix.plot_ext(ax, 'ext', color='k', lw=2, label='Total')\n",
    "\n",
    "ax.legend(loc='upper right', frameon=False)\n",
    "plt.loglog()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "516fe636",
   "metadata": {},
   "source": [
    "## Helper functions for common dust models\n",
    "\n",
    "The `xdust.grainpop` module also provides ready-made `GrainPop` objects for two common dust models, so you don't need to assemble them by hand.\n",
    "\n",
    "### make_MRN\n",
    "\n",
    "`make_MRN` returns a mixture of silicate and graphite grains following a power-law size distribution with a maximum grain size of 0.3 micron (Mathis, Rumpl & Nordsieck 1977). Following Draine's recommendation, the graphitic grain population is split 1/3 parallel and 2/3 perpendicular; by default, 60% of the dust mass is silicate, following the recommendation in Corrales et al. (2016). Use the `fsil` keyword to change the silicate mass fraction.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "845a8588",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrn = xdust.make_MRN(md=dust_mass_column)\n",
    "mrn.calculate_ext(UV_IR_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d641dc53",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "mrn.plot_ext(ax, 'all', frameon=False)\n",
    "plt.loglog()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f85ce990",
   "metadata": {},
   "source": [
    "The individual components can be plotted alongside the total by looping over `GrainPop.keys`, the same way you would iterate over the keys of a dictionary.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "599abf22",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "mrn.plot_ext(ax, 'ext', color='k', lw=2, label='total')\n",
    "\n",
    "print(\"GrainPop keys:\", mrn.keys)\n",
    "\n",
    "for key in mrn.keys:\n",
    "    mrn[key].plot_ext(ax, 'ext', ls='--', label=key)\n",
    "\n",
    "ax.legend(loc='upper right', frameon=False)\n",
    "plt.loglog()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da8a2d12",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrn.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68e2ae27",
   "metadata": {},
   "source": [
    "### make_MRN_RGDrude\n",
    "\n",
    "The Drude approximation treats the complex index of refraction as if the solid were a mass of free electrons, making it relatively insensitive to the specific compound. `make_MRN_RGDrude` returns a `SingleGrainPop` using the `CmDrude` composition together with the Rayleigh-Gans scattering approximation, which is most relevant at X-ray wavelengths.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34e93346",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrn_xray = xdust.make_MRN_RGDrude(md=dust_mass_column)\n",
    "mrn_xray.calculate_ext(xray_grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c6f9673",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = plt.subplot(111)\n",
    "mrn_xray.plot_ext(ax, 'all')\n",
    "plt.loglog()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26a02c1e",
   "metadata": {},
   "source": [
    "The plot above shows that the Rayleigh-Gans-plus-Drude approximation produces an extinction model that decays smoothly as $E^{-2}$. There is no absorption component in the RG-Drude model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ba4d1b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mrn_xray.info()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}