{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "44181c41",
   "metadata": {},
   "source": [
    "[![View notebook](https://img.shields.io/static/v1?label=render%20on&logo=github&color=87ce3e&message=GitHub)](https://github.com/open-atmos/PyPartMC/blob/main/examples/lognorm_ex.ipynb)   \n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/open-atmos/PyPartMC/blob/main/examples/lognorm_ex.ipynb)    \n",
    "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/open-atmos/PyPartMC.git/main?urlpath=lab/tree/examples/lognorm_ex.ipynb)   \n",
    "[![ARM JupyterHub](https://img.shields.io/static/v1?label=launch%20in&logo=jupyter&color=lightblue&message=ARM+JupyterHub)](https://jupyterhub.arm.gov/hub/user-redirect/git-pull?repo=https%3A//github.com/open-atmos/PyPartMC&branch=main&urlPath=) (requires [logging in with ARM account](https://www.arm.gov/capabilities/computing-resources) and directing Jupyter to a notebook within the cloned repo)\n",
    "[![Voila](https://img.shields.io/static/v1?label=Voil%C3%A0&logo=jupyter&color=teal&message=web+app)](https://mybinder.org/v2/gh/open-atmos/PyPartMC/main?urlpath=voila%2Frender%2Fexamples%2Flognorm_ex.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a19b8c92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This file is a part of PyPartMC licensed under the GNU General Public License v3\n",
    "# Copyright (C) 2022 University of Illinois Urbana-Champaign\n",
    "# Authors: https://github.com/open-atmos/PyPartMC-examples/graphs/contributors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7c55b23d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "if 'google.colab' in sys.modules:\n",
    "    !pip --quiet install open-atmos-jupyter-utils\n",
    "    from open_atmos_jupyter_utils import pip_install_on_colab\n",
    "    pip_install_on_colab('PyPartMC', 'PySDM')\n",
    "elif 'JUPYTER_IMAGE' in os.environ and '.arm.gov' in os.environ['JUPYTER_IMAGE']:\n",
    "    !pip --quiet install PyPartMC PySDM open-atmos-jupyter-utils\n",
    "    _pypartmc_path = !pip show PyPartMC | fgrep Location | cut -f2 -d' '\n",
    "    sys.path.extend(_pypartmc_path if _pypartmc_path[0] not in sys.path else [])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4dcc3eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "import ipywidgets as widgets\n",
    "import numpy as np\n",
    "from PySDM.environments import Box\n",
    "from PySDM import Builder\n",
    "from PySDM.backends import CPU\n",
    "from PySDM.initialisation import equilibrate_wet_radii\n",
    "from matplotlib import pyplot\n",
    "from IPython.display import display, clear_output\n",
    "from open_atmos_jupyter_utils import show_plot\n",
    "from PyPartMC import si\n",
    "import PyPartMC as ppmc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9eec9565",
   "metadata": {},
   "outputs": [],
   "source": [
    "sliders = {\n",
    "    'n_part': {\n",
    "        'widget': widgets.IntSlider(min=8, max=128, value=64),\n",
    "        'label': 'Number of Computational Particles [#]'\n",
    "    },\n",
    "    'temp': {\n",
    "        'widget': widgets.FloatSlider(min=250, max=350),\n",
    "        'label': 'Temperature [K]'\n",
    "    },\n",
    "    'RH_percent': {\n",
    "        'widget': widgets.FloatSlider(min=0, max=100.0, value=55, readout_format='.1f'),\n",
    "        'label': 'Relative Humidity [%]'\n",
    "    },\n",
    "    'kappa': {\n",
    "        'widget': widgets.FloatSlider(min=0, max=2, value=1),\n",
    "        'label': 'Kappa []'\n",
    "    },\n",
    "    'mode_1_n_per_cc': {\n",
    "        'widget': widgets.IntSlider(min=0, max=100000, value=50000),\n",
    "        'label': 'Mode 1 Number [#/cc]'\n",
    "    },\n",
    "    'mode_1_gsd': {\n",
    "        'widget': widgets.FloatSlider(min=1.1, max=5, value=1.3),\n",
    "        'label' : 'Mode 1 Geometric Standard Deviation'\n",
    "    },\n",
    "    'mode_1_gm_microns': {\n",
    "        'widget': widgets.FloatSlider(min=0.001, max=10, value=0.9, readout_format='.3f'),\n",
    "        'label': 'Mode 1 Geometric Mean Diameter [microns]'\n",
    "    },\n",
    "    'mode_2_n_per_cc' : {\n",
    "        'widget': widgets.IntSlider(min=0, max=100000, value=80000),\n",
    "        'label': 'Mode 2 Number [#/cc]'\n",
    "    },\n",
    "    'mode_2_gsd': {\n",
    "        'widget': widgets.FloatSlider(min=1.1, max=5, value=2),\n",
    "        'label': 'Mode 2 Geometric Standard Deviation'\n",
    "    },\n",
    "    'mode_2_gm_microns': {\n",
    "        'widget': widgets.FloatSlider(min=0.001, max=10, value=5.8, readout_format='.3f'),\n",
    "        'label':'Mode 2 Geometric Mean Diameter [microns]'\n",
    "    },\n",
    "    'log_base': {\n",
    "        'widget': widgets.RadioButtons(options=['10','e','none'], value='10'),\n",
    "        'label': 'Distribution function logarithm base'\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77772d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dn_ddp(diam, num, geom_mean, geom_stdev):\n",
    "    return ((num / (np.sqrt(2*np.pi)*diam*np.log(geom_stdev))) *\n",
    "        np.exp(-(np.log(diam) - np.log(geom_mean))** 2 / (2*np.log(geom_stdev)** 2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4e793239",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dn_dlogdp(diam, num, geom_mean, geom_stdev):\n",
    "    return np.log(10) * diam * dn_ddp(diam, num, geom_mean, geom_stdev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "54097605",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dn_dlndp(diam, num, geom_mean, geom_stdev):\n",
    "    return diam * dn_ddp(diam, num, geom_mean, geom_stdev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7607e7f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_bases = {\n",
    "    '10': {\n",
    "        'func': dn_dlogdp,\n",
    "        'y_unit': 1/si.cm**3,\n",
    "        'y_label': '$dN/dlogD_p$ [$cm^{-3}$]'\n",
    "    },\n",
    "    'e': {\n",
    "        'func': dn_dlndp,\n",
    "        'y_unit': 1/si.cm**3,\n",
    "        'y_label': '$dN/dlnD_p$ [$cm^{-3}$]'\n",
    "    },\n",
    "    'none': {\n",
    "        'func': dn_ddp,\n",
    "        'y_unit': (1/si.um)*(1/(si.cm**3)),\n",
    "        'y_label': r'$dN/dD_p$ [$\\mu m^{-1} cm^{-3}$]'\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d5b424ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "linestyles = {\n",
    "    'PyPartMC': 'dashed',\n",
    "    'PySDM': 'solid'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9af4df17",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pypartmc(dry_diameters, temp, rel_humid, kappa):\n",
    "    env_state = ppmc.EnvState({\n",
    "        'rel_humidity': rel_humid,\n",
    "        'latitude': 0.,\n",
    "        'longitude': 0.,\n",
    "        'altitude': 0.,\n",
    "        'start_time': 0.,\n",
    "        'start_day': 0\n",
    "    })\n",
    "    \n",
    "    env_state.set_temperature(temp)\n",
    "    \n",
    "    composition = (\n",
    "        {\"H2O\": [1000 * si.kg / si.m**3, 0, 18e-3 * si.kg / si.mol, 0]},\n",
    "        {\"XXX\": [np.nan * si.kg / si.m**3, 0, np.nan * si.kg / si.mol, kappa]}\n",
    "    )\n",
    "    \n",
    "    aero_data = ppmc.AeroData(composition)\n",
    "    \n",
    "    dry_volumes = (np.pi / 6) * dry_diameters**3\n",
    "    aero_particles = [\n",
    "        ppmc.AeroParticle(aero_data, np.array([0, 1])*volume) for volume in dry_volumes\n",
    "    ]\n",
    "\n",
    "    for aero_particle in aero_particles:\n",
    "        ppmc.condense_equilib_particle(env_state, aero_data, aero_particle)\n",
    "\n",
    "    wet_volumes = [np.sum(particle.volumes) for particle in aero_particles]\n",
    "    wet_diameters = ((6 / np.pi) * np.asarray(wet_volumes))**(1/3)\n",
    "\n",
    "    return wet_diameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4d9881b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pysdm(dry_diameters, temp, rel_humid, kappa):\n",
    "    r_dry = dry_diameters / 2\n",
    "    builder = Builder(n_sd=0, backend=CPU())\n",
    "    environment = Box(dt=np.nan, dv=np.nan)\n",
    "    environment.register(builder)\n",
    "    environment['T'] = temp\n",
    "    environment['RH'] = rel_humid\n",
    "    kappa_times_dry_volume = kappa * (np.pi / 6) * dry_diameters**3\n",
    "    return 2 * equilibrate_wet_radii(\n",
    "        r_dry=r_dry,\n",
    "        environment=environment,\n",
    "        kappa_times_dry_volume=kappa_times_dry_volume\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ee131ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = {\n",
    "    'PyPartMC': pypartmc,\n",
    "    'PySDM': pysdm\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "56fd1c71",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8413898068594929b63f25b06825f179",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Number of Computational Particles [#]:'), IntSlider(value=64, max=128, min=8)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6ae915030f7410292141561b880b036",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Temperature [K]:'), FloatSlider(value=250.0, max=350.0, min=250.0)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ccef33aae7484b94bf014631e97a5ce7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Relative Humidity [%]:'), FloatSlider(value=55.0, readout_format='.1f')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "210dc96de734412e969291856705b875",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Kappa []:'), FloatSlider(value=1.0, max=2.0)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b367f53a27d14a5fa96b7c44f342877f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Mode 1 Number [#/cc]:'), IntSlider(value=50000, max=100000)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "52c9258492e24a4581b6621f148818c5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Mode 1 Geometric Standard Deviation:'), FloatSlider(value=1.3, max=5.0, min=1.1)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6339251b9dd4d819c1d00fc965e6a64",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Mode 1 Geometric Mean Diameter [microns]:'), FloatSlider(value=0.9, max=10.0, min=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "074657939eee416eaad66e74a0db2041",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Mode 2 Number [#/cc]:'), IntSlider(value=80000, max=100000)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0774b29ac6384f1eb7d9b4ba4e519a88",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Mode 2 Geometric Standard Deviation:'), FloatSlider(value=2.0, max=5.0, min=1.1)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d31861643fad408cbdc92b9ac5a4f1ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Mode 2 Geometric Mean Diameter [microns]:'), FloatSlider(value=5.8, max=10.0, min=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "df434a9a1721464b95c35ce8ee79a0cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(Label(value='Distribution function logarithm base:'), RadioButtons(options=('10', 'e', 'none'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa23ee7ebb8c4defbe29892795ed986a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Button(description='Calculate', style=ButtonStyle())"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b709cc132b6f4d60af07434521f4abcc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def distribution():\n",
    "    Mode = namedtuple(\"Mode\", (\"norm_factor\", \"geom_mean\", \"geom_stdev\"))\n",
    "    modes = (\n",
    "        Mode(\n",
    "            norm_factor=sliders['mode_1_n_per_cc']['widget'].value/si.cm**3,\n",
    "            geom_mean=sliders['mode_1_gm_microns']['widget'].value*si.um,\n",
    "            geom_stdev=sliders['mode_1_gsd']['widget'].value\n",
    "        ),\n",
    "        Mode(\n",
    "            norm_factor=sliders['mode_2_n_per_cc']['widget'].value/si.cm**3,\n",
    "            geom_mean=sliders['mode_2_gm_microns']['widget'].value*si.um,\n",
    "            geom_stdev=sliders['mode_2_gsd']['widget'].value\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    with output:\n",
    "        clear_output(wait=True)\n",
    "        \n",
    "        fig = pyplot.figure()\n",
    "        \n",
    "        log_base = log_bases[sliders['log_base']['widget'].value]\n",
    "        \n",
    "        fig.add_subplot(xscale='log')\n",
    "        x_unit = si.um\n",
    "        y_unit = log_base['y_unit']\n",
    "        \n",
    "        dry_diameters = np.logspace(-0.5, 1.5, sliders['n_part']['widget'].value) * si.um\n",
    "        \n",
    "        dist_func = log_base['func']\n",
    "        \n",
    "        y_sum = np.zeros_like(dry_diameters)\n",
    "        for mode in modes:\n",
    "            y_sum += dist_func(dry_diameters, *mode)/y_unit\n",
    "            \n",
    "        pyplot.plot(dry_diameters/x_unit, y_sum, label=\"(dry)\", linewidth=3)\n",
    "        \n",
    "        for model, func in models.items():\n",
    "        \n",
    "            wet_diameters = func(\n",
    "                dry_diameters,\n",
    "                temp = sliders['temp']['widget'].value,\n",
    "                rel_humid = sliders['RH_percent']['widget'].value / 100,\n",
    "                kappa = sliders['kappa']['widget'].value\n",
    "            )\n",
    "            pyplot.plot(wet_diameters/x_unit, y_sum, label=f\"(wet) Model={model}\",\n",
    "                linestyle=linestyles[model], marker='.')\n",
    "        \n",
    "        pyplot.xlabel(r'Diameter, $D_p$ [$\\mu m$]')\n",
    "        pyplot.ylabel(log_base['y_label'])\n",
    "        pyplot.grid()\n",
    "        pyplot.legend()\n",
    "        show_plot(\"spectrum.pdf\")\n",
    "\n",
    "def on_button_clicked(_):\n",
    "    distribution()\n",
    "\n",
    "for slider in sliders.values():\n",
    "    hbox = widgets.HBox([widgets.Label(value=slider['label']+':'), slider['widget']])\n",
    "    display(hbox)\n",
    "    \n",
    "button = widgets.Button(description='Calculate')\n",
    "output = widgets.Output()\n",
    "display(button, output)\n",
    "button.on_click(on_button_clicked)\n",
    "button.click()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b68af7ae",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}