animation_03_correlation_cmip6_era5.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Correlation between CMIP6 and ERA5\n",
    "\n",
    "\n",
    "This notebooks show the correlation of precipitation between CMIP6 model simulation and observation data from ERA5.\n",
    "\n",
    "* Open CMIP6 and ERA5\n",
    "* Data interpolation\n",
    "* Calculate correlation\n",
    "* Make correlation animation across 35 years\n",
    "\n",
    "The following material uses Coupled Model Intercomparison Project (CMIP6) and ERA5 data. \n",
    "\n",
    "Please visit the NCI Data Catalogue for more information on the [CMIP6](http://dx.doi.org/10.25914/5b98afc88531e) and [ERA5](http://dx.doi.org/10.25914/5f48874388857) collections.\n",
    "\n",
    "---\n",
    "\n",
    "- Authors: NCI Virtual Research Environment Team\n",
    "- Keywords: CMIP6, ERA5, correlation, animation\n",
    "- Create Date: 2020-Jul\n",
    "    \n",
    "---\n",
    "\n",
    "This notebook is licenced under the [Creative Commons Attribution 4.0 International license](https://creativecommons.org/licenses/by/4.0/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load in the relevant Python libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cftime\n",
    "import cartopy.crs as ccrs\n",
    "from dask.distributed import Client\n",
    "import glob\n",
    "from IPython.display import display, HTML\n",
    "from matplotlib import pyplot as plt, animation\n",
    "import numpy as np\n",
    "import os\n",
    "from PIL import Image\n",
    "import xarray as xr\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook was designed for use on NCI's Virtual Desktop Infrastructure (VDI)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Open CMIP6 data\n",
    "\n",
    "Let's take a look at some CMIP6 precipitation data from the [IPSL-CM6A-LR](https://cmc.ipsl.fr/ipsl-climate-models/ipsl-cm6/) climate model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "### location of our CMIP6 example data\n",
    "file = '/g/data/oi10/replicas/CMIP6/CMIP/IPSL/IPSL-CM6A-LR/historical/r1i1p1f1/Amon/pr/gr/v20180803/pr_Amon_IPSL-CM6A-LR_historical_r1i1p1f1_gr_185001-201412.nc'\n",
    "\n",
    "### open this dataset using xarray\n",
    "ds_cmip6 = xr.open_dataset(file)\n",
    "\n",
    "### convert time axis to 360_day calendar -- see https://unidata.github.io/cftime/api.html for more information on using cftime.num2date\n",
    "units = 'months since 1850-1-15 00:00:00'\n",
    "time_360_ref = cftime.num2date(np.arange(0, (2014-1850+1)*12), units, '360_day')\n",
    "ds_cmip6 = ds_cmip6.assign_coords(time=time_360_ref)\n",
    "ds_cmip6"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ERA5 observation data \n",
    "\n",
    "Now let's take a look at some ERA5 single level monthly averaged precipitation data found in project rt52:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### location of our ERA5 example data\n",
    "directory = '/g/data/rt52/era5/single-levels/monthly-averaged/tp/'\n",
    "\n",
    "### define files we want to read in\n",
    "era5_files_all = sorted(list(glob.glob(os.path.join(directory,'*/*.nc'))))\n",
    "era5_files = era5_files_all[:-2]\n",
    "\n",
    "### open a multi file dataset using xarray:\n",
    "ds_tp_mon = xr.open_mfdataset(era5_files,combine='by_coords',parallel=True)\n",
    "\n",
    "ds_tp_mon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### rename ERA5 coordinates so that they are the same as our CMIP6 coordinates\n",
    "ds_tp_mon = ds_tp_mon.rename({'latitude':'lat'}).rename({'longitude':'lon'}).rename({'tp':'pr'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### make sure indexes along dimension 'time' are equal\n",
    "units = 'months since 1979-1-15 00:00:00'\n",
    "time_360_ref = cftime.num2date(np.arange(0, (2020-1979+1)*12), units, '360_day')\n",
    "time_360_ref_new = time_360_ref[0:time_360_ref.size-2] ### ensuring time co-ordines are same size (i.e. both end in October 2020)\n",
    "\n",
    "ds_tp_mon=ds_tp_mon.assign_coords(time=time_360_ref_new)\n",
    "ds_tp_mon=ds_tp_mon.assign_coords(lon=ds_tp_mon.lon+180)\n",
    "\n",
    "ds_tp_mon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data interpolation \n",
    "\n",
    "Now we can remap our ERA5 dataset(**ds_tp_mon**) so that it's the same resolution as our CMIP6 dataset (**ds_cmip6**):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_lon = ds_cmip6.lon\n",
    "new_lat = ds_cmip6.lat\n",
    "\n",
    "ds_tp_mon_remap = ds_tp_mon.interp(lat=new_lat, lon=new_lon)\n",
    "ds_tp_mon_remap"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate correlation\n",
    "\n",
    "Let's use the numpy [corrcoef](https://numpy.org/doc/stable/reference/generated/numpy.corrcoef.html) function to calculate the correlation between our two example datasets. We can use xarray's [apply_ufunc](https://xarray.pydata.org/en/stable/examples/apply_ufunc_vectorize_1d.html) to apply an unvectorised function to xarray objects:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_corrcoef(inarray1, inarray2):\n",
    "    return np.corrcoef(inarray1, inarray2)[0, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Let's use the year 2000 as a test case\n",
    "\n",
    "cmip_test = ds_cmip6['pr'].sel(time=slice('2000','2000')).load()\n",
    "era_test = ds_tp_mon_remap['pr'].sel(time=slice('2000','2000')).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.seterr(divide='ignore', invalid='ignore') ### Ignore \"invalid value encountered in true_divide\" warning messages\n",
    "\n",
    "cc = xr.apply_ufunc(\n",
    "\n",
    "    get_corrcoef,    # First the function\n",
    "\n",
    "    cmip_test, era_test,        # then the input data arrays\n",
    "\n",
    "    input_core_dims=[['time'], ['time']],\n",
    "\n",
    "    vectorize=True\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting\n",
    "\n",
    "Let's use cartopy to plot the correlation of Precipitation for CMIP6 vs ERA5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Choose a year to plot\n",
    "year=2000\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.axes(projection=ccrs.PlateCarree())\n",
    "ax.set_title('Correlation of Pr CMIP6 vs ERA5 '+str(year)+'')\n",
    "ax.coastlines()  ## add in coastlines\n",
    "colorlabel=np.linspace(-1, 1, 9)\n",
    "cc.plot(ax=ax, vmin=-1, vmax=1, cmap='coolwarm', add_colorbar=True, add_labels=True, cbar_kwargs={'orientation':'horizontal', 'pad':0.05, 'shrink':1, 'label':'Correlation', 'ticks':colorlabel, 'spacing':'proportional'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make animation\n",
    "\n",
    "Now we can create a CMIP6 vs ERA5 precipitation correlation animation over our 35 year overlapping period from 1979 to 2014. First let's create a PNG file correlation plot per year: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = []\n",
    "np.seterr(divide='ignore', invalid='ignore') ### Ignore \"invalid value encountered in true_divide\" warning messages\n",
    "\n",
    "### create output directory if it doesn't exist already\n",
    "outdir = './output'\n",
    "if not os.path.exists(outdir):\n",
    "    os.mkdir(outdir)\n",
    "\n",
    "for i in range(36):  \n",
    "    year=1979+i   ### for years 1979-2014\n",
    "    cc = xr.apply_ufunc(\n",
    "\n",
    "    get_corrcoef, \n",
    "\n",
    "    ds_cmip6['pr'].sel(time=slice(str(year),str(year))).load(), ds_tp_mon_remap['pr'].sel(time=slice(str(year),str(year))).load(),        # then the input data arrays\n",
    "\n",
    "    input_core_dims=[['time'], ['time']],\n",
    "\n",
    "    vectorize=True\n",
    "\n",
    "    )\n",
    "    ax.append([])\n",
    "\n",
    "    ax[i] = plt.axes(projection=ccrs.PlateCarree())\n",
    "    ax[i].coastlines()\n",
    "    cc.plot(ax=ax[i],\n",
    "        vmin=-1, vmax=1,\n",
    "        cmap='coolwarm',    # Change the colormap back to 'bwr'\n",
    "        cbar_kwargs={\n",
    "            'extend':'neither' # Don't extend the colorbar in either direction. Other possibilities\n",
    "                               # would be 'both', 'min', or 'max'\n",
    "        }\n",
    "    )\n",
    "    plt.title(f\"Correlation of Pr CMIP6 vs ERA5 {year}\")\n",
    "    plt.savefig(f\"output/Correlation_of_Pr_CMIP6_vs_ERA5_{year}.png\")\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we can convert our PNG files into a GIF:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### create GIF from PNG files\n",
    "\n",
    "img, *imgs = [Image.open(f) for f in sorted(glob.glob('output/Correlation_of_Pr_CMIP6_vs_ERA5_*.png'))]\n",
    "img.save(fp='output/correlation.gif', format='GIF', append_images=imgs,\n",
    "         save_all=True, duration=200, loop=0)\n",
    "\n",
    "### clean up output directory\n",
    "del_files = glob.glob('output/Correlation_of_Pr_CMIP6_vs_ERA5_*.png')\n",
    "for file in del_files:\n",
    "    os.remove(file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(HTML(\"<img src='output/correlation.gif' />\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "This notebook shows correlation maps between CMIP6 and ERA5 data, which demonstrate the areas where model simulations are consistent with observations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference\n",
    "\n",
    "https://vimeo.com/112794571"
   ]
  }
 ],
 "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}