{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c33c085b",
   "metadata": {},
   "source": [
    "# What is post-processing? \n",
    "\n",
    "To run a full radiative transfer model and create a spectrum, at a minimum we need: \n",
    "\n",
    "1. temperature-pressure profiles \n",
    "2. chemical abundance profiles \n",
    "3. basic planetary/stellar properties (radius, mass, etc) \n",
    "\n",
    "This would be considered a basic cloud-free model. We could also increase complexity by adding a cloud model, which would add: \n",
    "\n",
    "1. optical depth of the cloud \n",
    "2. single scattering albedo profile \n",
    "3. asymmetry profile \n",
    "\n",
    "\n",
    "In reality, chemistry, temperature, and clouds are all related to each other. Temperature dictates cloud formation, clouds affect the energy budget, which feeds back to the temperature. Pairing all three models together is often referred to as a \"self-consistent\" model. Self-consistent models (especially in 2 or 3 dimensions) become very computation expensive.\n",
    "\n",
    "Therefore often times, especially in the case of 3D modeling, we don't get back all of these inputs. Yet, we there are many cases where we would still to explore the effect that, for example, clouds would have on the spectrum. This is where post-processing can help. We can take, for example, a starting temperature map, and compute a chemical profile based on that map (or similarly for clouds). \n",
    "\n",
    "In these next tutorials, we will explore post-processing of chemistry, and clouds. We will discuss different regridding methods to speed up calculations. Here is a basic diagram that illustrates the various post-processing steps that lead up to computing a full spectrum. \n",
    "\n",
    "![workflow-phase.png](workflow-phase.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "628a4099",
   "metadata": {},
   "source": [
    "# Post-Process 3D Chemistry Input\n",
    "\n",
    "\n",
    "In this notebook you will learn: \n",
    "1. How to post- process chemistry to an existing xarray\n",
    "\n",
    "\n",
    "You should already know: \n",
    "1. How to format and regrid an `xarray` dataset\n",
    "\n",
    "Next notebook you will learn:\n",
    "1. How to create spectra and phase curves from 3D input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eac7940b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "from picaso import justdoit as jdi\n",
    "from picaso import justplotit as jpi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50335ae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcm_out =  jdi.HJ_pt_3d(as_xarray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fad48ac0",
   "metadata": {},
   "source": [
    "## Post-Process Chemistry: User Options\n",
    "\n",
    "1. User-defined input: this would only be used to explore simplistic cases (e.g. 100% H2O, or 50/50 H2/H2O). It also might be the case that you have abundance models from elsewhere (e.g. 3D model or GCM) and want to add it to your pressure/temperature `xarray` \n",
    "\n",
    "2. Computationally intensive route: Usually GCM output is at a very high spatial resolution -- higher than what is needed to post-process spectra. If you want, you can grab chemistry for every single `lat` x `lon` point in your grid first. Then, regrid after. \n",
    "    - Pro: you would only have to run it once, then you could regrid to different spatial resolutions later. \n",
    "    - Con: computationally expensive (e.g. 128 x 64 pts x ~1 second = 2 hours per model .. though the exact time depends on how many altitude points you have) \n",
    "3. Computationally efficient route: Regrid first, then compute chemistry. \n",
    "    - Pro: would save potentially hundreds of chemistry computations \n",
    "    - Con: if you aren't happy with your initial binning choice, you may have to repeat the calculation a few times to get it right\n",
    "\n",
    "    \n",
    "In this demo, we will do option #1 and #3 so that it can be self-contained and computationally fast, but you should explore what works best for your problem. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e1a2c7",
   "metadata": {},
   "source": [
    "## Post-Process Chemistry: User Defined Input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be953726",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create coords\n",
    "lon = gcm_out.coords.get('lon').values\n",
    "lat = gcm_out.coords.get('lat').values\n",
    "pres = gcm_out.coords.get('pressure').values\n",
    "\n",
    "fake_chem_H2O = np.random.rand(len(lon), len(lat),len(pres))*0.1+0.1 # create fake data\n",
    "fake_chem_H2 = 1-fake_chem_H2O # create data\n",
    "\n",
    "# put data into a dataset\n",
    "ds_chem = xr.Dataset(\n",
    "    data_vars=dict(\n",
    "        H2O=([\"lon\", \"lat\",\"pressure\"], fake_chem_H2O,{'units': 'v/v'}),\n",
    "        H2=([\"lon\", \"lat\",\"pressure\"], fake_chem_H2,{'units': 'v/v'}),\n",
    "    ),\n",
    "    coords=dict(\n",
    "        lon=([\"lon\"], lon,{'units': 'degrees'}),#required\n",
    "        lat=([\"lat\"], lat,{'units': 'degrees'}),#required\n",
    "        pressure=([\"pressure\"], pres,{'units': 'bar'})#required*\n",
    "    ),\n",
    "    attrs=dict(description=\"coords with vectors\"),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32705cc3",
   "metadata": {},
   "source": [
    "### Add chemistry to an existing GCM xarray "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65bebc81",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_gcm = gcm_out.update(ds_chem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a8697cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_gcm['H2O'].isel(pressure=10).plot(x='lon',y='lat')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "613b208f",
   "metadata": {},
   "source": [
    "### Auto-regrid all 3D input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc112db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d = jdi.inputs()\n",
    "case_3d.phase_angle(0, num_tangle=20, num_gangle=20)\n",
    "case_3d.atmosphere_3d(all_gcm, regrid=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9d00e75",
   "metadata": {},
   "source": [
    "Note there is no print warning about missing abundance info since we have added it directly to the input `xarray`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cad53a0",
   "metadata": {},
   "source": [
    "### See regridded abundance input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "722afe05",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.inputs['atmosphere']['profile']['H2O'].isel(pressure=10).plot(\n",
    "    x='lon',y='lat')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d7c58cc",
   "metadata": {},
   "source": [
    "## Post-Process Chemistry: Chemical Equilibrium\n",
    "\n",
    "We will run this example on a very coarse (5x5) grid to make it faster.\n",
    "\n",
    "First step is the same a above but you will notice a warning that we have yet to add chemistry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be39455b",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcm_no_chem =  jdi.HJ_pt_3d(as_xarray=True)\n",
    "case_3d = jdi.inputs()\n",
    "case_3d.phase_angle(0, num_tangle=5, num_gangle=5)\n",
    "case_3d.atmosphere_3d(gcm_no_chem, regrid=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48f71582",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.chemeq_3d(n_cpu=3)#cpu sets parallelization (1=no parallelization)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56dd1be7",
   "metadata": {},
   "source": [
    "### Plot 3D Chemistry \n",
    "\n",
    "Now we can see structure that intuitively makes sense. Near the equator, where temperatures are hotter, there is less CH4. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66541f56",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.inputs['atmosphere']['profile']['CH4'].isel(pressure=50).plot(\n",
    "    x='lon',y='lat')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6edbcb1e",
   "metadata": {},
   "source": [
    "Make adjustments to `x` and `y` axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3dfae3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = jpi.plt.subplots(figsize=(6, 4))\n",
    "case_3d.inputs['atmosphere']['profile']['CH4'].isel(lat=0).plot(x='lon',y='pressure',\n",
    "                                                                ax=ax)\n",
    "ax.set_ylim([3e2,1e-6])\n",
    "ax.set_yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3d00241",
   "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
  },
  "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": 5
}