{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Options for K$_z$\n",
    "\n",
    "In this tutorial you will learn: \n",
    "\n",
    "1. The hierarchy of K$_z$ assumptions \n",
    "\n",
    "You should already be familiar with: \n",
    "\n",
    "1. How to pick condensates to run in a cloud model \n",
    "2. What is $f_{sed}$ and what number is right for my model? \n",
    "3. What are the chemical limitations in the code \n",
    "4. How to compute initial Mie scattering grid\n",
    "5. Generally how to run the code \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.io import output_notebook \n",
    "from bokeh.plotting import show, figure\n",
    "output_notebook()\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import astropy.units as u\n",
    "\n",
    "#cloud code\n",
    "import virga.justdoit as jdi\n",
    "import virga.justplotit as jpi\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initial parameters we will keep fixed for the tutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mieff_directory = '/data/virga/'\n",
    "fsed = 0.3\n",
    "metallicity = 1 #atmospheric metallicity relative to Solar\n",
    "mean_molecular_weight = 2.2 # atmospheric mean molecular weight"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Toy Model: Supply Constant K$_z$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum_planet = jdi.Atmosphere(['KCl', 'MgSiO3'],fsed=fsed,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "sum_planet.gravity(gravity=100.00, gravity_unit=u.Unit('m/(s**2)'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constant_kz = 1e10\n",
    "df = jdi.brown_dwarf()\n",
    "sum_planet.ptk(df = df.loc[:,['pressure','temperature']],constant_kz = constant_kz) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_constant = sum_planet.compute(as_dict=True, \n",
    "                               directory=mieff_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Informed: Compute K$_z$ from convective heat flux \n",
    "\n",
    "The convective heat flux is supplied through an external radiative transfer model which partitions the transport of interior heat between radiative and convective fluxes. As an example, we will upload one from the Sonora Grid. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = jdi.brown_dwarf()\n",
    "sum_planet = jdi.Atmosphere(['MgSiO3','KCl'],fsed=fsed,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "sum_planet.gravity(gravity=100.00, gravity_unit=u.Unit('m/(s**2)'))\n",
    "sum_planet.ptk(df=df)\n",
    "out_chf = sum_planet.compute(as_dict=True, \n",
    "                               directory=mieff_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Turn on correction for latent heat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum_planet = jdi.Atmosphere(['MgSiO3','KCl'],fsed=fsed,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "sum_planet.gravity(gravity=100.00, gravity_unit=u.Unit('m/(s**2)'))\n",
    "sum_planet.ptk(df = df, latent_heat=True)\n",
    "out_latent = sum_planet.compute(as_dict=True, \n",
    "                               directory=mieff_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Turn on correction for latent heat & convective overshoot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sum_planet = jdi.Atmosphere(['MgSiO3','KCl'],fsed=fsed,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "sum_planet.gravity(gravity=100.00, gravity_unit=u.Unit('m/(s**2)'))\n",
    "sum_planet.ptk(df = df, latent_heat=True, convective_overshoot=1/3.)\n",
    "out_latent_cos = sum_planet.compute(as_dict=True, \n",
    "                               directory=mieff_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze different $K_z$ formalism via convective heat flux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "legend = ['Constant','CHF-Only','CHF+Latent','CFH+Latent+COS']\n",
    "fig =jpi.pressure_fig(plot_height=250,x_axis_label='Kz cm2/s',x_axis_type='log',x_range=[1e8,5e11])\n",
    "for i,idf in enumerate([out_constant, out_chf, out_latent, out_latent_cos ]):\n",
    "    fig.line(idf['kz'],idf['pressure'],line_width=3,color=jpi.colpals.Colorblind8[i],legend_label=legend[i])\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.all_optics_1d([out_constant, out_chf, out_latent, out_latent_cos ], wave_range=[0.3,1],\n",
    "                       legend = legend))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More informed: Input K$_z$ from external model \n",
    "\n",
    "(e.g. Global Circulation Model) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = jdi.hot_jupiter()\n",
    "\n",
    "sum_planet = jdi.Atmosphere(['MgSiO3','KCl'],fsed=1,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "sum_planet.gravity(gravity=25.00, gravity_unit=u.Unit('m/(s**2)'))\n",
    "sum_planet.ptk(df=df)\n",
    "hj_real = sum_planet.compute(directory=mieff_directory,as_dict=True)\n",
    "sum_planet.ptk(df=df.loc[:,['pressure','temperature']],constant_kz=1e10)\n",
    "hj_constant = sum_planet.compute(directory=mieff_directory,as_dict=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze different $K_z$ formalism via external input\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#get full dictionary output \n",
    "show(jpi.all_optics_1d([hj_real,hj_constant],[0.3,1] , legend=['Real','Constant']))"
   ]
  },
  {
   "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
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}