{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple visualisation of a CMIP6 dataset with Python Xarray\n",
    "\n",
    "  - this is an adoption of [NordicESMhub example](https://nordicesmhub.github.io/NEGI-Abisko-2019/training/CMIP6_example.html) to use CMIP6 data stored in the DKRZ data pool "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = '/work/ik1017/CMIP6/data/CMIP6/CMIP/NCAR/CESM2/historical/r1i1p1f1/Amon/tas/gn/v20190308/tas_Amon_CESM2_historical_r1i1p1f1_gn_185001-201412.nc'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import python packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import cartopy.crs as ccrs\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import cftime"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open dataset\n",
    "\n",
    "- Use `xarray` python package to analyze netCDF dataset\n",
    "- `open_dataset` allows to get all the metadata without loading data into memory. \n",
    "- with `xarray`, we only load into memory what is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dset = xr.open_dataset(filename, decode_times=True, use_cftime=True)\n",
    "print(dset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get metadata corresponding to near-surface air temperature (tas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dset['tas'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dset.time.values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Select time\n",
    "\n",
    "- Select a specific time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dset['tas'].sel(time=cftime.DatetimeNoLeap(1850, 1, 15, 12, 0, 0, 0, 2, 15)).plot(cmap = 'coolwarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- select the nearest time. Here from 1st April 1950"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dset['tas'].sel(time=cftime.DatetimeNoLeap(1850, 4, 1), method='nearest').plot(cmap='coolwarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customize plot\n",
    "\n",
    "### Set the size of the figure and add coastlines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=[30,13])\n",
    "\n",
    "# Set the projection to use for plotting\n",
    "ax = plt.subplot(1, 1, 1, projection=ccrs.PlateCarree())\n",
    "ax.coastlines()\n",
    "\n",
    "# Pass ax as an argument when plotting. Here we assume data is in the same coordinate reference system than the projection chosen for plotting\n",
    "# isel allows to select by indices instead of the time values\n",
    "dset['tas'].isel(time=0).plot.pcolormesh(ax=ax, cmap='coolwarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Change plotting projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=[10,10])\n",
    "\n",
    "# We're using cartopy and are plotting in Orthographic projection \n",
    "# (see documentation on cartopy)\n",
    "ax = plt.subplot(1, 1, 1, projection=ccrs.Orthographic(0, 90))\n",
    "ax.coastlines()\n",
    "\n",
    "# We need to project our data to the new Orthographic projection and for this we use `transform`.\n",
    "# we set the original data projection in transform (here PlateCarree)\n",
    "dset['tas'].isel(time=0).plot(ax=ax, transform=ccrs.PlateCarree(), cmap='coolwarm')\n",
    "\n",
    "# One way to customize your title\n",
    "plt.title(dset.time.values[0].strftime(\"%B %Y\"), fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose the extent of values\n",
    "- Fix your minimum and maximum values in your plot and\n",
    "- Use extend so values below the minimum and max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=[10,10])\n",
    "\n",
    "ax = plt.subplot(1, 1, 1, projection=ccrs.Orthographic(0, 90))\n",
    "ax.coastlines()\n",
    "\n",
    "# Fix extent\n",
    "minval = 240\n",
    "maxval = 300\n",
    "\n",
    "# pass extent with vmin and vmax parameters\n",
    "dset['tas'].isel(time=0).plot(ax=ax, vmin=minval, vmax=maxval, transform=ccrs.PlateCarree(), cmap='coolwarm')\n",
    "\n",
    "# One way to customize your title\n",
    "plt.title(dset.time.values[0].strftime(\"%B %Y\"), fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiplots\n",
    "\n",
    "### Faceting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "proj_plot = ccrs.Orthographic(0, 90)\n",
    "\n",
    "p = dset['tas'].sel(time = dset.time.dt.year.isin([1850, 2014])).plot(x='lon', y='lat', \n",
    "                                                                 transform=ccrs.PlateCarree(),\n",
    "                                                                 aspect=dset.dims[\"lon\"] / dset.dims[\"lat\"],  # for a sensible figsize\n",
    "                                                                 subplot_kws={\"projection\": proj_plot},\n",
    "                                                                 col='time', col_wrap=6, robust=True, cmap='PiYG')\n",
    "# We have to set the map's options on all four axes\n",
    "for ax,i in zip(p.axes.flat,  dset.time.sel(time = dset.time.dt.year.isin([1850, 2014])).values):\n",
    "    ax.coastlines()\n",
    "    ax.set_title(i.strftime(\"%B %Y\"), fontsize=18)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combine plots with different projections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(1, figsize=[20,10])\n",
    "\n",
    "# Fix extent\n",
    "minval = 240\n",
    "maxval = 300\n",
    "\n",
    "# Plot 1 for Northern Hemisphere subplot argument (nrows, ncols, nplot)\n",
    "# here 1 row, 2 columns and 1st plot\n",
    "ax1 = plt.subplot(1, 2, 1, projection=ccrs.Orthographic(0, 90))\n",
    "\n",
    "# Plot 2 for Southern Hemisphere\n",
    "# 2nd plot \n",
    "ax2 = plt.subplot(1, 2, 2, projection=ccrs.Orthographic(180, -90))\n",
    "\n",
    "tsel = 0\n",
    "for ax,t in zip([ax1, ax2], [\"Northern\", \"Southern\"]):\n",
    "    map = dset['tas'].isel(time=tsel).plot(ax=ax, vmin=minval, vmax=maxval, \n",
    "                                           transform=ccrs.PlateCarree(), \n",
    "                                           cmap='coolwarm', \n",
    "                                           add_colorbar=False)\n",
    "    ax.set_title(t + \" Hemisphere \\n\" , fontsize=15)\n",
    "    ax.coastlines()\n",
    "    ax.gridlines()\n",
    "\n",
    "# Title for both plots\n",
    "fig.suptitle('Near Surface Temperature\\n' + dset.time.values[tsel].strftime(\"%B %Y\"), fontsize=20)\n",
    "\n",
    "\n",
    "cb_ax = fig.add_axes([0.325, 0.05, 0.4, 0.04])\n",
    "\n",
    "cbar = plt.colorbar(map, cax=cb_ax, extend='both', orientation='horizontal', fraction=0.046, pad=0.04)\n",
    "cbar.ax.tick_params(labelsize=25)\n",
    "cbar.ax.set_ylabel('K', fontsize=25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}