{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# CMIP6 Surface Temperature/Precipitation Analysis (3-hourly data)\n",
    "\n",
    "[O'Gorman & Schneider (2009)](https://www.pnas.org/content/106/35/14773):\n",
    "> Global warming is expected to lead to a large increase in atmospheric water vapor content and to changes in the hydrological cycle, which include an intensification of precipitation extremes. The intensity of precipitation extremes is widely held to increase proportionately to the increase in atmospheric water vapor content.\n",
    "\n",
    "O'Gorman & Schneider (2009) in CMIP3, O'Gorman (2015) in CMIP5 show P-T scaling in GCMs.\n",
    "\n",
    "\n",
    "[Westra et al. (2013a)](https://journals.ametsoc.org/doi/10.1175/JCLI-D-12-00502.1), among others, have confirmed from long‐term records that the median intensity of observed annual maximum daily rainfall is increasing a rate of 5.9% to 7.7% per °C globally averaged near‐surface atmospheric temperature, but with significant variation by latitude. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import intake\n",
    "import xarray as xr\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import xgcm\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = 12, 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client, progress\n",
    "\n",
    "from dask_kubernetes import KubeCluster\n",
    "cluster = KubeCluster(n_workers=20)\n",
    "client = Client(cluster)\n",
    "cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_url = 'https://raw.githubusercontent.com/pangeo-data/pangeo-datastore/master/intake-catalogs/master.yaml'\n",
    "master_cat = intake.Catalog(cat_url)\n",
    "list(master_cat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A3hr_cat = master_cat.cmip6.A3hr.get()\n",
    "list(A3hr_cat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = 'GISS_E2_1_G.historical.r1i1p1f1'\n",
    "model = 'IPSL_CM6A_LR.historical.r1i1p1f1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_pr = A3hr_cat[f'{model}.pr'].to_dask()\n",
    "ds_pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_tas = A3hr_cat[f'{model}.tas'].to_dask()\n",
    "ds_tas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fix time offset\n",
    "\n",
    "These two datasets are offset by 1.5 hours. This makes it impossible to compute joint statistics. We can use xgcm to fix this. Below we interpolate the `tas` values to the `pr` times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_pr = ds_pr.rename({'time': 'time_center'}).reset_coords(drop=True)\n",
    "ds_tas = ds_tas.rename({'time': 'time_right'}).reset_coords(drop=True)\n",
    "ds = xr.merge([ds_pr, ds_tas])\n",
    "grid = xgcm.Grid(ds, periodic=False,\n",
    "                 coords={'T': {'center': 'time_center', 'right': 'time_right'}})\n",
    "tas_interp = grid.interp(ds.tas, 'T', boundary='extend').chunk({'time_center': 900})\n",
    "ds = ds.drop('time_right')\n",
    "ds['tas'] = tas_interp\n",
    "ds = ds.rename({'time_center': 'time'})\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(ncols=2, figsize=(12, 4))\n",
    "ds.tas[0].plot(ax=axs[0])\n",
    "ds.pr[0].plot(ax=axs[1]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_bins = np.arange(270,315,2)\n",
    "pr_bins = np.logspace(-7, -2, 200)\n",
    "\n",
    "t_center = 0.5*(t_bins[:-1] + t_bins[1:])\n",
    "pr_center = 0.5*(pr_bins[:-1] + pr_bins[1:])\n",
    "\n",
    "def wrapped_hist2d(a, b):\n",
    "    h2d, _, _ = np.histogram2d(a.ravel(), b.ravel(),\n",
    "                           bins=[pr_bins, t_bins])\n",
    "    return h2d[None, :, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#region = dict(lat=slice(-20, 20)) # tropics\n",
    "region = dict(lat=slice(30, 50), lon=slice(230, 300)) # USA\n",
    "\n",
    "ds_reg = ds.sel(**region)\n",
    "\n",
    "import dask.array as dsa\n",
    "h2d_full = dsa.map_blocks(wrapped_hist2d, ds_reg.pr.data, ds_reg.tas.data,\n",
    "                          dtype='i8',\n",
    "                          chunks=(1, len(pr_center), len(t_center)))\n",
    "h2d_full"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2d = xr.DataArray(h2d_full.sum(axis=0).compute(),\n",
    "                   dims=['pr_bin', 'tas_bin'],\n",
    "                   coords={'pr_bin': pr_center,\n",
    "                           'tas_bin': t_center})\n",
    "h2d.coords['pr_bin'].attrs.update(ds.pr.attrs)\n",
    "h2d.coords['tas_bin'].attrs.update(ds.tas.attrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.colors import LogNorm\n",
    "h2d.where(h2d>0).plot(yscale='log', norm=LogNorm())\n",
    "plt.title('Joint Probability Distribution');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2d_normed = h2d / h2d.sum(dim='pr_bin')\n",
    "h2d_cum = h2d_normed.cumsum(dim='pr_bin')\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(12, 7))\n",
    "h2d_cum.plot.contourf(yscale='log', levels=np.arange(0,1,0.1), cmap='Blues_r')\n",
    "cs = h2d_cum.plot.contour(yscale='log', levels=[0.7, 0.8, 0.9, 0.99, 0.999], colors='k')\n",
    "ax.clabel(cs)\n",
    "\n",
    "for scale in np.arange(-16, -10, 0.3):\n",
    "    pr_cc = np.exp(0.068*t_bins)*10**scale\n",
    "    plt.plot(t_bins, pr_cc, color='m', linestyle='--', linewidth=1)\n",
    "\n",
    "\n",
    "ax.set_xlim(285, 310)\n",
    "ax.set_ylim(pr_bins[0], pr_bins[-1]);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}