{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gamma vs Log-Normal Particle Size Distribution\n",
    "\n",
    "Author: Elspeth K.H. Lee (https://github.com/ELeeAstro)\n",
    "\n",
    "In this notebook we compare the two particle size distributions available in `virga`:\n",
    "\n",
    "1. **Log-normal** (default) — parameterised by geometric standard deviation `sig`\n",
    "2. **Gamma** (Christie et al. 2022, Appendix B) — shape parameter `A` is derived from the same `sig` via the trigamma function, so the two distributions have matched variance in log-space\n",
    "\n",
    "We run both on the same Hot Jupiter PT profile and compare particle radii, optical depth, and the dn/dr distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.io import output_notebook\n",
    "from bokeh.plotting import show\n",
    "output_notebook()\n",
    "import numpy as np\n",
    "import astropy.units as u\n",
    "\n",
    "import virga.justdoit as jdi\n",
    "import virga.justplotit as jpi\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the built-in Hot Jupiter PT profile. Both runs use identical atmospheric parameters — only the particle size distribution (`dist`) differs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mieff_dir = '/Users/nbatalh1/Documents/data/virga'\n",
    "\n",
    "metallicity = 1          # atmospheric metallicity relative to Solar\n",
    "mean_molecular_weight = 2.2  # atmospheric mean molecular weight\n",
    "\n",
    "# --- Log-normal run (default) ---\n",
    "a_ln = jdi.Atmosphere(['MnS','Cr','MgSiO3'],\n",
    "                      fsed=1, mh=metallicity, mmw=mean_molecular_weight)\n",
    "a_ln.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "a_ln.ptk(df=jdi.hot_jupiter())\n",
    "out_ln = jdi.compute(a_ln, as_dict=True, directory=mieff_dir)\n",
    "\n",
    "# --- Gamma run (Christie et al. 2022, Appendix B) ---\n",
    "a_gm = jdi.Atmosphere(['MnS','Cr','MgSiO3'],\n",
    "                      fsed=1, mh=metallicity, mmw=mean_molecular_weight,\n",
    "                      dist='gamma')\n",
    "a_gm.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "a_gm.ptk(df=jdi.hot_jupiter())\n",
    "out_gm = jdi.compute(a_gm, as_dict=True, directory=mieff_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"lognormal dist: {a_ln.dist}\")\n",
    "print(f\"gamma dist:     {a_gm.dist}\")\n",
    "print(f\"gamma A:        {a_gm.gamma_A:.5f}\")  # should be ~2.54278 for sig=2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PT Profile with Condensation Curves\n",
    "\n",
    "Same atmosphere for both runs — this is just a sanity check that the PT profile and condensation curves are identical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.pt(out_ln, plot_height=450))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Particle Radii: Log-normal vs Gamma\n",
    "\n",
    "The left panel shows the mean particle radius profile with pressure. The right panel shows the full dn/dr distribution at a selected pressure level.\n",
    "\n",
    "For `sig=2` (the default), Christie et al. (2022) find the effective radii are similar between the two distributions — this is a good first check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, dndr = jpi.radii([out_ln, out_gm], at_pressure=1e-3,\n",
    "                       compare=True, legend=['lognormal', 'gamma'])\n",
    "show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"lognormal rg (top layer):\", out_ln['mean_particle_r'][0])\n",
    "print(\"gamma     rg (top layer):\", out_gm['mean_particle_r'][0])\n",
    "print(\"lognormal reff (top layer):\", out_ln['droplet_eff_r'][0])\n",
    "print(\"gamma     reff (top layer):\", out_gm['droplet_eff_r'][0])\n",
    "print(\"gamma_A:\", out_gm['scalar_inputs']['gamma_A'])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optical Depth, Single Scattering, Asymmetry\n",
    "\n",
    "Compare the full wavelength-resolved optics between the two distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.all_optics(out_ln))\n",
    "show(jpi.all_optics(out_gm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1D Optics (wavelength-averaged)\n",
    "\n",
    "Averaged over 1–2 microns for a quick visual comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.all_optics_1d(out_ln, wave_range=[1,2]))\n",
    "show(jpi.all_optics_1d(out_gm, wave_range=[1,2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cumulative Optical Depth by Gas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.opd_by_gas(out_ln))\n",
    "show(jpi.opd_by_gas(out_gm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mean Mass Mixing Ratio of the Condensates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show(jpi.condensate_mmr(out_ln))\n",
    "show(jpi.condensate_mmr(out_gm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Extended Comparison of Different Distribtuions over Sigma and Fsed \n",
    "\n",
    "The sections above used the default `sig=2`. Christie et al. (2022) show that the two distributions agree well at `sig=2` but diverge at larger values — the gamma distribution peak stays close to `r_w` while the lognormal peak shifts significantly. The sections below explore this systematically."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sweep over sigma\n",
    "\n",
    "We run both distributions for `sig = [1.5, 2.0, 2.5, 3.0]` with all other parameters fixed. This is the key test from Christie et al. (2022): the distributions should agree at `sig=2` and diverge as `sig` increases, with the gamma distribution producing fewer very large particles than the lognormal."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Particle radii and dn/dr at 1e-3 bar for each sigma — lognormal vs gamma overlaid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig_values = [1.1, 1.5, 2.0, 2.5, 3.0]\n",
    "\n",
    "outputs_ln_sig = []\n",
    "outputs_gm_sig = []\n",
    "\n",
    "for sig in sig_values:\n",
    "    # lognormal\n",
    "    a = jdi.Atmosphere(['MnS','Cr','MgSiO3'],\n",
    "                       fsed=1, mh=metallicity, mmw=mean_molecular_weight, sig=sig)\n",
    "    a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "    a.ptk(df=jdi.hot_jupiter())\n",
    "    outputs_ln_sig.append(jdi.compute(a, as_dict=True, directory=mieff_dir))\n",
    "\n",
    "    # gamma\n",
    "    a = jdi.Atmosphere(['MnS','Cr','MgSiO3'],\n",
    "                       fsed=1, mh=metallicity, mmw=mean_molecular_weight, sig=sig,\n",
    "                       dist='gamma')\n",
    "    a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "    a.ptk(df=jdi.hot_jupiter())\n",
    "    outputs_gm_sig.append(jdi.compute(a, as_dict=True, directory=mieff_dir))\n",
    "\n",
    "print(\"Done. sig values:\", sig_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.layouts import column\n",
    "\n",
    "all_figs = []\n",
    "for i, sig in enumerate(sig_values):\n",
    "    fig, _ = jpi.radii([outputs_ln_sig[i], outputs_gm_sig[i]], at_pressure=1e-3,\n",
    "                        compare=True, legend=['lognormal', 'gamma'])\n",
    "    all_figs.append(fig)\n",
    "\n",
    "# Link x/y ranges across rows so all panels share the same scale\n",
    "p1_ref = all_figs[0].children[0]   # mean radius vs pressure\n",
    "p2_ref = all_figs[0].children[1]   # dn/dr distribution\n",
    "for fig in all_figs[1:]:\n",
    "    fig.children[0].x_range = p1_ref.x_range\n",
    "    fig.children[1].x_range = p2_ref.x_range\n",
    "    fig.children[1].y_range = p2_ref.y_range\n",
    "\n",
    "show(column(*all_figs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sweep over fsed\n",
    "\n",
    "We fix `sig=2` and sweep `fsed = [0.1, 0.5, 1.0, 3.0]`. Lower `fsed` means more lofted, larger particles — this tests whether the distribution choice matters more in the thick cloud regime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fsed_values = [0.1, 0.5, 1.0, 3.0]\n",
    "\n",
    "outputs_ln_fsed = []\n",
    "outputs_gm_fsed = []\n",
    "\n",
    "for fsed in fsed_values:\n",
    "    # lognormal\n",
    "    a = jdi.Atmosphere(['MnS','Cr','MgSiO3'],\n",
    "                       fsed=fsed, mh=metallicity, mmw=mean_molecular_weight)\n",
    "    a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "    a.ptk(df=jdi.hot_jupiter())\n",
    "    outputs_ln_fsed.append(jdi.compute(a, as_dict=True, directory=mieff_dir))\n",
    "\n",
    "    # gamma\n",
    "    a = jdi.Atmosphere(['MnS','Cr','MgSiO3'],\n",
    "                       fsed=fsed, mh=metallicity, mmw=mean_molecular_weight,\n",
    "                       dist='gamma')\n",
    "    a.gravity(gravity=7.460, gravity_unit=u.Unit('m/(s**2)'))\n",
    "    a.ptk(df=jdi.hot_jupiter())\n",
    "    outputs_gm_fsed.append(jdi.compute(a, as_dict=True, directory=mieff_dir))\n",
    "\n",
    "print(\"Done. fsed values:\", fsed_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.layouts import column\n",
    "\n",
    "all_figs_r = []\n",
    "all_figs_m = []\n",
    "for i, fsed in enumerate(fsed_values):\n",
    "    fig_r, _ = jpi.radii([outputs_ln_fsed[i], outputs_gm_fsed[i]], at_pressure=1e-3,\n",
    "                          compare=True, legend=['lognormal', 'gamma'])\n",
    "    fig_m = jpi.condensate_mmr([outputs_ln_fsed[i], outputs_gm_fsed[i]])\n",
    "    all_figs_r.append(fig_r)\n",
    "    all_figs_m.append(fig_m)\n",
    "\n",
    "# Link ranges for radii panels so all rows share the same scale\n",
    "p1_ref = all_figs_r[0].children[0]   # mean radius vs pressure\n",
    "p2_ref = all_figs_r[0].children[1]   # dn/dr distribution\n",
    "for fig in all_figs_r[1:]:\n",
    "    fig.children[0].x_range = p1_ref.x_range\n",
    "    fig.children[1].x_range = p2_ref.x_range\n",
    "    fig.children[1].y_range = p2_ref.y_range\n",
    "\n",
    "show(column(*[f for pair in zip(all_figs_r, all_figs_m) for f in pair]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optics comparison\n",
    "\n",
    "For `sig=2` (where distributions agree) and `sig=3` (where they diverge), compare 1D wavelength-averaged optics. Christie et al. (2022) find negligible differences in transmission spectra at `sig=2` but larger differences at wider distributions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.layouts import column\n",
    "\n",
    "all_figs = []\n",
    "for label, out in [('lognormal sig=2', out_ln), ('gamma sig=2', out_gm),\n",
    "                   ('lognormal sig=3', outputs_ln_sig[3]), ('gamma sig=3', outputs_gm_sig[3])]:\n",
    "    print(f\"── {label} ──\")\n",
    "    all_figs.append(jpi.all_optics_1d(out, wave_range=[0.5, 5]))\n",
    "\n",
    "show(column(*all_figs))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pic312",
   "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.12.2"
  },
  "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
}