{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clouds: Young Planet Spectroscopy\n",
    "\n",
    "What you will learn:\n",
    "\n",
    "1. How are clouds expected to shape planet spectra across a range of effective temperatures?\n",
    "2. How does this impede our ability to differentiate formation scenarios\n",
    "\n",
    "What you should know: \n",
    "\n",
    "1. What do formation models predict for the effective temperatures of young planets across different masses?\n",
    "2. Given identical luminosity and age, can formation scenarios and mass be determined?\n",
    "3. How do we dissect spectroscopy of planet atmospheres in order to infer atmospheric physical properties such as abundance and climate profiles?\n",
    "\n",
    "**Additional setup requirements:**\n",
    "\n",
    "This is a **bonus** notebook and will require one last setup step. This notebook relies on the cloud code `virga`, which should already be installed with your `PICASO` installation. However, in order for `virga` to work you need to **download** the cloud optical condensates (about 10 Mb): \n",
    "\n",
    "https://zenodo.org/record/3992294#.YN1CXxNKjlw\n",
    "\n",
    "Once you do that you will have to define the path to the condensates (see this next cell) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(action='ignore')\n",
    "import picaso.justdoit as jdi\n",
    "import picaso.justplotit as jpi\n",
    "\n",
    "import virga \n",
    "\n",
    "jpi.output_notebook()\n",
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "#point to your sonora profile grid that you untared (see above cell #2)\n",
    "sonora_profile_db = '/data/sonora_profile/'\n",
    "#see additional path here!!!\n",
    "mieff_dir = '/data/virga/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wave_range = [1,5]\n",
    "opa = jdi.opannection(wave_range=wave_range)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How do clouds affect an atmospheric spectrum compared to the cloud free cases?\n",
    "\n",
    "In the previous workbook, we computed a sequence of spectra as a function of temperature along the hot start formation scenario. Now we will repeat the exercise while adding in cloud models. \n",
    "\n",
    "We will use the cloud model, `virga` which is baked into `PICASO`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case_study = jdi.evolution_track(mass=8, age='all')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a feasible subset of these. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i_to_compute = case_study['hot'].index[0::15]#take every 15th value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cloud free spectra\n",
    "\n",
    "Let's run PICASO in a loop with the different effective temperatures and gravities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yph = jdi.inputs()\n",
    "#let's keep the star fixed\n",
    "yph.star(opa, 5000,0,4.0,radius=1, radius_unit=jdi.u.Unit('R_sun'))\n",
    "yph.phase_angle(0)\n",
    "\n",
    "#Let's stick the loop in right here!\n",
    "hot_output={} #easy mechanism to save all the output\n",
    "for i in i_to_compute:\n",
    "    Teff = case_study['hot'].loc[i,'Teff']\n",
    "    grav = case_study['hot'].loc[i,'grav_cgs']\n",
    "    yph.gravity(gravity= grav, \n",
    "                gravity_unit=jdi.u.Unit('cm/s**2'))\n",
    "    yph.sonora(sonora_profile_db,  Teff)\n",
    "    hot_case = yph.spectrum(opa,calculation='thermal', full_output=True)\n",
    "    hot_output[f'{Teff}_{grav}'] = hot_case"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cloudy spectra\n",
    "\n",
    "Now that we know how to run cloudy models, let's recompute our sequence with clouds. \n",
    "\n",
    "There are a few more parameters that we have to specify with regards to the cloud code: \n",
    "\n",
    "- metallicity and mean molecular weight (should be consistent with your atmosphere profile) \n",
    "- f$_{sed}$: sedimentation efficiency, which controls the vertical extent of the cloud deck. For more information see: https://natashabatalha.github.io/virga/notebooks/1_GettingStarted.html#How-to-pick-f_{sed}\n",
    "- K$_{zz}$: vertical mixing \n",
    "- Gas condensates: in general, you should think carefully about what gases you want to set as condensable. For this exercise we will just look at a subset four potential condensates for simplicity. See [Gao et al. 2021](#References) for a more in depth discussion of condensates in exoplanet atmospheres."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "yph = jdi.inputs()\n",
    "#let's keep the star fixed\n",
    "yph.star(opa, 5000,0,4.0,radius=1, radius_unit=jdi.u.Unit('R_sun'))\n",
    "yph.phase_angle(0)\n",
    "\n",
    "#Let's stick the loop in right here!\n",
    "cldy_hot_output={} #easy mechanism to save all the output\n",
    "cld_output={}\n",
    "for i in i_to_compute:\n",
    "    Teff = case_study['hot'].loc[i,'Teff']\n",
    "    grav = case_study['hot'].loc[i,'grav_cgs']\n",
    "    yph.gravity(gravity= grav, \n",
    "                gravity_unit=jdi.u.Unit('cm/s**2'))\n",
    "    yph.sonora(sonora_profile_db,  Teff)\n",
    "    p=yph.inputs['atmosphere']['profile']['pressure']\n",
    "    t=yph.inputs['atmosphere']['profile']['temperature']\n",
    "    \n",
    "    #NEW CLOUD STUFF\n",
    "    metallicity = 1 #1xSolar \n",
    "    mean_molecular_weight = 2.2\n",
    "    fsed=1\n",
    "    gas_condensates = ['H2O','MnS','Mg2SiO4','Al2O3']\n",
    "\n",
    "    #for the cloud code we have to supply a kzz value, which describes the degree of mixing \n",
    "    yph.inputs['atmosphere']['profile']['kz'] = [1e9]*len(p)\n",
    "    \n",
    "    cld_output[f'{Teff}_{grav}'] = yph.virga(gas_condensates, mieff_dir, fsed=fsed,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "    \n",
    "    hot_case = yph.spectrum(opa,calculation='thermal', full_output=True)\n",
    "    cldy_hot_output[f'{Teff}_{grav}'] = hot_case"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's plot the sequence!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wno,spec=[],[]\n",
    "fig = jpi.figure(height=500,width=500, y_axis_type='log',\n",
    "                 x_axis_label='Wavelength(um)',y_axis_label='Flux (erg/s/cm2/cm)')\n",
    "delta_fig = jpi.figure(height=500,width=500, #y_axis_type='log',\n",
    "                 x_axis_label='Wavelength(um)',y_axis_label='% Diff Cloud Free - Cloudy')\n",
    "#CLOUD FREE\n",
    "for i,ikey in enumerate(hot_output.keys()):\n",
    "    x,y1 = jdi.mean_regrid(hot_output[ikey]['wavenumber'],\n",
    "                          hot_output[ikey]['thermal'], R=150)\n",
    "    t,g=tuple(ikey.split('_'));g=int(np.log10(float(g))*1000)/1000\n",
    "\n",
    "    a=fig.line(1e4/x,y1,color=jpi.pals.Spectral11[::-1][i],line_width=3,\n",
    "               alpha=0.75, line_dash='dashed',legend_label='cloud free')\n",
    "\n",
    "    #CLOUDYyY\n",
    "    x,y2 = jdi.mean_regrid(cldy_hot_output[ikey]['wavenumber'],\n",
    "                          cldy_hot_output[ikey]['thermal'], R=150)\n",
    "    a=fig.line(1e4/x,y2,color=jpi.pals.Spectral11[::-1][i],line_width=3,\n",
    "               legend_label='cloudy')\n",
    "\n",
    "    delta_fig.line(1e4/x,100*(y1-y2)/y1,color=jpi.pals.Spectral11[::-1][i],line_width=3)\n",
    "\n",
    "fig.legend.location='bottom_right'\n",
    "    \n",
    "jpi.show(jpi.row(fig,delta_fig))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At what temperatures and wavelengths are the contribution from clouds dominant? Where are the clouds not a factor?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is the dominant cloud species for each case? \n",
    "\n",
    "We saved all of the cloud output in the dictionary `cld_output`. You can explore all the output fields here: https://natashabatalha.github.io/virga/notebooks/2_RunningTheCode.html#Exploring-dict-Output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cld_output['2276.0_3271.9'].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`virga` has several different functions for exploring this output. For example you can plot the optical depth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.show(virga.justplotit.opd_by_gas(cld_output['2276.0_3271.9']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are interested in the full sequence of optical depths, let's pull out these optical depth plots for a few of the cases. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figs = []\n",
    "\n",
    "for i,ikey in enumerate(list(hot_output.keys())[::3]):\n",
    "\n",
    "    t,g=tuple(ikey.split('_'));g=int(np.log10(float(g))*1000)/1000\n",
    "    \n",
    "    figs += [virga.justplotit.opd_by_gas(cld_output[ikey], \n",
    "                                         title=f'Teff={t}K,logg={g}',plot_width=300)]\n",
    "    \n",
    "    figs[i].legend.location='bottom_left'\n",
    "\n",
    "\n",
    "jpi.show(jpi.row(figs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Does this sequence agree with what is observed in the spectra? Why or why not?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting the optical contribution from the cloud in relation to the molecular gas\n",
    "\n",
    "There is one last output that we will explore in order to understand how the clouds affect the molecular opacity. `picaso` outputs the per layer opacity of both the gas and the cloud. Therefore, we can create a heat map in order to understand the contribution from both."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hot_case['full_output']['taugas'].shape #a nlayer by nwaves by ngauss (1 for resampled line-by-line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a simple heat map plotting function that will plot the opacity of the molecule and the opacity of the cloud. Step through a few of the cases to understand the interplay between the gas and the cloud opacity. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hot_output.keys() #explore a few of these cases "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "fig, ax = plt.subplots(ncols=2,figsize=(15,5))\n",
    "\n",
    "ikey = '1449.8_9637.7'\n",
    "\n",
    "for it, itau in enumerate(['taugas','taucld']):\n",
    "\n",
    "    tau_bin = []\n",
    "    for i in range(hot_case['full_output'][itau].shape[0]):\n",
    "        x,y = jdi.mean_regrid(cldy_hot_output[ikey]['wavenumber'],\n",
    "                              cldy_hot_output[ikey]['full_output'][itau][i,:,0], R=150)\n",
    "        tau_bin += [[y]]\n",
    "\n",
    "    tau_bin = np.array(np.log10(tau_bin))[:,0,:]\n",
    "    X,Y = np.meshgrid(1e4/x,cldy_hot_output[ikey]['full_output']['layer']['pressure'])\n",
    "    Z = tau_bin\n",
    "    pcm=ax[it].pcolormesh(X, Y, Z)\n",
    "    cbar=fig.colorbar(pcm, ax=ax[it])\n",
    "    pcm.set_clim(-3.0, 3.0)\n",
    "    ax[it].set_title(itau)\n",
    "    ax[it].set_yscale('log')\n",
    "    ax[it].set_ylim([1e2,1e-3])\n",
    "    ax[it].set_ylabel('Pressure(bars)')\n",
    "    ax[it].set_ylabel('Wavelength(um)')\n",
    "    cbar.set_label('log Opacity')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References \n",
    "\n",
    "Gao, Peter, et al. \"Aerosols in Exoplanet Atmospheres.\" Journal of Geophysical Research: Planets 126.4 (2021): e2020JE006655."
   ]
  }
 ],
 "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
}