{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Correlated-K Tables vs. Monochromatic Opacities\n",
    "\n",
    "Throughout the tutorials, we have always used monochromatic opacities. If you are interested in switching to correlated-K tables versus using the currently provided opacities on Zenodo, that is possible. However, we currently are only supporting our pre-mixed correlated-k-tables that you can find  on [Zenodo]() and details are included in [Marley et al. 2020](https://ui.adsabs.harvard.edu/abs/2021ApJ...920...85M/abstract).  \n",
    "\n",
    "Before completing this notebook you will have to: \n",
    "\n",
    "1. [Download at least one or multiple of the k-table folder](https://zenodo.org/record/5590989#.Yzy1qezMJb9)\n",
    "    File should be of the format `sonora_2020_feh+XXX_co_YYY`, where XXX defines the Fe/H and YYY describes the C/O ratio. **Inside that directory there should be at least two files: ascii_data, and full_abunds**\n",
    "\n",
    "\n",
    "2. [Download Sonora PT profiles](https://zenodo.org/record/1309035#.Y1MHYezMLvW) (if this is unfamiliar please see [Brown Dwarf Tutorial](https://natashabatalha.github.io/picaso/notebooks/6_BrownDwarfs.html))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#for number #1 above point to the directory \n",
    "ck_db = '/data/kcoeff_2020_v3/sonora_2020_feh+000_co_100.data.196/'\n",
    "#for number #2 above, point to the directory\n",
    "sonora_profile_db = '/data/sonora_profile/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#picaso\n",
    "from picaso import justdoit as jdi \n",
    "from picaso import justplotit as jpi\n",
    "jpi.output_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#let's do two calculations \n",
    "#first, the regular monochromatic opacities\n",
    "opacity_mono = jdi.opannection() #lets just use all defaults"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Call `opannection` and supply a correlated-k table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that when we call the `opannection` with the `ck=True`, we must also pass it one of the directories from Zenodo. The directory name describes the chemistry that this has been calculated for. The `ck_db` we have defined  is for 1xSolar, Solar C/O. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#second the correlated k table\n",
    "opacity_ck = jdi.opannection(ck=True, \n",
    "                             ck_db=ck_db\n",
    "                            ) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#lets use the same example Jupiter PT profile\n",
    "pt= jdi.pd.read_csv(jdi.jupiter_pt(), delim_whitespace=True,usecols=[0,1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should already be familiar with the following code block. Here we are loading the input class, setting the gravity, stellar parameters, and sonora profile. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "calc = {'ck':opacity_ck,'mono':opacity_mono}\n",
    "cases = {i: jdi.inputs() for i in calc.keys()}\n",
    "\n",
    "#phase angle \n",
    "for i in calc.keys(): cases[i].phase_angle(0) #radians\n",
    "\n",
    "#define gravity\n",
    "for i in calc.keys(): cases[i].gravity(radius=1, \n",
    "                                    radius_unit=jdi.u.Unit('R_jup'),\n",
    "                                    mass=1,\n",
    "                                    mass_unit=jdi.u.Unit('M_jup'))\n",
    "\n",
    "#define star \n",
    "for i in calc.keys(): cases[i].star(calc[i], 5000,0,4.0,\n",
    "                                   radius=1,\n",
    "                                   radius_unit=jdi.u.Unit('R_sun'),\n",
    "                                   semi_major=5,\n",
    "                                   semi_major_unit=jdi.u.Unit('au')) \n",
    "    \n",
    "\n",
    "#just grabbing any Teff so we can use a sonora pressure-temperature profile\n",
    "Teff = 1000\n",
    "for i in calc.keys():cases[i].sonora(sonora_profile_db, Teff, chem='grid')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The differences between the two procedures comes in when we specify the chemistry. Here we introduce a new function called `premix_atmosphere`. This is going to pull the precomputed chemistry that was used to compute the correlated k tables. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cases['ck'].premix_atmosphere(calc['ck'], \n",
    "            df = cases['ck'].inputs['atmosphere']['profile'].loc[:,['pressure','temperature']])\n",
    "\n",
    "#now let's pass that to the monochromatic opacities for consistency\n",
    "cases['mono'].atmosphere(df=cases['ck'].inputs['atmosphere']['profile'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = {}\n",
    "df['ck'] = cases['ck'].spectrum(calc['ck'],full_output=True, calculation='thermal')\n",
    "df['mono'] = cases['mono'].spectrum(calc['mono'],full_output=True, calculation='thermal')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare the optical depth for various layers\n",
    "\n",
    "One big difference you will notice between `taugas` for a monochromatic opacity calculation and a correlated-k calculation, is that `taugas` will have an extra dimension for CK. Those correspond to each of the gauss points. Now, be careful not to confuse the gauss points that we use for the disk integration and these gauss points. They are different! \n",
    "\n",
    "Below is an example of summing the `tau`'s in order to compare with the monochromatic opacities. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = jpi.figure(y_axis_type='log',y_range=[1e-3,1e4],x_range=[1,14],\n",
    "              x_axis_label='Wavelength (um)', \n",
    "              y_axis_label='Optical Depth per Layer')\n",
    "\n",
    "regridx = df['ck']['wavenumber'][df['ck']['wavenumber']<1e4/1]\n",
    "\n",
    "for i in range(df['mono']['full_output']['taugas'].shape[0])[::15]:\n",
    "    x,y = jdi.mean_regrid(df['mono']['wavenumber'],df['mono']['full_output']['taugas'][i,:,0], newx=regridx)\n",
    "    f.line(1e4/x,y, line_dash='solid', line_width=4, color=jpi.Colorblind8[0], legend_label='mono')\n",
    "    plot = 0 \n",
    "    \n",
    "    #we have 8 gauss points so we need to sum these before plotting them \n",
    "    for ig in range(8):\n",
    "        plot +=  calc['ck'].gauss_wts[ig]*df['ck']['full_output']['taugas'][:,:,ig]\n",
    "    f.line(1e4/df['ck']['wavenumber'],plot[i,:],color=jpi.Colorblind8[1]\n",
    "           ,line_width=2, legend_label='CK')\n",
    "jpi.show(f)    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare the final result of the monochromatic and ck generated - spectra\n",
    "\n",
    "If all was perfect and the Correlated-K tables and the picaso opacity file were derived from the same set of cross sections, then these should be identical. Small deviations might be from: 1) different line lists sources, or 2) the resampling factor might be too low. The latter could be resolved by not using the resampling parameter in `opannection`.  \n",
    " \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#flip units just to show an example of unit converstion\n",
    "wno_ck = df['ck']['wavenumber']\n",
    "for i in list(calc.keys()):\n",
    "    compare = 'thermal'\n",
    "    x, y = df[i]['wavenumber'] , df[i][compare]\n",
    "    if i == 'mono': x,y = jdi.mean_regrid(x, y, newx=wno_ck[jdi.np.where((wno_ck>1e4/14)&\n",
    "                                                                      (wno_ck<1e4/1))]) #wavenumber, erg/cm2/s/Hz\n",
    "        \n",
    "    xmicron = 1e4/x\n",
    "\n",
    "    flamy = y*1e-8 #per anstrom instead of per cm\n",
    "    sp = jdi.psyn.ArraySpectrum(xmicron, flamy,\n",
    "                                waveunits='um',\n",
    "                                fluxunits='FLAM')\n",
    "    sp.convert(\"um\")\n",
    "    sp.convert('Fnu') #erg/cm2/s/Hz\n",
    "\n",
    "    x = sp.wave #micron\n",
    "    y= sp.flux #erg/cm2/s/Hz\n",
    "    df[i]['fluxnu'] = y\n",
    "    \n",
    "    df[i]['regridy'] =  y\n",
    "    df[i]['regridx'] = x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = [1e4/df[i]['regridx'] for i in calc.keys()]\n",
    "y = [df[i]['regridy'] for i in calc.keys()]\n",
    "jpi.show(jpi.spectrum(x, \n",
    "                  y,\n",
    "                  legend=list(calc.keys()), plot_width=1000,\n",
    "                  y_axis_type='log',x_range=[1,10],y_range=[1e-8,1e-5]\n",
    "                 ,x_axis_type='log'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Chemistry\n",
    "\n",
    "In this case we supplied it the same chemistry. However, there may be cases where you want to extract the chemistry from the pre-computed files. This is how you would go about doing so. \n",
    "\n",
    "For the below code snippet, I am pulling out the abundances that are higher than 0.1 ppm. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = jpi.figure(y_range=[50,1e-4], y_axis_type='log',x_axis_type='log',\n",
    "               x_range=[1e-8, 1], y_axis_label='Pressure (bar)', \n",
    "              x_axis_label='Mixing Ratio (v/v)')\n",
    "\n",
    "cols = jpi.pals.Category20[20]\n",
    "ii=0\n",
    "for i in cases['ck'].inputs['atmosphere']['profile'].keys(): \n",
    "    if i in ['pressure','temperature']:\n",
    "        continue\n",
    "    if i in cases['ck'].inputs['atmosphere']['profile'].keys(): \n",
    "        for j ,line in zip(['ck', 'mono'], ['dotted','solid']):\n",
    "            try: \n",
    "                x = cases[j].inputs['atmosphere']['profile'][i]\n",
    "            except:\n",
    "                pass\n",
    "            if x.max()>1e-7:\n",
    "                y = cases[j].inputs['atmosphere']['profile']['pressure']\n",
    "                f.line(x,y, line_width=3,line_dash=line,color=cols[ii])\n",
    "                if j == 'mono':ii+=1\n",
    "jpi.show(f)"
   ]
  },
  {
   "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.11.5"
  },
  "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": 4
}