{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing Transmission Spectroscopy\n",
    "\n",
    "Computing transmission spectroscopy is pretty much the same deal as the thermal emission/reflected light. There are just a couple tweaks we will have to make in order to make sure we have the full set of info. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "#picaso\n",
    "from picaso import justdoit as jdi \n",
    "from picaso import justplotit as jpi\n",
    "\n",
    "#plotting\n",
    "jpi.output_notebook()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use new Hot Jupiter template to guide us through the exercise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opa = jdi.opannection(wave_range=[0.6,5])\n",
    "\n",
    "case1 = jdi.inputs()\n",
    "\n",
    "case1.phase_angle(0) \n",
    "\n",
    "\n",
    "#here we are going to have to specify gravity through R and M since we need it in the Flux calc\n",
    "case1.gravity(mass=1, mass_unit=jdi.u.Unit('M_jup'), \n",
    "              radius=1.2, radius_unit=jdi.u.Unit('R_jup'))\n",
    "\n",
    "#here we are going to have to specify R as well\n",
    "case1.star(opa, 4000,0.0122,4.437,radius=0.7, radius_unit = jdi.u.Unit('R_sun') )\n",
    "\n",
    "#atmo\n",
    "case1.atmosphere(filename = jdi.HJ_pt(), delim_whitespace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference Pressure \n",
    "\n",
    "The reference pressure sets the level of your radius. It is needed to compute the altitude and altitude-dependent gravity. PICASO default is 1 bar but for gas giant calculations its more common to use 10 bar. This is easy to swap. If you are running terrestrial planet models you want your reference pressure to be the highest pressure in your user-specified grid. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case1.approx(p_reference=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Return ``PICASO`` Full Ouput"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df= case1.spectrum(opa, full_output=True,calculation='transmission') #note the new last key \n",
    "\n",
    "wno, rprs2  = df['wavenumber'] , df['transit_depth']\n",
    "wno, rprs2 = jdi.mean_regrid(wno, rprs2, R=150)\n",
    "full_output = df['full_output']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we are getting a few error messages that our example `HJ_cld` had some unrecognized molecules. Looking them over, this looks okay! `x` is not a molecule, and neither is `Kzz`. The code has accurately identified non-molecules in our input."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyzing Transit Spectroscopy Output\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Can use the same plotting functionality as reflected light\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.show(jpi.pt(full_output,plot_width=500))\n",
    "jpi.show(jpi.mixing_ratio(full_output,plot_width=500))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Standard (Rp/Rs)^2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.show(jpi.spectrum(wno,rprs2*1e6,plot_width=500))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See Molecular Contribution of Spectra\n",
    "\n",
    "Sometimes it is helpful to compute your spectra by excluding contribution from certain molecules. This will help you see where they affect the spectra. You can do this by using the `exclude_mol` feature, which excludes the opacity of a certain molecule without changing the calculation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
   "metadata": {},
   "source": [
    "## Adding Grey Cloud\n",
    "\n",
    "If you want to test the observability of your features against varying strengths of cloud optical depth, you can use the `cloud` box model function \n",
    "\n",
    "Remember: \n",
    "\n",
    "- `p` is the maximum cloud pressure in log. Max Pressure = 10^p bars \n",
    "- `dp` is the thickness in log. Minimum Pressure = 10^(p-dp) bars\n",
    "- In order to simulate an optically thick cloud deck, set `opd` arbitrarily large (>1)\n",
    "- In transit, the asymmetry and single scattering albedo are not used and can be set to anything between 0-1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case1.clouds(p = [1], dp=[4], opd=[1], g0=[0], w0=[0])\n",
    "\n",
    "df_gry= case1.spectrum(opa, full_output=True,calculation='transmission') #note the new last key \n",
    "\n",
    "wno, rprs2_gry  = df_gry['wavenumber'] , df_gry['transit_depth']\n",
    "wno, rprs2_gry = jdi.mean_regrid(wno, rprs2_gry, R=150)\n",
    "full_output_gry = df_gry['full_output']\n",
    "\n",
    "jpi.show(jpi.spectrum([wno, wno],[1e6*(rprs2-np.mean(rprs2))\n",
    "                              , (rprs2_gry-np.mean(rprs2_gry))*1e6],\n",
    "                  legend=['Cloud free','Grey cloud'],\n",
    "                  plot_width=500))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Cloud From Model Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case1.clouds(filename = jdi.HJ_cld(), delim_whitespace=True)\n",
    "\n",
    "df_model= case1.spectrum(opa, full_output=True,calculation='transmission') #note the new last key \n",
    "\n",
    "wno, rprs2_model = df_model['wavenumber'] , df_model['transit_depth']\n",
    "wno, rprs2_model = jdi.mean_regrid(wno, rprs2_model, R=150)\n",
    "full_output_model = df_model['full_output']\n",
    "\n",
    "jpi.show(jpi.spectrum([wno, wno,wno],[1e6*(rprs2-np.mean(rprs2))\n",
    "                              , (rprs2_gry-np.mean(rprs2_gry))*1e6\n",
    "                             , (rprs2_model-np.mean(rprs2_model))*1e6],legend=['Cloud Free','Grey Cloud','Virga Cloud'],\n",
    "                  plot_width=500))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding Arbitrary Rayleigh Scattering \n",
    "\n",
    "All models will include Rayleigh scattering that is computed based on the atmospheric composition. H$_2$ gas giants would have Rayleigh scattering from hydrogen, Earth-like composition planets would have Rayleigh scattering from nitrogen. Several hot Jupiter observations have revealed slopes that are enhanced, relative to Rayleigh. \n",
    "\n",
    "To simulate an enhanced scattering slope you can add your own parameterized opacity directly to the cloud routine. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = 4.5 #slope dependence for arbitrary rayleigh\n",
    "ref_wave = 0.3 #reference wave to set at what wavelength tau=ref_tau\n",
    "ref_tau = 5e-2\n",
    "\n",
    "wno_cld = jdi.get_cld_input_grid()\n",
    "wave= 1e4/wno_cld\n",
    "enhanced_rayleigh_opd =  np.concatenate([ref_tau*(wave/ref_wave)**(-alpha) for i in range(case1.nlevel-1)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case1.clouds(df=pd.DataFrame({\n",
    "    'opd':enhanced_rayleigh_opd,\n",
    "    'g0':enhanced_rayleigh_opd*0,\n",
    "    'w0':enhanced_rayleigh_opd*0\n",
    "}))\n",
    "\n",
    "df_ray= case1.spectrum(opa, full_output=True,calculation='transmission') #note the new last key \n",
    "\n",
    "wno, rprs2_ray = df_ray['wavenumber'] , df_ray['transit_depth']\n",
    "wno, rprs2_ray = jdi.mean_regrid(wno, rprs2_ray, R=150)\n",
    "full_output_ray = df_ray['full_output']\n",
    "\n",
    "jpi.show(jpi.spectrum([wno, wno],[1e6*(rprs2-np.mean(rprs2))\n",
    "                             , (rprs2_ray-np.mean(rprs2_ray))*1e6],\n",
    "                  legend=['H2 Rayleigh','Enhanced Rayleigh'],\n",
    "                  plot_width=600,x_axis_type='log'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "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,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_output=df['full_output']\n",
    "fig, ax, um, CF_bin = jpi.transmission_contribution(full_output ,R=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Customize Transmission Contribution plot "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = jpi.plt.subplots(figsize=(10, 10))\n",
    "\n",
    "smap = ax.pcolormesh(um, full_output['layer']['pressure'], CF_bin, norm=jpi.colors.LogNorm(vmax=np.max(CF_bin)\n",
    "                                                                                  , vmin =np.max(CF_bin)/100 ))\n",
    "ax.set_ylim(np.max(full_output['layer']['pressure']),\n",
    "            np.min(full_output['layer']['pressure']))\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylabel('Pressure (bar)')\n",
    "ax.set_xlabel('Wavelength ($\\mu$m)')\n",
    "jpi.plt.colorbar(smap, label='Transmission CF')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 2
}