{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SED evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import constants\n",
    "import pylab as pl\n",
    "import numpy as np\n",
    "from matplotlib import rcParams, cm\n",
    "rcParams['figure.figsize'] = 12, 8\n",
    "\n",
    "# Imports needed for component separation\n",
    "from fgbuster import CMB, Dust, Synchrotron, AnalyticComponent, MixingMatrix"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `AnalyticComponenet`\n",
    "The `component_model` module collects classes for SED evaluation. It has mainly two classes:\n",
    "\n",
    "* `Component`, abstract class which defines the shared API\n",
    "\n",
    "    - Evaluation of the SED as well as its first and second order derivatives with respect to the free parameters\n",
    "    - Keep track of the free parameters\n",
    "    \n",
    "* `AnalyticComponent`\n",
    "    \n",
    "    - Construct a `Component` from an analytic expression\n",
    "\n",
    "We'll now focus on the latter. \n",
    "An `AnalyticComponent` can be constructed from anything that can be parsed by `sympy`, which include all the most common mathematical functions. A variable called `nu` has to be used to expresses the frequency dependence."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construction\n",
    "As an example, consider the SED of a modified back body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "analytic_expr = ('(nu / nu0)**(1 + beta_d) *'\n",
    "                 '(exp(nu0 / temp * h_over_k) -1)'\n",
    "                 '/ (exp(nu / temp * h_over_k) - 1)'\n",
    "                 )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp = AnalyticComponent(analytic_expr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The free parameters are automatically detected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp.params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As in this case, some of them are not really free but we did not want to clutter the expression by hardcoding the value. Their value can be specified at construction time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h_over_k_val = constants.h * 1e9 / constants.k  # Assumes frequencies in GHz\n",
    "nu0_val = 353.\n",
    "\n",
    "comp = AnalyticComponent(analytic_expr,  # Same as before\n",
    "                         h_over_k=h_over_k_val, nu0=nu0_val)  # but now keyword arguments fix some variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and they are no longer free parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp.params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation\n",
    "Efficient evaluators of the analytic expression are created at construction time. The expression can be evaluated both for a single value and for an array of values. (Note that the free parameters are passed in the same order they appear in `comp.params`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nu_scalar = 50\n",
    "beta_scalar = 1.6\n",
    "temp_scalar = 16.\n",
    "\n",
    "comp.eval(nu_scalar, beta_scalar, temp_scalar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nu_vector = np.logspace(1, 3, 20)\n",
    "\n",
    "sed = comp.eval(nu_vector, beta_scalar, temp_scalar)\n",
    "\n",
    "print('Output shape is :', sed.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true
   },
   "outputs": [],
   "source": [
    "_ = pl.loglog(nu_vector, sed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also the free parameters can be arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_vector = np.linspace(1.2, 2.2, 4)\n",
    "temp_vector = np.linspace(15, 27, 4)\n",
    "\n",
    "seds = comp.eval(nu_vector, beta_vector, temp_vector)\n",
    "\n",
    "print('Output shape is :', seds.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colors_map = cm.ScalarMappable()\n",
    "colors_map.set_array(beta_vector)\n",
    "colors_map.autoscale()\n",
    "colors = lambda x: colors_map.to_rgba(x)\n",
    "for sed, beta, temp in zip(seds, beta_vector, temp_vector):\n",
    "    pl.loglog(nu_vector, sed, c=colors(beta), label='Temp: %.0f   Beta: %.1f' % (temp, beta))\n",
    "_ = pl.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "They can be anything that is bradcating-compatible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sedss = comp.eval(nu_vector, beta_vector, temp_vector[:, np.newaxis])\n",
    "\n",
    "print('Output shape is :', sedss.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true
   },
   "outputs": [],
   "source": [
    "linewidths = lambda x: (x - temp_vector.min()) / 4. + 1\n",
    "\n",
    "# Legend kludge\n",
    "for temp in temp_vector:\n",
    "    pl.plot([353.]*2, [1, 1.001], c='k', lw=linewidths(temp),\n",
    "              label='Temp: %.0f' % temp)\n",
    "for beta in beta_vector:\n",
    "    pl.plot([353.]*2, [1, 1.001], c=colors(beta), lw=2,\n",
    "              label='Beta: %.1f' % beta)\n",
    "# Plotting\n",
    "for temp, seds in zip(temp_vector, sedss):\n",
    "    for beta, sed in zip(beta_vector, seds):\n",
    "        pl.loglog(nu_vector, sed,\n",
    "                  c=colors(beta), lw=linewidths(temp))\n",
    "\n",
    "_ = pl.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation is reasonably quick. For example, evaluating a pixel-dependent SED over 5 frequency bands at `nside=1024` takes few seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "comp.eval(nu_vector[::4], np.linspace(1., 2., 12*1024**2), 20.).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Derivatives\n",
    "Another important feature of the `AnalyticComponent`s is that the derivatives with respect to the parameters are computed analytically at construction time and evaluate with respect to all the parameters using"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diffs = comp.diff(nu_vector, 1.6, 20.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for diff, param in zip(diffs, comp.params):\n",
    "    pl.semilogx(nu_vector, diff, label=param)\n",
    "_ = pl.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Predefined Components\n",
    "Popular analytic SEDs are ready to be used. For example, we construct the very same `Component` we have been using above with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dust = Dust(353., units='K_RJ')\n",
    "print(dust.params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.allclose(comp.eval(nu_vector, beta_vector, temp_vector), \n",
    "            dust.eval(nu_vector, beta_vector, temp_vector)\n",
    "           )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As for the components, parameters can be fixed at construction time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dust_fixed_temp = Dust(353., temp=temp_scalar, units='K_RJ')\n",
    "print(dust_fixed_temp.params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.allclose(comp.eval(nu_vector, beta_vector, temp_scalar), \n",
    "            dust_fixed_temp.eval(nu_vector, beta_vector)\n",
    "           )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mixing Matrix\n",
    "The mixing matrix allows to collect components and evaluate them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixing_matrix = MixingMatrix(Dust(353.), CMB(), Synchrotron(70.))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixing_matrix.params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixing_matrix.components"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As the `Component`s, the SED and its derivative can be easily evaluated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mm_val = mixing_matrix.eval(nu_vector, 1.6, 20., -3)\n",
    "\n",
    "print('Output shape is :', mm_val.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for sed, name in zip(mm_val.T, mixing_matrix.components):\n",
    "    pl.loglog(nu_vector, sed, label=name)\n",
    "_ = pl.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mm_diffs = mixing_matrix.diff(nu_vector, 1.6, 20., -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for mm_diff, param in zip(mm_diffs, mixing_matrix.params):\n",
    "    plot = pl.loglog(nu_vector, mm_diff, label=param)\n",
    "    pl.loglog(nu_vector, -mm_diff, ls='--', c=plot[-1].get_color())\n",
    "_ = pl.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that only the column affected by the derivative is reported in the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mm_diffs[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each parameter, the index of the column affected can be retrieved from the `comp_of_dB` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mixing_matrix.comp_of_dB"
   ]
  }
 ],
 "metadata": {
  "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}