{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ex 3: Manipulating Spectra Part 2 - `Spectra`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "import sys\n",
    "sys.path.append(\"..\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%autoreload\n",
    "\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.rc_file('matplotlibrc')\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "from scipy.integrate import quad\n",
    "\n",
    "import darkhistory.physics as phys\n",
    "import darkhistory.utilities as utils\n",
    "import darkhistory.spec.spectools as spectools\n",
    "from darkhistory.spec.spectrum import Spectrum\n",
    "from darkhistory.spec.spectra import Spectra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `Spectra` Class - Introduction\n",
    "---------------------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) class is simply a collection of [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) objects. The main use of [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) is as a list of spectra that indexes spectra at different redshifts (e.g. the low-energy photon spectrum deposited at each redshift) or at different injection energies (e.g. the spectrum of photons produced by injecting an electron at different energies). \n",
    "\n",
    "To demonstrate how to use this class, consider the extragalactic photon flux coming from the smooth distribution of dark matter throughout the universe, decaying with some spectrum $dN/dE'$. The $'$ notation is used to indicate the spectrum *at production*. For a delta function photon spectrum from the decay, $dN/dE' = 2 \\delta(E' - m_{\\text{DM}}/2)$, the photon flux (i.e. the number of photons received per area per energy per steradian) can be evaluated exactly as\n",
    "\n",
    "$$ \\frac{d^2 \\Phi}{d\\Omega \\, dE} = \\frac{1}{4\\pi} \\frac{\\Gamma \\Omega_{\\text{DM}} \\rho_c}{m_{\\text{DM}} H_0 \\sqrt{\\Omega_m}} \\cdot \\frac{4}{m_{\\text{DM}}} \\sqrt{\\frac{2E}{m_{\\text{DM}}}} \\frac{\\Theta(m_{\\text{DM}}/2 - E)}{\\sqrt{1 + \\kappa(2E/m_{\\text{DM}})^3}} $$\n",
    "\n",
    "where $\\Gamma$ is the decay width, $\\Omega_\\text{DM} \\rho_c$ and $\\Omega_m \\rho_c$ are the energy density of dark matter and all matter today, $\\rho_c$ is the critical density today, $H_0$ is the Hubble constant, and $\\kappa \\equiv \\Omega_\\Lambda/\\Omega_m \\sim 3$ is the ratio of the energy densities of dark energy to all matter. $\\Theta(x)$ is the usual step function. Instead of evaluating this analytically, we will now instead compute a discrete approximation, and see that we recover some approximation of the final analytic result. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main attributes of the [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) class are very similar to those of [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html), and are as follows: \n",
    "\n",
    "1. `eng`, length *M:* the energy abscissa of the spectra. All spectra stored within a single [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) class must have the same abscissae;\n",
    "\n",
    "2. `in_eng`, length *N:* the injection energy abscissa of the spectra. This should be initialized appropriately if the spectra correspond to secondary spectra of injected particles at different *injection energies*; otherwise, it is initialized to an array of -1;\n",
    "\n",
    "3. `rs`, length *N:* the redshift abscissa of the spectra. This should be initialized appropriately if the spectra corresponds to the same kind of output at different *redshifts*; otherwise, it is initialized to an array of -1; \n",
    "\n",
    "4. `grid_vals`, shape *(N,M)*: the 2D grid storing all of the data. The first dimension corresponds to either `in_eng` or `rs`, while the second dimension is `eng`; \n",
    "\n",
    "5. `spec_type`, which can be either `'N'` or `'dNdE'`, and determines how the data is stored in `grid_vals`, and\n",
    "\n",
    "6. `N_underflow` and `eng_underflow`, length *N*, to store the underflow in number and energy. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To begin, we will initialize a list of [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) objects, each containing a delta function at half the dark matter mass, taken to be $m_{\\text{DM}}$ = 1 GeV. This can be efficiently done by [*rebin_N_arr()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectools/darkhistory.spec.spectools.rebin_N_arr.html), which takes a list of energies (second input), a list containing the number of particles at that energy (first input), and an output energy abscissa (third input), and returns a [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) object with the output energy abscissa and spectrum approximating the initial particle input, conserving both total number and total energy. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll work in GeV for the this part of the code. \n",
    "mDM = 1\n",
    "\n",
    "# Energy and redshift abscissa.\n",
    "# Redshift is ordered from largest to smallest. \n",
    "eng = 10**np.arange(-3, 0, 0.05)\n",
    "rs_arr  = 10**np.arange(3, 0, -0.003)\n",
    "\n",
    "\n",
    "# Create the list of Spectrum objects. \n",
    "spec_arr = [\n",
    "    spectools.rebin_N_arr(np.array([2]), np.array([mDM/2]), eng)\n",
    "    for rs in rs_arr\n",
    "]\n",
    "\n",
    "# Assign the correct redshift to each of the Spectrum objects. \n",
    "# The first Spectrum object is assigned the largest redshift.\n",
    "for spec,rs in zip(spec_arr,rs_arr):\n",
    "    spec.rs = rs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to use the [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) class. Given a list of [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) objects with identical energy abscissa, we can create a [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "photon_specs = Spectra(spec_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object can also be directly initialized from a 2D array. Here, we initialize with a 3x4 matrix, with appropriate entries for *rs*, *in_eng* and *eng*: \n",
    "\n",
    "```python\n",
    "    test_arr = np.ones_like((3, 4))\n",
    "    test_spectra = Spectra(\n",
    "        test_arr, \n",
    "        eng=np.array([1, 10, 100, 1000]),  # must be length 4\n",
    "        in_eng=np.array([3, 30, 300])      # must be length 3\n",
    "    )\n",
    "    # test_spectra.rs will be automatically initialized to array([-1, -1, -1]). \n",
    "    # test_spectra.grid_vals == test_arr will be True. \n",
    "```\n",
    "\n",
    "Initializing with a collection of [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) objects which have been correctly initialized will correctly populate *rs*, *in_eng* and *eng* from these objects.\n",
    "\n",
    "A few simple functions exist to get attributes from the [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) objects stored in a [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object. The user can run the following lines to inspect these attributes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('The energy abscissa of each Spectrum in photon_specs is: ')\n",
    "print(photon_specs.eng)\n",
    "print('*******************************************************')\n",
    "print('The redshift of each Spectrum in photon_specs is: ')\n",
    "print(photon_specs.rs)\n",
    "print('*******************************************************')\n",
    "print('The injection energy of each Spectrum in photon_specs is: ')\n",
    "print(photon_specs.in_eng)\n",
    "print('*******************************************************')\n",
    "print('The 3rd Spectrum in photon_specs has N values (should add up to 2): ')\n",
    "print(photon_specs[2].N)\n",
    "print('*******************************************************')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point we simply have a list of spectra with a delta-function spectrum at $m_\\text{DM}/2$, produced at different redshifts. To turn this into a flux that is observed on Earth, we need to do some cosmology. \n",
    "\n",
    "The spectrum of photons produced per volume per time, is given by\n",
    "\n",
    "$$ \\frac{dN}{dV' dt' dE'} = \\frac{dN}{dE'} \\Gamma n_{\\chi,0}(1+z')^3 $$\n",
    "\n",
    "where $n_{\\chi,0}$ is the dark matter number density today. To convert this into a flux observed *today*, we need to relate the $'$ variables to the present day. We have $(1+z')^3 dV' = dV$ from the expansion, and $dE' = (1+z') dE$ from the redshifting of photons. Thus the flux observed today is\n",
    "\n",
    "$$ \n",
    "\\begin{align}\n",
    "\\frac{d^2 \\Phi}{d\\Omega dE} &= \\frac{1}{4\\pi}  \\int_0^\\infty dt' \\left. \\frac{dN}{dV' \\, dt' \\, dE'} \\right|_{E' = E(1+z)} \\frac{dV'}{dV} \\frac{dE'}{dE} \\\\\n",
    "&= \\frac{\\Gamma n_{\\chi,0}}{4\\pi} \\int_0^\\infty d \\log(1+z') \\frac{dt'}{d \\log(1+z')} \\left. \\frac{dN}{dE'} \\right|_{E' = E(1+z')} (1+z') \\\\\n",
    "&= \\frac{ \\Gamma n_{\\chi,0}}{4\\pi H_0 \\sqrt{\\Omega_m}} \\int_0^\\infty d \\log(1+z') \\left. \\frac{dN}{dE'} \\right|_{E' = E(1+z')} \\frac{1}{\\sqrt{1+z'} \\sqrt{1 + \\kappa(1+z')^{-3}}}\n",
    "\\end{align}\n",
    "$$\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now obtain the integral numerically. First, note that the spectrum \n",
    "\n",
    "$$ \\left. \\frac{dN}{dE'} \\right|_{E' = E(1+z')} = 2 \\delta(E(1+z') - m_\\text{DM}/2)$$ \n",
    "\n",
    "is simply all of the delta-function spectra in `photon_specs` redshifted from their starting redshifts (which are tabulated in `rs_arr`) down to $z = 0$, and then divided by $1+z'$. You can check that this extra factor is necessary based on the normalization of each spectrum (redshifting of spectra preserves total number in DarkHistory, and only shifts energies). \n",
    "\n",
    "We can then treat all of these spectra as approximating the produced spectrum within some bin of width $\\Delta \\log(1+z')$:\n",
    "\n",
    "$$\n",
    "\\frac{d^2 \\Phi}{d \\Omega dE} \\approx \\frac{\\Gamma n_{\\chi,0}}{4 \\pi H_0 \\sqrt{\\Omega_m}} \\sum_i \\left.\\frac{dN}{dE'}\\right|_{E'=E(1+z_i')} \\frac{\\Delta \\log(1+z_i')}{\\sqrt{1+z_i'} \\sqrt{1 + \\kappa(1+z_i')^{-3}}}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's perform the redshifting. This can be done with [*Spectra.redshift()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.redshift), which takes a list of floats to redshift to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "photon_specs.redshift(np.ones_like(photon_specs.rs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then we divide by the redshift value for the correct normalization. Note that you can manipulate [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) very naturally with standard mathematical operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "photon_specs /= rs_arr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, each of these spectra must be multiplied by a factor of\n",
    "\n",
    "$$ \\frac{\\Delta \\log(1+z_i)}{\\sqrt{1+z_i} \\sqrt{1 + \\kappa(1 + z_i)^{-3}}} $$\n",
    "\n",
    "To find $\\Delta \\log z_i$, we use the convenience function [*spectools.get_log_bin_width()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectools/darkhistory.spec.spectools.get_log_bin_width.html), which takes a redshift array and computes the log bin width of each bin. Note the minus sign, since our redshifts are currently stored in descending order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dlogz = -spectools.get_log_bin_width(rs_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now simply compute an array of this factor, and then multiply `photon_specs` by it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa = phys.omega_lambda/phys.omega_m\n",
    "\n",
    "fac = dlogz/(np.sqrt(rs_arr)*np.sqrt(1 + kappa*rs_arr**(-3)))\n",
    "photon_specs *= fac"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also multiply in the remaining prefactors (here we choose $\\Gamma$ to be in natural units):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This choice of the width (in GeV) is roughly a lifetime of 1e25 seconds.\n",
    "Gamma = 1e-50\n",
    "\n",
    "prefac = (\n",
    "    (1/(4*np.pi)) * Gamma * phys.rho_DM/(mDM*1e9)\n",
    "    / (phys.H0 * np.sqrt(phys.omega_m))\n",
    "    )*phys.c/(phys.hbar*1e-9)\n",
    "\n",
    "photon_specs *= prefac"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that all units of fundamental parameters in [*physics*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/darkhistory.physics.html) are given in eV (mass, energy and temperature), cm and seconds, and the final value of the photon flux has been multiplied by appropriate numerical factors to give `photon_specs` units of $\\text{cm}^{-2} \\text{ GeV}^{-1} \\text{ s}^{-1}$. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summing of Spectra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last remaining thing to do is to perform the integral by summing up the discrete spectra. This can be done using [*Spectra.sum_specs()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.sum_specs), which returns a [*Spectrum*](https://darkhistory.readthedocs.io/en/development/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) object: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_photon_flux = photon_specs.sum_specs()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now to compare this with the analytic expression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analytic_flux(eng):\n",
    "    factor = 4/mDM*np.sqrt(2*eng/mDM)*1/np.sqrt(1 + kappa*(2*eng/mDM)**3)\n",
    "    factor[eng > mDM/2] = 0\n",
    "    return np.array(prefac*factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure() \n",
    "\n",
    "plt.loglog()\n",
    "\n",
    "discrete_spec, = plt.plot(\n",
    "    obs_photon_flux.eng, obs_photon_flux.dNdE, label='Discrete Spectrum'\n",
    ")\n",
    "\n",
    "dNdE_analytic = analytic_flux(obs_photon_flux.eng)\n",
    "\n",
    "analytic_spec, = plt.plot(\n",
    "    obs_photon_flux.eng, dNdE_analytic, 'go', marker='o', markersize='6', \n",
    "    markerfacecolor='w', markevery=2,\n",
    "    label='Analytic Spectrum'\n",
    ")\n",
    "\n",
    "plt.legend(handles=[discrete_spec, analytic_spec])\n",
    "\n",
    "\n",
    "plt.title('Extragalactic Photon Flux')\n",
    "plt.xlabel('Photon Energy [GeV]')\n",
    "plt.ylabel(r'$d^2\\Phi/d\\Omega\\, dE$ [cm$^{-2}$ GeV$^{-1}$ sr$^{-1}$ s$^{-1}$]')\n",
    "plt.axis([1e-3, 1, 1e-6, 3e-4])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[*Spectra.sum_specs()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.sum_specs) can also be called with an argument specifying some weight to be assigned to each of the spectra. So the preceding analysis could have been done using the following code:\n",
    "\n",
    "```python\n",
    "# Create the list of Spectrum objects. \n",
    "spec_arr = [\n",
    "    spectools.rebin_N_arr(np.array([2/rs]), np.array([mDM/2]), eng)\n",
    "    for rs in rs_arr\n",
    "           ]\n",
    "\n",
    "# Assign the correct redshift to each of the Spectrum objects. \n",
    "# The first Spectrum object is assigned the largest redshift.\n",
    "for spec,rs in zip(spec_arr,rs_arr):\n",
    "    spec.rs = rs\n",
    "    # Redshifting\n",
    "    spec.redshift(1)\n",
    "    \n",
    "photon_flux_2 = Spectra(spec_arr)\n",
    "# Weighted sum\n",
    "obs_photon_flux_2 = photon_flux_2.sum_specs(fac*prefac)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpolation and Plotting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The two key advantages of using the  object is interpolation and plotting. Given a [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object with spectra corresponding to different redshifts, [*Spectra.at_rs()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.at_rs) takes in an array of new redshifts, and performs a linear interpolation over the whole [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) to give you a new [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object, with spectra at the new redshifts. \n",
    "\n",
    "To demonstrate this, we return to the CMB energy spectrum once again, and generate a [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object storing the number of CMB photons per unit energy per co-moving volume at several different redshifts. We do this by initializing the [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object with an empty array, and then using [*Spectra.append()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.append) to add more [*Spectrum*](https://darkhistory.readthedocs.io/en/development/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html) objects to it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the output energy abscissa. \n",
    "eng = 10**((np.arange(120)-90)*(1/10))\n",
    "# Define the redshift abscissa.\n",
    "rs_arr = 1+np.arange(100)\n",
    "\n",
    "# Analytic CMB spectrum\n",
    "\n",
    "CMB_specs = Spectra([])\n",
    "for rs in rs_arr:\n",
    "    spec = spectools.discretize(eng, phys.CMB_spec, phys.TCMB(1)*rs)\n",
    "    # Change the result to per co-moving volume\n",
    "    spec *= rs**3\n",
    "    spec.rs = rs\n",
    "    CMB_specs.append(spec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can interpolate to get the spectrum at any intermediate redshift within the abscissa of the [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object, using [*Spectra.at_rs()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.at_rs). This returns a new [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object with just the spectra at the specified redshift. The interpolation is done in log redshift. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rs_selected = np.array([2.495, 30.4903, 99.910])\n",
    "CMB_specs_selected = CMB_specs.at_rs(rs_selected)\n",
    "print('CMB_specs_selected has the following redshift abscissa: ', \n",
    "     CMB_specs_selected.rs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the result. [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) includes a function [*Spectra.plot()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html), meant to make plotting of multiple spectra at the same time easy. The axis of the plot is passed to the function, and keyword arguments to `matplotlib` can simply be passed to the function to change the plot properties. \n",
    "\n",
    "For a start, we call [*Spectra.plot()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) for `CMB_specs_selected` by passing just the axis and keyword arguments for plotting: this simply plots every spectrum inside it (**NB: Doing this for a large [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) object will take a long time!**). We overlay the analytic result for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analytic_spec = [\n",
    "    phys.CMB_spec(\n",
    "        CMB_specs_selected.eng, phys.TCMB(1)*rs\n",
    "    )*rs**3 for rs in rs_selected\n",
    "]\n",
    "\n",
    "ax = plt.subplot(1,1,1)\n",
    "plt.rc('text', usetex=True)\n",
    "\n",
    "spectra_plot = CMB_specs_selected.plot(ax, linewidth=2)\n",
    "\n",
    "plt.gca().set_prop_cycle(None)\n",
    "\n",
    "analytic_plot = [\n",
    "    plt.plot(\n",
    "        CMB_specs_selected.eng, spec, 'o',\n",
    "        marker='o', markersize='6', markevery=3, markerfacecolor='w'\n",
    "    ) for spec in analytic_spec\n",
    "]\n",
    "\n",
    "plt.legend(spectra_plot, ('1+z = 2.495','1+z = 30.4903', '1+z = 99.91'))\n",
    "\n",
    "\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "\n",
    "\n",
    "plt.title('CMB Spectrum')\n",
    "plt.xlabel('Photon Energy [eV]')\n",
    "plt.ylabel(r'$dn_\\gamma/dE$ [eV$^{-1}$ cm$^{-3}$]')\n",
    "plt.axis([1e-8, 100, 1e-7, 1e17])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[*Spectra.plot()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) can be used to plot selected spectra, for example at arbitrary redshifts within the redshift abscissa. Linear interpolation is done automatically to pick out the correct spectra. \n",
    "\n",
    "[*Spectra.plot()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) also takes the argument `fac`, which allows the user to multiply the spectra by some arbitrary factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.loglog()\n",
    "\n",
    "ax = plt.gca()\n",
    "\n",
    "rs_to_plot = np.array([90, 45, 22.5, 11.25])\n",
    "spectra_plot = CMB_specs.plot(ax, ind=rs_to_plot, indtype='rs', fac=CMB_specs.eng**2, linewidth=2)\n",
    "\n",
    "plt.legend(spectra_plot, ('1+z = 90','1+z = 45', '1+z = 22.5', '1+z = 11.25'))\n",
    "\n",
    "plt.gca().set_prop_cycle(None)\n",
    "\n",
    "analytic_spec = [\n",
    "    phys.CMB_spec(\n",
    "        CMB_specs.eng, phys.TCMB(1)*rs\n",
    "    )*CMB_specs.eng**2*rs**3 for rs in rs_to_plot\n",
    "]\n",
    "\n",
    "analytic_plot = [\n",
    "    plt.plot(\n",
    "        CMB_specs_selected.eng, spec, 'o',\n",
    "        marker='o', markersize='6', markevery=3, markerfacecolor='w'\n",
    "    ) for spec in analytic_spec\n",
    "]\n",
    "\n",
    "\n",
    "plt.title('CMB Spectrum')\n",
    "plt.xlabel('Photon Energy [eV]')\n",
    "plt.ylabel(r'$E^2 \\, dn_\\gamma/dE$ [eV cm$^{-3}$]')\n",
    "plt.axis([1e-8, 100, 1e-7, 1e17])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also make plots without interpolation, by selecting the spectra stored in the `Spectra` object by specifying the index. There are many different ways to specify what to plot, and the user should refer to the documentation to find out more. Below, we plot the spectra between two indices in steps: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "ax = plt.gca()\n",
    "ax.loglog()\n",
    "\n",
    "# Plots index 20 to 99, steps of 10 (i.e. 20, 30, 40, 50, 60, 70, 80, 90)\n",
    "spectra_plot = CMB_specs.plot(ax, ind=(20,100), step=10, indtype='ind', linewidth=2)\n",
    "\n",
    "plt.title('CMB Spectrum')\n",
    "plt.xlabel('Photon Energy [eV]')\n",
    "plt.ylabel(r'$dn_\\gamma/dE$ [eV$^{-1}$ cm$^{-3}$]')\n",
    "plt.axis([1e-8, 100, 1e5, 1e17])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `Spectra` - Sums, Integration over each Spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Integration over each spectrum can be done using [*Spectra.integrate_each_spec()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.integrate_each_spec), with the result being an array. There are many reasons to perform an integration over each spectrum; here, we will show only a trivial example, integrating over the spectrum to get the total number. To do this, we use [*Spectra.integrate_each_spec()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.integrate_each_spec) with a weight given by $E \\, d \\log E$. Like [*Spectrum*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectrum/darkhistory.spec.spectrum.Spectrum.html), [*Spectra*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html) also has [*Spectra.totN()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.totN) and [*Spectra.toteng()*](https://darkhistory.readthedocs.io/en/latest/_autosummary/darkhistory/spec/spectra/darkhistory.spec.spectra.Spectra.html?highlight=spectra#darkhistory.spec.spectra.Spectra.toteng) functions, which return an array of numbers instead, one for each spectrum. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_bin_width = spectools.get_log_bin_width(CMB_specs.eng)\n",
    "total_1 = CMB_specs.integrate_each_spec(CMB_specs.eng * log_bin_width)\n",
    "total_2 = CMB_specs.totN()\n",
    "print('Are the two results equal (up to floating point accuracy)?')\n",
    "print(np.all(np.isclose(total_1, total_2, rtol=1e-15)))"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "cite2c": {
   "citations": {
    "6232408/PF4E6VIV": {
     "DOI": "10.1007/JHEP11(2013)193",
     "URL": "https://doi.org/10.1007/JHEP11(2013)193",
     "abstract": "We present constraints on decaying and annihilating dark matter (DM) in the 4 keV to 10 GeV mass range, using published results from the satellites HEAO-1, INTEGRAL, COMPTEL, EGRET, and the Fermi Gamma-ray Space Telescope. We derive analytic expressions for the gamma-ray spectra from various DM decay modes, and find lifetime constraints in the range 1024 − 1028 sec, depending on the DM mass and decay mode. We map these constraints onto the parameter space for a variety of models, including a hidden photino that is part of a kinetically mixed hidden sector, a gravitino with R-parity violating decays, a sterile neutrino, DM with a dipole moment, and a dark pion. The indirect constraints on sterile-neutrino and hidden-photino DM are found to be more powerful than other experimental or astrophysical probes in some parts of parameter space. While our focus is on decaying DM, we also present constraints on DM annihilation to electron-positron pairs. We find that if the annihilation is p-wave suppressed, the galactic diffuse constraints are, depending on the DM mass and velocity at recombination, more powerful than the constraints from the Cosmic Microwave Background.",
     "accessed": {
      "day": 5,
      "month": 4,
      "year": 2019
     },
     "author": [
      {
       "family": "Essig",
       "given": "Rouven"
      },
      {
       "family": "Kuflik",
       "given": "Eric"
      },
      {
       "family": "McDermott",
       "given": "Samuel D."
      },
      {
       "family": "Volansky",
       "given": "Tomer"
      },
      {
       "family": "Zurek",
       "given": "Kathryn M."
      }
     ],
     "container-title": "Journal of High Energy Physics",
     "container-title-short": "J. High Energ. Phys.",
     "id": "6232408/PF4E6VIV",
     "issue": "11",
     "issued": {
      "day": 26,
      "month": 11,
      "year": 2013
     },
     "journalAbbreviation": "J. High Energ. Phys.",
     "language": "en",
     "page": "193",
     "page-first": "193",
     "title": "Constraining light dark matter with diffuse X-ray and gamma-ray observations",
     "type": "article-journal",
     "volume": "2013"
    }
   }
  },
  "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.12.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
