{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "628a4099",
   "metadata": {},
   "source": [
    "# Post-Process 3D Clouds Input\n",
    "\n",
    "\n",
    "There are many cases were you may want to post-process components from a 3D GCM output. Post-processing is used to compensate for simplicities, or assumptions within a GCM. For example, you may want to post-process clouds on your output spectra, even if your initial GCM run did not include the effects of clouds. \n",
    "\n",
    "In this notebook you will learn: \n",
    "1. How to post- process clouds using simple box model\n",
    "2. How to post- process clouds using `virga`\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 Clouds: 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 Clouds: User Defined Input \n",
    "\n",
    "Very similar to chemistry process, we can create an xarray that has our cloud output. Note that for clouds there is a fourth dimension (wavelength) "
   ]
  },
  {
   "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",
    "wno_grid = np.linspace(1e4/2,1e4/1,10)#cloud properties are defined on a wavenumber grid\n",
    "\n",
    "#create box-band cloud model \n",
    "fake_opd = np.zeros((len(lon), len(lat),len(pres), len(wno_grid))) # create fake data\n",
    "where_lat = np.where(((lat>-50) & (lat<50)))#creating a grey cloud band\n",
    "where_pres = np.where(((pres<0.01) & (pres>.001)))#creating a grey cloud band \n",
    "for il in where_lat[0]:\n",
    "    for ip in where_pres[0]:\n",
    "        fake_opd[:,il,ip,:]=10 #optical depth of 10 (>>1)\n",
    "\n",
    "#make up asymmetry and single scattering properties \n",
    "fake_asymmetry_g0 = 0.8+ np.zeros((len(lon), len(lat),len(pres), len(wno_grid))) \n",
    "fake_ssa_w0 = 0.9+ np.zeros((len(lon), len(lat),len(pres), len(wno_grid))) \n",
    "\n",
    "# put data into a dataset\n",
    "ds_cld= xr.Dataset(\n",
    "    data_vars=dict(\n",
    "        opd=([\"lon\", \"lat\",\"pressure\",\"wno\"], fake_opd,{'units': 'depth per layer'}),\n",
    "        g0=([\"lon\", \"lat\",\"pressure\",\"wno\"], fake_asymmetry_g0,{'units': 'none'}),\n",
    "        w0=([\"lon\", \"lat\",\"pressure\",\"wno\"], fake_ssa_w0,{'units': 'none'}),\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",
    "        wno=([\"wno\"], wno_grid,{'units': 'cm^(-1)'})#required for clouds\n",
    "    ),\n",
    "    attrs=dict(description=\"coords with vectors\"),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32705cc3",
   "metadata": {},
   "source": [
    "### Visualize Cloud Input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a8697cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_cld['opd'].isel(pressure=where_pres[0][0],wno=0).plot(x='lon',y='lat')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "613b208f",
   "metadata": {},
   "source": [
    "### Add clouds to `picaso` bundle\n",
    "\n",
    "As a recap, in the steps below we: \n",
    "1. Initiate a new run  (jdi.inputs)\n",
    "2. Define the geometry of the run  (phase_angle) \n",
    "3. Set the pressure-temperature-lat-lon profile (atmosphere_3d) \n",
    "\n",
    "Note here we are not worrying about chemical abundances (see post process chem tutorial). We will combine both clouds and chemistry steps in the next tutorial on 3D Spectra. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc112db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#first step is identical to what's been done in the past \n",
    "case_3d = jdi.inputs()\n",
    "case_3d.phase_angle(0, num_tangle=4, num_gangle=4)\n",
    "#turning off alerts since we already went through this \n",
    "case_3d.atmosphere_3d(gcm_out, regrid=True, plot=False,verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9d00e75",
   "metadata": {},
   "source": [
    "Now we can add cloud box model, which will be automatically regridded to the inputs based on the `phase_angle` function ran above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b85e75f",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.clouds_3d(ds_cld,iz_plot=where_pres[0][0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e104739",
   "metadata": {},
   "source": [
    "Stylize plots to see reversed pressure axis in log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44b02500",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = jpi.plt.subplots(figsize=(6, 4))\n",
    "case_3d.inputs['clouds']['profile']['opd'].isel(\n",
    "    lat=2,wno=0).plot(\n",
    "    x='lon',y='pressure',ax=ax)\n",
    "ax.set_ylim([3e2,1e-6])\n",
    "ax.set_yscale('log')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cad53a0",
   "metadata": {},
   "source": [
    "### See other regridded variables\n",
    "\n",
    "All the elements of your data bundle have been regridded! Though for now this is not too interesting since we have set the asymmetry to a constant value of 0.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "722afe05",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.inputs['clouds']['profile']['g0'].isel(pressure=10,wno=0).plot(\n",
    "    x='lon',y='lat')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d7c58cc",
   "metadata": {},
   "source": [
    "## Post-Process Clouds: `virga` cloud model\n",
    "\n",
    "We will run this example on a very coarse (5x5) grid to make it faster. Note that similar to the `chemeq_3d` routine, running `virga_3d` will just run `ngangle`x`ntangle` 1d models. **There is no transport of particles in the vertical/horizontal direction.**  \n",
    "\n",
    "\n",
    "**NOTE!!** If you are not familiar with `virga` models it is highly recommended that you start in one-dimension first. [You can review the tutorials here](https://natashabatalha.github.io/virga/notebooks/1_GettingStarted.html). This tutorial will not explain the required `virga` inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be39455b",
   "metadata": {},
   "outputs": [],
   "source": [
    "opacity = jdi.opannection(wave_range=[1,2])\n",
    "gcm_out =  jdi.HJ_pt_3d(as_xarray=True,add_kz=True)\n",
    "case_3d = jdi.inputs()\n",
    "case_3d.phase_angle(0, num_tangle=10, num_gangle=10)\n",
    "case_3d.atmosphere_3d(gcm_out, regrid=True)\n",
    "case_3d.gravity(radius=1,radius_unit=jdi.u.Unit('R_jup'), \n",
    "                mass=1, mass_unit=jdi.u.Unit('M_jup')) #any astropy units available\n",
    "case_3d.star(opacity,5000,0,4.0, radius=1, radius_unit=jdi.u.Unit('R_sun')) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46bee49c",
   "metadata": {},
   "source": [
    "Run Virga! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8351a892",
   "metadata": {},
   "outputs": [],
   "source": [
    "mieff_directory = '/data/virga/'\n",
    "clds = case_3d.virga_3d(['MnS'], mieff_directory,fsed=1,kz_min=1e10,\n",
    "                        verbose=False,n_cpu=3, full_output=True\n",
    "                       )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56dd1be7",
   "metadata": {},
   "source": [
    "Plot optical depth as a function of longitude\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7225cc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = jpi.plt.subplots(figsize=(6, 4))\n",
    "case_3d.inputs['clouds']['profile']['opd'].isel(lat=5,wno=0).plot(\n",
    "    x='lon',y='pressure',ax=ax)\n",
    "ax.set_ylim([3e2,1e-6])\n",
    "ax.set_yscale('log')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24c6d6f8",
   "metadata": {},
   "source": [
    "### Trouble Shooting Negative Kz Values (smooth_kz) \n",
    "\n",
    "GCM runs produce both positive and negative K_z values. A negative K_z represents downdrafts, which cannot be accounted for in `virga`. In fact, `virga` has a default minimum value of 1e5 cm$^2$/s. For the above run, note I used a very large kz_min value of 1e10cm$^2$/s. The smaller the k_z value, the  values create small particles. Let's take at some of the difficulty with smoothing over K_z profiles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6042ecc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcm_out =  jdi.HJ_pt_3d(as_xarray=True,add_kz=True)\n",
    "test = jdi.inputs()\n",
    "test.phase_angle(0, num_tangle=10, num_gangle=10)\n",
    "test.atmosphere_3d(gcm_out, regrid=True, verbose=False, plot=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05ac9657",
   "metadata": {},
   "source": [
    "Explore some profiles in our 3d Kz grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1227ccd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.output_notebook()\n",
    "df = test.inputs['atmosphere']['profile'].isel(lon=0, lat=0\n",
    "                            ).to_pandas(\n",
    "                            ).reset_index(\n",
    "                            ).drop(\n",
    "                            ['lat','lon'],axis=1\n",
    "                            ).sort_values('pressure')\n",
    "\n",
    "fig = jpi.figure(y_axis_type='log',y_range=[1e2,1e-6],x_axis_label='Kz (cm2/s)',\n",
    "                 y_axis_label='pressure (bars)',\n",
    "                 width=300, height=300, x_range=[-1e11, 1e11])\n",
    "fig.line(df['kz']*1e4,df['pressure'])\n",
    "jpi.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1827f3bd",
   "metadata": {},
   "source": [
    "Experiment with splines "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5302ef56",
   "metadata": {},
   "outputs": [],
   "source": [
    "x=np.log10(df['pressure'].values)\n",
    "df['logp']=x\n",
    "y = np.log10(df['kz'].values)\n",
    "x=x[y>0]\n",
    "y = y[y>0]\n",
    "spl = jdi.UnivariateSpline(x, y,ext=3)\n",
    "df['splKz'] = spl(np.log10(df['pressure'].values))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bd81b2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = jpi.figure(y_axis_type='log',y_range=[1e2,1e-6], x_axis_type='log',\n",
    "                 x_axis_label='Kz (cm2/s)',\n",
    "                 y_axis_label='pressure (bars)',\n",
    "                 width=300, height=300, x_range=[1e5, 1e11])\n",
    "fig.line(df['kz']*1e4,df['pressure'])\n",
    "fig.line(10**df['splKz']*1e4, df['pressure'],color='red')\n",
    "fig.circle(10**y*1e4, 10**x,color='red',size=7)\n",
    "jpi.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb4bb69",
   "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
}