{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9eb2dfad",
   "metadata": {},
   "source": [
    "# Contribution Functions\n",
    "\n",
    "Contribution functions are helpful to understand either 1) what molecules are contributing to the overall resultant spectrum, or 2) at what pressure levels you are sensitive to in thermal emission\n",
    "\n",
    "In this tutorial you will learn how to get molecular contribution plots and thermal contribution heatmaps. This will help you in analyzing your spectra. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "632888af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import picaso.justdoit as jdi\n",
    "import picaso.justplotit as jpi\n",
    "jpi.output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a13dc2b1",
   "metadata": {},
   "source": [
    "## Initial Model Setup\n",
    "\n",
    "Let's do a basic run with a thermal emission and transmission spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7eee08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "opa=jdi.opannection()\n",
    "case1 = jdi.inputs()\n",
    "case1.star(opa, 5000,0,4.0,radius=1, radius_unit=jdi.u.Unit('R_sun'))\n",
    "case1.gravity(radius = 1, radius_unit=jdi.u.Unit('R_jup'),\n",
    "           mass = 1, mass_unit=jdi.u.Unit('M_jup'))\n",
    "case1.phase_angle(0)\n",
    "case1.atmosphere(filename = jdi.HJ_pt(), delim_whitespace=True)\n",
    "df = case1.spectrum(opa, calculation='thermal+transmission', full_output=True)\n",
    "full_output=df['full_output']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "141240e7",
   "metadata": {},
   "source": [
    "## Individual Species Contribution \n",
    "\n",
    "What does get contribution output?? \n",
    "\n",
    "`taus_per_layer`\n",
    "- Each dictionary entry is a nlayer x nwave that represents the per layer optical depth for that molecule. \n",
    "\n",
    "`cumsum_taus`\n",
    "- Each dictionary entry is a nlevel x nwave that represents the cumulative summed opacity for that molecule. \n",
    "\n",
    "`tau_p_surface` \n",
    "- Each dictionary entry is a nwave array that represents the pressure level where the cumulative opacity reaches the value specified by the user through `at_tau`.\n",
    "\n",
    "Note: If you do not see a molecule that you have added as input, check to make sure it is propertly formatted (e.g. Sodium must be Na not NA, Titanium Oxide must be TiO not TIO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c042dd8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = jdi.get_contribution(case1, opa, at_tau=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "564f24b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "out.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcbdc2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.show(jpi.molecule_contribution(out, opa, plot_width=700, x_axis_type='log'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a629367",
   "metadata": {},
   "outputs": [],
   "source": [
    "out['taus_per_layer']['H2O'].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b98d91bc",
   "metadata": {},
   "source": [
    "## See Species Contribution w/ \"leave-one-out\" Method\n",
    "\n",
    "Though the tau~1 surface is very informative it is sometimes easier to just look at your spectrum with the opacity contribution of a single molecule removed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f25d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "#atmo\n",
    "w,f,l =[],[],[]\n",
    "for iex in ['CH4','H2O','CO2',None]:\n",
    "    case1.atmosphere(filename = jdi.HJ_pt(),exclude_mol=iex, delim_whitespace=True)\n",
    "    df= case1.spectrum(opa, full_output=True,calculation='transmission') #note the new last key \n",
    "    wno, rprs2  = df['wavenumber'] , df['transit_depth']\n",
    "    wno, rprs2 = jdi.mean_regrid(wno, rprs2, R=150)\n",
    "    w +=[wno]\n",
    "    f+=[rprs2]\n",
    "    if iex==None: \n",
    "        leg='all'\n",
    "    else: \n",
    "        leg = f'No {iex}'\n",
    "    l+=[leg]\n",
    "jpi.show(jpi.spectrum(w,f,legend=l))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d30d46f",
   "metadata": {},
   "source": [
    "## Transmission Contribution Function \n",
    "\n",
    "The transmission contribution function allows you to determine what pressures your model is sensitive to. The transmission contribution function is defined by Eqn. 8 [Molliere et al 2019](https://arxiv.org/pdf/1904.11504.pdf).\n",
    "\n",
    "Note that this will take a little to run as it is effectively re-computing a transmission spectrum with each layer removed. By comparing the original spectrum to a spectrum with the layer removed, we can see how much that layer contributes to the overall spectrum. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "542a95e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax, um, CF_bin = jpi.transmission_contribution(full_output ,R=100, \n",
    "                                       norm=jpi.colors.LogNorm(vmin=0.005, vmax=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acdbbc6f",
   "metadata": {},
   "source": [
    "## Thermal Contribution Function\n",
    "\n",
    "Contribution functions give us an understanding of where pressures flux is being emitted (e.g. [Lothringer et al 2018](https://iopscience.iop.org/article/10.3847/1538-4357/aadd9e#apjaadd9es3-3-1) Figure 12) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37c6a3d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax, CF = jpi.thermal_contribution(df['full_output'], \n",
    "                                       norm=jpi.colors.LogNorm(vmin=1e7, vmax=1e16))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9952765a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.12"
  },
  "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": true,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}