{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "38eb226c",
   "metadata": {},
   "source": [
    "# 3D Spectra\n",
    "\n",
    "From the previous two tutorials you should understand how: \n",
    "\n",
    "1. How to convert your GCM input to `PICASO`'s required `xarray`\n",
    "2. How to post-process output to append to your GCM output\n",
    "\n",
    "Here you will learn: \n",
    "\n",
    "1. How to run a 3D spectrum\n",
    "2. How to analyze the resulting output\n",
    "\n",
    "**Citation: [Adams et al. 2022 \"Spatially Resolved Modeling of Optical Albedos for a Sample of Six Hot Jupiters\" ApJ](https://ui.adsabs.harvard.edu/abs/2022ApJ...926..157A/abstract)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55e310b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from picaso import justdoit as jdi\n",
    "from picaso import justplotit as jpi\n",
    "jpi.output_notebook()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fdeafad",
   "metadata": {},
   "source": [
    "## Run Cloud-Free 3D Spectra\n",
    "\n",
    "By now you should have completed all the 3D inputs notebooks. This will just take you through completing the last step or running 3D spectra and analyzing your output. \n",
    "\n",
    "Note most of the framework here follows the regular PICASO framework. If you need a refresher on the general workflow, please see earlier tutorials."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7067bf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "opacity = jdi.opannection(wave_range=[1.1,1.7])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "395dc3cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcm_out =  jdi.HJ_pt_3d(as_xarray=True)\n",
    "case_3d = jdi.inputs()\n",
    "case_3d.phase_angle(0, num_gangle=10, num_tangle=10)\n",
    "case_3d.atmosphere_3d(gcm_out,regrid=True,plot=False,verbose=False)\n",
    "case_3d.chemeq_3d(n_cpu=5) #parallelize chemistry run\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": "d97c69d5",
   "metadata": {},
   "source": [
    "### Predicting Compute Runtimes for 3D Calculations\n",
    "\n",
    "The speed of the code will be nearly directly proportional to the number of gangles and tangles as on radiative transfer calculation has to be computed for each facet of the \"disco ball\". "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce568e38",
   "metadata": {},
   "outputs": [],
   "source": [
    "out3d = case_3d.spectrum(opacity,calculation='thermal',dimension='3d',full_output=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d1373d8",
   "metadata": {},
   "source": [
    "## Analyze 3D Spectra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146889bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#same old same old\n",
    "wno,fpfs = jdi.mean_regrid(out3d['wavenumber'],out3d['fpfs_thermal'],R=100)\n",
    "jpi.show(jpi.spectrum(wno, fpfs*1e6, plot_width=500,y_axis_type='log'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c2f1fd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.show(jpi.flux_at_top(out3d ,\n",
    "                     pressures = [5,1,0.8],\n",
    "                     ng=0,nt=0,R=100,\n",
    "                     plot_height=500, plot_width=500))#note the addition of ng and nt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf9528a",
   "metadata": {},
   "source": [
    "### 3D Map of Flux\n",
    "\n",
    "We can use the same plot as we used for reflected light. In 3D this plot is quite a bit more interesting since we can see variations across the disk from out input file.\n",
    "\n",
    "Note below I specify thermal. The default is reflected, I just chose this as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f5ae65",
   "metadata": {},
   "outputs": [],
   "source": [
    "jpi.disco(out3d['full_output'] ,wavelength=[1.1, 1.4], calculation='thermal')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb057d2c",
   "metadata": {},
   "source": [
    "### 3D Map of Gas Opacity Sources \n",
    "\n",
    "In addition to disco there is a general map function that can plot anything in `picaso`’s full output dict. Below is an example of the `taugas` where we must specify both a pressure and wavelength to plot the map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e094e574",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Since Taugas is [nlayer, nwave, nlong, nlat] we can specify a list of pressures and a SINGLE wavelength.\n",
    "jpi.map(out3d['full_output'],plot='taugas',pressure=[5,1,0.8],wavelength=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e6e9224",
   "metadata": {},
   "source": [
    "## Run Cloudy 3D Spectra\n",
    "\n",
    "If you have not already, please consult the tutorial on post-processing cloud inputs. Below, we will use `virga` to post-process clouds on our example GCM input. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1a4e93",
   "metadata": {},
   "outputs": [],
   "source": [
    "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')) \n",
    "case_3d.chemeq_3d(n_cpu=5)\n",
    "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": "code",
   "execution_count": null,
   "id": "6255491e",
   "metadata": {},
   "outputs": [],
   "source": [
    "out3d_cld = case_3d.spectrum(opacity,\n",
    "                             calculation='thermal',dimension='3d',full_output=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b02bf920",
   "metadata": {},
   "source": [
    "### 3D $\\tau$~1 Pressure-Level Map\n",
    "\n",
    "Here we can break down a full map of the molecular, cloud and rayleigh opacity. Whereas in map we had to specify a specific pressure and wavelength, here we can specify an optical depth and it will produce a map of what pressure level corresponds to that optical depth. We tend to see photons that originate from tau=1. Therefore, it is interesting to plot a tau=1 map so we can see if we are seeing different pressure levels at different locations across the disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d2114af",
   "metadata": {},
   "outputs": [],
   "source": [
    "#NOTE, we haven't run clouds and are not at short wavelenghts, \n",
    "#so the cloud and rayleigh opacity will not be too interesting\n",
    "jpi.taumap(out3d_cld['full_output'], wavelength=1.4, at_tau=1)"
   ]
  }
 ],
 "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
}