{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Approximations for Spherical Harmonics Radiative Transfer in Thermal Emission \n",
    "\n",
    "In [Rooney et al 2023](add-link) we rigorously derive the spherical harmonics method for thermal emission and benchmark the 2-term and 4-term method (SH4) against [Toon et al. 1989](https://ui.adsabs.harvard.edu/abs/1989JGR....9416287T/abstract) and CDISORT. Here, we provide the code to reproduce the analysis that compares Toon89 with the higher fidelity 4-term spherical harmonics method for thermal emission spectroscopy. \n",
    "\n",
    "Note that all comparisons with `CDISORT` are precomputed following Rooney et al's calculations, which used [V1 opacities](https://zenodo.org/record/3759675#.Y_aJROzMI8Y). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import astropy.units as u\n",
    "\n",
    "#picaso\n",
    "from picaso import justdoit as jdi \n",
    "from picaso import justplotit as jpi\n",
    "\n",
    "jpi.output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Setting up Brown Dwarf Comparison\n",
    "\n",
    "Within the PICASO repository there exists a simple benchmark brown dwarf case that we used in the paper to compare the code. We will start by setting that up. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wave_range = [.7,14]\n",
    "opa = jdi.opannection(wave_range=wave_range)#, resample=100)\n",
    "bd = jdi.inputs(calculation='browndwarf')\n",
    "\n",
    "bd.phase_angle(0)\n",
    "grav = 200\n",
    "bd.gravity(gravity=grav , gravity_unit=u.Unit('m/s**2'))\n",
    "bd.surface_reflect(0,opa.wno)\n",
    "\n",
    "#brown dwarf PT and CLD provide a pressure-temperature profile and cloud profile \n",
    "#from a standard brown dwarf case with Teff~1270 K and fsed = 1 courtesy of C. Morley @ UT Austin\n",
    "bd.atmosphere(filename=jdi.brown_dwarf_pt(), delim_whitespace=True)\n",
    "bd.clouds(filename=jdi.brown_dwarf_cld(), delim_whitespace=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the `jdi.approx` function to setup different test cases\n",
    "\n",
    "Similar to the reflected light problem, we can use the `approx` key to setup different methods of computing the thermal radiative transfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = []; labels = []; \n",
    "# PICASO Original Methodology using Toon source function technique \n",
    "dfs += [bd.spectrum(opa, full_output=True, \n",
    "                    calculation='thermal')\n",
    "       ]; labels += [\"PICASO Toon89\"]\n",
    "\n",
    "# 2-term Spherical harmonics\n",
    "bd.approx(rt_method = 'SH', stream=2)\n",
    "two_lin = bd.spectrum(opa, full_output=True, \n",
    "                      calculation='thermal')\n",
    "dfs += [two_lin]; labels += [\"SH2\"]\n",
    "\n",
    "# 4-term Spherical Harmonics\n",
    "bd.approx(rt_method = 'SH', stream=4)\n",
    "four_lin = bd.spectrum(opa, full_output=True, \n",
    "                       calculation='thermal')\n",
    "dfs += [four_lin]; labels += [\"SH4\"]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Simple regridding and plotting, as we normally do"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = []; ys = []; labs = [];\n",
    "for df, i in zip(dfs, range(len(dfs))):\n",
    "    x,y = df['wavenumber'], df['thermal'] #units of erg/cm2/s/cm\n",
    "    xflux,yflux = jdi.mean_regrid(x, y, R=150)\n",
    "    xs += [xflux]\n",
    "    ys += [yflux]\n",
    "    labs += [labels[i]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### External Testing with `CDISROT` for Validation\n",
    "\n",
    "Reproduction of Figure 2b, Rooney et al. \n",
    "\n",
    "In Rooney et al., we tested `PICASO` against a higher order code, `CDISORT`. Here is a code snippet to grab it from the code base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # read in my disort profile\n",
    "disort_output = jdi.os.path.join(jdi.__refdata__, 'base_cases','testing','cdisort_output_1270_cloudy.spec')\n",
    "son = pd.read_csv(disort_output,\n",
    "                    delim_whitespace=True, skiprows=2,header=None,names=['1','2','3'])\n",
    "sonx, sony, flx =  np.array(1e4/son['1']), np.array(son['2']), np.array(son['3'])\n",
    "sonx_,sony_ = jdi.mean_regrid(sonx, sony*1e1, newx=xs[0]) \n",
    "xs += [sonx_]; ys += [sony_]; labs += ['DISORT16']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = jpi.spectrum(xs,ys,legend=labs\n",
    "                  ,plot_width=800,x_range=wave_range,x_axis_type='log')\n",
    "jpi.show(fig)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up Jupiter-like Case\n",
    "\n",
    "Within the PICASO repository there exists a simple benchmark code jupiter-like case that we used in the paper to compare the code. We will start by setting that up.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wave_range = [5,14]\n",
    "opa = jdi.opannection(wave_range=wave_range)#, resample=100)\n",
    "planet = jdi.inputs()\n",
    "\n",
    "planet.phase_angle(0)\n",
    "grav = 25\n",
    "planet.gravity(gravity=grav , gravity_unit=u.Unit('m/s**2'))\n",
    "# bd.surface_reflect(0,opa.wno)\n",
    "planet.star(opa, 5000,0,4.0) #opacity db, pysynphot database, temp, metallicity, logg\n",
    "\n",
    "planet.atmosphere(filename= jdi.jupiter_pt(), delim_whitespace=True)\n",
    "\n",
    "planet.clouds( filename= jdi.jupiter_cld(), delim_whitespace=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the `jdi.approx` function to setup different test cases\n",
    "\n",
    "Similar to the reflected light problem, we can use the `approx` key to setup different methods of computing the thermal radiative transfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfs = []; labels = []; \n",
    "# PICASO\n",
    "dfs += [planet.spectrum(opa, full_output=True, calculation='thermal')\n",
    "                        ]; labels += [\"PICASO\"]\n",
    "\n",
    "# 2-stream\n",
    "planet.approx(rt_method = 'SH', stream=2)\n",
    "two_lin = planet.spectrum(opa, full_output=True, calculation='thermal')\n",
    "dfs += [two_lin]; labels += [\"SH2\"]\n",
    "\n",
    "# 4-stream \n",
    "planet.approx(rt_method = 'SH', stream=4)\n",
    "four_lin = planet.spectrum(opa, full_output=True, calculation='thermal')\n",
    "dfs += [four_lin]; labels += [\"SH4\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = []; ys = []; labs = [];\n",
    "for df, i in zip(dfs, range(len(dfs))):\n",
    "    x,y = df['wavenumber'], df['thermal'] #units of erg/cm2/s/cm\n",
    "    xflux,yflux = jdi.mean_regrid(x, y, R=150)\n",
    "    xs += [xflux]\n",
    "    ys += [yflux]\n",
    "    labs += [labels[i]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### External Testing with `CDISROT` for Validation\n",
    "\n",
    "Reproduction of Figure bb, Rooney et al. \n",
    "\n",
    "In Rooney et al., we tested `PICASO` against a higher order code, `CDISORT`. Here is a code snippet to grab it from the code base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # read in my disort profile\n",
    "disort_output = jdi.os.path.join(jdi.__refdata__, 'base_cases','testing','cdisort_output_jupiter_cloudy.spec')\n",
    "son = pd.read_csv(disort_output,\n",
    "                    delim_whitespace=True, skiprows=2,header=None,names=['1','2','3'])\n",
    "sonx, sony, flx =  np.array(1e4/son['1']), np.array(son['2']), np.array(son['3'])\n",
    "sonx_,sony_ = jdi.mean_regrid(sonx, sony*1e1, newx=xs[0]) \n",
    "xs += [sonx_]; ys += [sony_]; labs += ['DISORT16']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = jpi.spectrum(xs,ys,legend=labs\n",
    "                  ,plot_width=800,x_range=wave_range,x_axis_type='log')\n",
    "jpi.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can directly compare all of them. Interestingly enough, this figure shows better agreement between the methods (compared to the Brown Dwarf case). As we explain in Rooney, this is because the Toon89 methodology is better suited for scattering regimes in the limit of single scattering -> 1 and -> 0. We will explore this further below. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dependence of Radiative Transfer Method on Scattering Parameters\n",
    "\n",
    "As we alluded to above, there is a large accuracy dependence on single scattering and on single scattering and asymmetry. We can run `PICASO` Toon89 methodology, along with SH2 and SH4 and compare with a precomputed cdisort 32-stream calculation. \n",
    "\n",
    "For this we will use a pre-built test function included in `picaso.test` (this will take some time as it is computing many radiative transfer calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import picaso.test as ptest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Toon quadrature\n",
    "Toon_quad = ptest.thermal_sh_test(method=\"toon\",\n",
    "                                  tau=0.2)\n",
    "# SH 2-term\n",
    "SH2 = ptest.thermal_sh_test(method=\"SH\", \n",
    "                            stream=2, tau=0.2)\n",
    "\n",
    "# SH 4-term\n",
    "SH4 = ptest.thermal_sh_test(method='SH',\n",
    "                            stream=4, tau=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load the cdisort pre computed test file that we have in the `PICASO` reference files. Note that since this is precomputed, minor differences might occur based on what it is in the Rooney paper. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_disort32 = jdi.pd.read_csv(jdi.os.path.join(jdi.__refdata__, 'base_cases','testing', \n",
    "                                         'cdisort32str_1270K_tau02.csv'),index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compare_picaso_disort32 = (data_disort32-Toon_quad)/data_disort32*100   \n",
    "compare_SH4_disort32 = (data_disort32-SH4)/data_disort32*100  \n",
    "compare_SH2_disort32 = (data_disort32-SH2)/data_disort32*100  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot heatmap comparing radiative transfer methods\n",
    "\n",
    "Reproduce Figure 6 in Rooney et al. 2023 Part II Thermal. \n",
    "\n",
    "The efficacy of Toon large depends on the strength of the scattering. When single scattering approaches 1 or 0, Toon89 can be a very effective RT method. However, for intermediate values large errors (>10%) are visible when compared to DISORT 32 stream. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "toon89_fig = jpi.rt_heatmap(\n",
    "    compare_picaso_disort32.iloc[:-1,[0,1,3,5,6,7,8,9,10,11,12,13,14]], \n",
    "    cmap_kwargs={'palette':jpi.pals.plasma(11)[::-1], 'low':-5,'high':60},\n",
    "    figure_kwargs={'title':'Toon89'}\n",
    ")\n",
    "\n",
    "sh4_fig = jpi.rt_heatmap(\n",
    "    compare_SH4_disort32.iloc[:-1,[0,1,3,5,6,7,8,9,10,11,12,13,14]], \n",
    "    cmap_kwargs={'palette':jpi.pals.viridis(11), 'low':-6,'high':2},\n",
    "    figure_kwargs={'title':'SH4'}\n",
    ")\n",
    "\n",
    "diff_fig = jpi.rt_heatmap((compare_picaso_disort32.iloc[:-1,[0,1,3,5,6,7,8,9,10,11,12,13,14]] - \n",
    "                    compare_SH4_disort32.iloc[:-1,[0,1,3,5,6,7,8,9,10,11,12,13,14]] \n",
    "                           ), \n",
    "    cmap_kwargs={'palette':jpi.pals.RdBu[11], 'low':-60,'high':60},\n",
    "    figure_kwargs={'title':'Diff between Toon89 and SH4'}\n",
    ")\n",
    "\n",
    "jpi.show(jpi.row([toon89_fig, sh4_fig, diff_fig]))"
   ]
  }
 ],
 "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
}