{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e7531268",
   "metadata": {},
   "source": [
    "# Phase Curves Part 1\n",
    "\n",
    "**Citation: [Robbins-Blanch et al. 2022 \"Cloudy and Cloud-free Thermal Phase Curves with PICASO: Applications to WASP-43b\" ApJ](http://arxiv.org/abs/2204.03545)**\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 compute a phase curve\n",
    "2. How to analyze the resulting output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fff56978",
   "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": "bc3a7121",
   "metadata": {},
   "source": [
    "## Run Thermal Phase Curve w/ 3D Input\n",
    "\n",
    "Computing a phase curve is identical to computing 3D spectra. The only difference is that you will need to **rotate the GCM longitude grid such that the visible portion of the planet changes as the phase changes.** \n",
    "\n",
    "We can experiment with doing this with our original xarray input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d52f69c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "opacity = jdi.opannection(wave_range=[1,1.7])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89de424a",
   "metadata": {},
   "outputs": [],
   "source": [
    "gcm_out =  jdi.HJ_pt_3d(as_xarray=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "381b4276",
   "metadata": {},
   "source": [
    "### Add Chemistry (randomized for purposes of the example)\n",
    "\n",
    "Add chemistry, if not already included in your GCM output. Here, we are using the user-defined input. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d048a86a",
   "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 = jdi.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",
    "all_gcm = gcm_out.update(ds_chem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d470969b",
   "metadata": {},
   "outputs": [],
   "source": [
    "all_gcm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57941ed7",
   "metadata": {},
   "source": [
    "Setup phase curve grid by defining the type of phase curve (thermal or reflected) as well as the phase angle grid. \n",
    "\n",
    "If you have not gone through the previous tutorials on xarray, post processing input, and computing 3D spectra, we highly recommend you look at that documentation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52520d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d = jdi.inputs()\n",
    "n_phases = 4\n",
    "min_phase = 0 \n",
    "max_phase = 2*np.pi\n",
    "phase_grid = np.linspace(min_phase,max_phase,n_phases)#between 0 - 2pi\n",
    "#send params to phase angle routine\n",
    "case_3d.phase_angle(phase_grid=phase_grid, \n",
    "                    num_gangle=6, num_tangle=6,calculation='thermal')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c9f89e2",
   "metadata": {},
   "source": [
    "### Determining the `zero_point` and the  `shift` parameter \n",
    "\n",
    "Starting assumptions: \n",
    "- User input lat=0, lon=0 is the substellar point and phase=0 is the nightside transit \n",
    "\n",
    "If you want to change the zero point use `zero_point`: \n",
    "- options for zero_point include: 'night_transit' (default) or 'secondary_eclipse'\n",
    "\n",
    "For each orbital `phase`, `picaso` will rotate the longitude grid `phase_i`+`shift_i`. For example:\n",
    "- For tidally locked planets `shift`=np.zeros(len(phases)). Meaning, the sub-stellar point never changes. \n",
    "- For non-tidally locked planets `shift` will depend on the eccentricity and orbital period of the planet \n",
    "\n",
    "`shift` must be input as an array of length `n_phase`. In this example, we will only explore tidally locked planets and will input a zero array. \n",
    "\n",
    "\n",
    "Use `plot=True` below to get a better idea of how your GCM grid is being shifted. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3169098c",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.atmosphere_4d(all_gcm, shift = np.zeros(n_phases), zero_point='night_transit',\n",
    "                                     plot=True,verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22e1d2cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.inputs['atmosphere']['profile']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e589b0b",
   "metadata": {},
   "source": [
    "Easy plotting different phases with `xarray` functionality (remember our H2O selection was random, so the plot will not be to enlightening) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb2286a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_3d.inputs['atmosphere']['profile']['temperature'].isel(pressure=52).plot(\n",
    "                                              x='lon', y ='lat',\n",
    "                                              col='phase',col_wrap=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d61e939f",
   "metadata": {},
   "source": [
    "Set gravity and stellar parameters as usual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78199766",
   "metadata": {},
   "outputs": [],
   "source": [
    "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": "1e24442a",
   "metadata": {},
   "source": [
    "All is setup. Proceed with phase curve calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1017c6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "allout = case_3d.phase_curve(opacity, n_cpu = 3,#jdi.cpu_count(),\n",
    "                             full_output=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a33ca4d9",
   "metadata": {},
   "source": [
    "## Analyze Phase Curves\n",
    "\n",
    "All original plotting tools still work by selecting an individual phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07dd17b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#same old same old\n",
    "wno =[];fpfs=[];legend=[]\n",
    "for iphase in allout.keys():\n",
    "    w,f = jdi.mean_regrid(allout[iphase]['wavenumber'],\n",
    "                               allout[iphase]['fpfs_thermal'],R=100)\n",
    "    wno+=[w]\n",
    "    fpfs+=[f*1e6]\n",
    "    legend +=[str(int(iphase*180/np.pi))]\n",
    "jpi.show(jpi.spectrum(wno, fpfs, plot_width=500,legend=legend,\n",
    "                     palette=jpi.pals.viridis(n_phases)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6014ce49",
   "metadata": {},
   "source": [
    "### Phase Curve Plotting Tools: Phase Snaps\n",
    "\n",
    "Phase snaps allows you to show snapshots of what is happening at each phase that was computed. x,y,z axes can be swapped out for flexibility. \n",
    "\n",
    "- Allowable x: `'longitude', 'latitude', 'pressure'`\n",
    "- Allowable y: `'longitude', 'latitude', 'pressure'` \n",
    "- Allowable z: `'temperature','taugas','taucld','tauray','w0','g0','opd'` (*the latter three are cloud properties*)\n",
    "- Allowable collapse: \n",
    "    Collapse tells picaso what to do with the extra axes. For instance, taugas is an array that is `[npressure,nwavelength,nlongitude,nlatitude]`. If we want `pressure x latitude` we have to collapse wavelength and the longitude axis. To do so we could either supply an integer value to select a single dimension. Or we could supply one of the following as `str` input: [np.mean, np.median, np.min, np.max]. If there are multiple axes you want to collapse, supply an ordered list. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85776266",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=jpi.phase_snaps(allout, x='longitude', y='pressure',z='temperature', \n",
    "                y_log=True, z_log=False, collapse='np.mean')#this will collapse the latitude axis by taking a mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd003f00",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=jpi.phase_snaps(allout, x='longitude', y='pressure',z='taugas', palette='PuBu',\n",
    "                y_log=True, z_log=True, collapse=['np.mean','np.mean'])\n",
    "#this will collapse the latitude axis by taking a mean, and the wavelength axis also by taking the mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbfe463a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=jpi.phase_snaps(allout, x='longitude', y='latitude',z='taugas', palette='PuBu',\n",
    "                y_log=False, z_log=True, collapse=[20,'np.mean'])\n",
    "#this will collapse the pressure axis by taking the 20th pressure grid point \n",
    "#and the wavelength axis by taking the mean\n",
    "print('Pressure at', allout[iphase]['full_output']['layer']['pressure'][20,0,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a308dfc",
   "metadata": {},
   "source": [
    "### Phase Curve Plotting Tools: Phase Curves\n",
    "\n",
    "We can use the same `collapse` feature as in the `phase_snaps` tool to collapse the wavelength axis, when plotting phase curves. For example, we can select a single wavelength point at a given resolution, or we can average over all wavelengths. \n",
    "\n",
    "Allowable collapse: \n",
    "- `'np.mean'` or `np.sum`\n",
    "- float or list of float: wavelength in microns (will find the nearest value to this wavelength). Must be in wavenumber range. \n",
    "\n",
    "For `float` option, user must specify resolution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e9fe2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_plot = 'fpfs_thermal'#or thermal \n",
    "collapse = [1.1,1.4,1.6]#micron\n",
    "R=100\n",
    "phases, all_curves, all_ws, fig=jpi.phase_curve(allout, to_plot, collapse=collapse, R=100)\n",
    "jpi.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "207e105b",
   "metadata": {},
   "source": [
    "## Run Cloudy Thermal Phase Curve w/ 3D Input\n",
    "\n",
    "Here we will create a hypothetical cloud path and add it to our fake H2O profile to demonstrate the basics of running a cloudy phase curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff93a3d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create coords\n",
    "lon = all_gcm.coords.get('lon').values\n",
    "lat = all_gcm.coords.get('lat').values\n",
    "pres = all_gcm.coords.get('pressure').values\n",
    "pres_layer = np.sqrt(pres[0:-1]*pres[1:])\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_layer), len(wno_grid))) # create fake data\n",
    "where_lat = np.where(((lat>-60) & (lat<60)))#creating a grey cloud path\n",
    "where_lon = np.where(lon>0)#creating a grey cloud path\n",
    "where_pres = np.where(((pres_layer<0.01) & (pres_layer>1e-4)))#creating a grey cloud band \n",
    "\n",
    "for il in where_lat[0]:\n",
    "    for ip in where_pres[0]:\n",
    "        for ilo in where_lon[0]:\n",
    "            fake_opd[ilo,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_layer), len(wno_grid))) \n",
    "fake_ssa_w0 = 0.9+ np.zeros((len(lon), len(lat),len(pres_layer), len(wno_grid))) \n",
    "\n",
    "# put data into a dataset\n",
    "ds_cld= jdi.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_layer,{'units': 'bar'}),#required\n",
    "        wno=([\"wno\"], wno_grid,{'units': 'cm^(-1)'})#required for clouds\n",
    "    ),\n",
    "    attrs=dict(description=\"coords with vectors\"),\n",
    ")\n",
    "ds_cld['opd'].isel(pressure=30,wno=0).plot(x='lon',y='lat')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7240a999",
   "metadata": {},
   "outputs": [],
   "source": [
    "case_cld = jdi.inputs()\n",
    "#setup geometry of run\n",
    "n_phases = 4\n",
    "min_phase = 0 \n",
    "max_phase = 2*np.pi\n",
    "phase_grid = np.linspace(min_phase,max_phase,n_phases)#between 0 - 2pi\n",
    "#send params to phase angle routine\n",
    "case_cld.phase_angle(phase_grid=phase_grid, \n",
    "                    num_gangle=6, num_tangle=6,calculation='thermal')\n",
    "\n",
    "#setup 4d atmosphere \n",
    "case_cld.atmosphere_4d(all_gcm, shift = np.zeros(n_phases), zero_point='night_transit',\n",
    "                                     plot=True,verbose=False)\n",
    "#no need to input shift here, it will take it from atmosphere_4d\n",
    "#in fact, clouds_4d must always be run AFTER atmosphere_4d\n",
    "case_cld.clouds_4d(ds_cld, iz_plot=30,\n",
    "                                     plot=True,verbose=False)\n",
    "#same old planet and star properties\n",
    "case_cld.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_cld.star(opacity,5000,0,4.0, radius=1, radius_unit=jdi.u.Unit('R_sun')) \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a578b8e6",
   "metadata": {},
   "source": [
    "Run cloudy phase curve!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1a1f1d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "cldout = case_cld.phase_curve(opacity, n_cpu = 3,#jdi.cpu_count(),\n",
    "                             full_output=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34db3934",
   "metadata": {},
   "outputs": [],
   "source": [
    "#same old same old\n",
    "wno =[];fpfs=[];legend=[]\n",
    "for iphase in cldout.keys():\n",
    "    w,f = jdi.mean_regrid(cldout[iphase]['wavenumber'],\n",
    "                               cldout[iphase]['fpfs_thermal'],R=100)\n",
    "    wno+=[w]\n",
    "    fpfs+=[f*1e6]\n",
    "    legend +=[str(int(iphase*180/np.pi))]\n",
    "jpi.show(jpi.spectrum(wno, fpfs, plot_width=500,legend=legend,\n",
    "                     palette=jpi.pals.viridis(n_phases)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4faacd57",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_plot = 'fpfs_thermal'#or thermal \n",
    "collapse = [1.1,1.4,1.6]#micron\n",
    "R=100\n",
    "phases, all_curves, all_ws, fig=jpi.phase_curve(cldout, to_plot, collapse=collapse, R=100)\n",
    "jpi.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670a2c43",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig=jpi.phase_snaps(cldout, x='longitude', y='pressure',z='taucld', palette='PuBu',\n",
    "                y_log=True, z_log=False, collapse=['np.max','np.mean'])\n",
    "#this will collapse the latitude axis by taking a mean, and the wavelength axis also by taking the mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8336058d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = jpi.plt.subplots(figsize=(6, 4))\n",
    "case_cld.inputs['clouds']['profile']['opd'].isel(phase=0,\n",
    "    lat=3,wno=0).plot(\n",
    "    x='lon',y='pressure',ax=ax)\n",
    "ax.set_ylim([3e2,1e-6])\n",
    "ax.set_yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e52840df",
   "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
}