{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding a Variable $f_\\text{sed}$\n",
    "\n",
    "In this tutorial you will:\n",
    "\n",
    "1. Learn how to use variable $f_\\text{sed}$ in your calculation\n",
    "2. Compare outputs of constant and variable $f_\\text{sed}$ \n",
    "\n",
    "Citation for this methodology: https://ui.adsabs.harvard.edu/abs/2021arXiv211005903R/abstract \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "from bokeh.io import output_notebook \n",
    "from bokeh.plotting import show, figure\n",
    "from bokeh.palettes import Colorblind\n",
    "output_notebook()\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import astropy.units as u\n",
    "import time\n",
    "\n",
    "#here is pyeddy\n",
    "import virga.justdoit as jdi\n",
    "import virga.justplotit as jpi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us again load in a Hot Jupiter PT profile and run virga for constant $f_\\text{sed}=1$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mieff_directory = \"/data/virga/\"\n",
    "metallicity = 1 #atmospheric metallicity relative to Solar\n",
    "mean_molecular_weight = 2.2 # atmospheric mean molecular weight\n",
    "\n",
    "#set the run \n",
    "a = jdi.Atmosphere(['MnS','Cr','MgSiO3','Fe'],\n",
    "                  fsed=1,mh=metallicity,\n",
    "                 mmw = mean_molecular_weight)\n",
    "\n",
    "#set the planet gravity\n",
    "a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "\n",
    "#Get preset pt profile for testing\n",
    "a.ptk(df = jdi.hot_jupiter())\n",
    "\n",
    "#get full dictionary output \n",
    "all_out = jdi.compute(a, as_dict=True, directory=mieff_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variable $f_\\text{sed}$ expression\n",
    "\n",
    "We have derived an altitude-dependent expression for fsed, given by\n",
    "\n",
    "$f_\\text{sed}(z) = \\alpha\\exp\\left(\\frac{z-z_*}{6\\beta H_0}\\right) + \\epsilon$\n",
    "\n",
    "where $\\alpha$ and $\\beta$ are free parameters, $z_*$ is the altitude at which $f_\\text{sed}(z_*)=\\alpha+\\epsilon$ (default=top of atmosphere), $H_0$ is the scale-height at the effective temperature of the planet and $\\epsilon$ is the minimum value of $f_\\text{sed}$ we wish to consider (default=1e-2). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Varying $\\alpha$ controls the value of $f_\\text{sed}$ at $z_*$ (here the top of the atmosphere)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pressure = all_out['pressure']\n",
    "z = all_out['altitude']\n",
    "H = all_out['scale_height']\n",
    "alpha = [1,10,20,100]; beta = [0.5]\n",
    "fig = jpi.plot_fsed(pressure, z, H, alpha, beta)\n",
    "fig.legend.location='top_left'\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Varying $\\beta$ controls the slope of $f_\\text{sed}$, or the rate of change of $f_\\text{sed}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = [10]; beta = [0.25,0.5,1,5]\n",
    "fig = jpi.plot_fsed(pressure, z, H, alpha, beta)\n",
    "fig.legend.location='top_left'\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running code with variable $f_\\text{sed}$\n",
    "\n",
    "To run VIRGA with this variable $f_\\text{sed}$ expression, we must set the parameterization to \"exp\" and define fsed $ =\\alpha$ and b $ =\\beta$ in the Atmosphere class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define alpha and beta\n",
    "alpha = 10 # generally choose alpha to be between 1 and 100\n",
    "beta = 0.5 # generally choose beta to be between 0.1 and 5\n",
    "\n",
    "#set the run \n",
    "a = jdi.Atmosphere(['MnS','Cr','MgSiO3','Fe'],\n",
    "                   mh=metallicity, mmw = mean_molecular_weight,\n",
    "                   fsed=alpha, b=beta, param=\"exp\")\n",
    "\n",
    "#set the planet gravity\n",
    "a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "\n",
    "#Get preset pt profile for testing\n",
    "a.ptk(df = jdi.hot_jupiter())\n",
    "\n",
    "#get full dictionary output \n",
    "all_out_var = jdi.compute(a, as_dict=True, directory=mieff_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualise the fsed used in the virga calculation using jpi.fsed_from_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = jpi.fsed_from_output(all_out_var,['variable fsed'])\n",
    "fig.legend.visible=False\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional flexibility\n",
    "\n",
    "You can also change the default parameters $z_*$, $\\epsilon$ and $H_0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha=10;\n",
    "beta=0.5;\n",
    "eps=1e-1; # this stops fsed being smaller than 1e-1\n",
    "alpha_pressure=1e-4; # z_* will be set to the altitude at 1e-4 bars\n",
    "Teff=2000; # scale-height H_0 will be defined at this temperature\n",
    "\n",
    "# change alpha, beta, epsilon and parameterisation\n",
    "a = jdi.Atmosphere(['MnS','Cr','MgSiO3','Fe'],\n",
    "                   mh=metallicity, mmw = mean_molecular_weight,\n",
    "                   fsed=alpha, b=beta, param='exp',\n",
    "                   eps=eps)\n",
    "\n",
    "a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "\n",
    "# change alpha_pressure and Teff \n",
    "a.ptk(df = jdi.hot_jupiter(), alpha_pressure=alpha_pressure, Teff=Teff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compare constant to variable $f_\\text{sed}$ calculations\n",
    "\n",
    "### 1. Timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#   compare solvers\n",
    "labels = [\"constant\", \"variable\"]\n",
    "param = ['const', 'exp']\n",
    "fsed = [1, 10] # fsed=1 for constant, alpha=10 for variable\n",
    "beta = 0.5 # note b is irrelevant when param='const'\n",
    "\n",
    "output = []\n",
    "for i, p, fs in zip(range(2), param, fsed):\n",
    "    print(labels[i] + \" fsed\")\n",
    "    t1 = time.time()\n",
    "    a = jdi.Atmosphere(['MnS','Cr','MgSiO3','Fe'],\n",
    "                   mh=metallicity, mmw = mean_molecular_weight,\n",
    "                   fsed=fs, b=beta, param=p)\n",
    "    \n",
    "    a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "    a.ptk(df = jdi.hot_jupiter())\n",
    "    \n",
    "    all_out = jdi.compute(a, as_dict=True, directory=mieff_directory)\n",
    "    t2 = time.time() - t1\n",
    "    print(\"time = %.2f seconds\" % t2)\n",
    "    output.append(all_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Introducing an altitude-dependent expression for $f_\\text{sed}$ does not have detrimental effects on the computational efficiency of VIRGA."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.  Mean Mass Mixing Ratio of the Condensates\n",
    "\n",
    "We can plot the mean mass mixing ratio of the condensates for each solver.\n",
    "We use a solid line for constant $f_\\text{sed}$ and a dashed line for the variable $f_\\text{sed}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = jpi.condensate_mmr(output, plot_height=400, plot_width=500,x_range=[1e-9,1e-2])\n",
    "fig.legend.location='bottom_left'\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The comparison plot shows the results for all of the gases for both solvers. For a clearer comparison, you can choose which gas to view."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.condensate_mmr(output, gas='Fe'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Particle Radii \n",
    "\n",
    "We can similarly analyse the particle size distributions produced by the two solvers.\n",
    "Here we plot \n",
    "1. the mean particle radius for each condensate with pressure\n",
    "2. the particle distribution at a given pressure level (1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, dndr = jpi.radii(output,at_pressure=1e-3)\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, you can choose which gas to view."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, dndr = jpi.radii(output,gas=\"MnS\",at_pressure=1e-3)\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Cumulative Optical Depth By Gas "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.opd_by_gas(output))\n",
    "show(jpi.opd_by_gas(output, gas=\"MgSiO3\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Visualize Optical Depth, Asymmetry, and Single Scattering\n",
    "\n",
    "Finally we can compare the cloud optical depth, single scattering and asymmetry parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(output)):\n",
    "    print(labels[i] + \" solver\")\n",
    "    show(jpi.all_optics(output[i]))"
   ]
  },
  {
   "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
}